diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index f65130c3198b..136c96d218a1 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -31,16 +31,16 @@ steps: - text: "What is the release version?" key: release-version - - group: "Build Python wheels" + - group: "Build CUDA 13.0 Python wheels" key: "build-wheels" steps: - - label: "Build wheel - aarch64 - CUDA 12.9" + - label: "Build wheel - aarch64 - CUDA 13.0" depends_on: ~ - id: build-wheel-arm64-cuda-12-9 + id: build-wheel-arm64-cuda-13-0 agents: queue: arm64_cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" @@ -48,13 +48,37 @@ steps: env: DOCKER_BUILDKIT: "1" - - label: "Build wheel - aarch64 - CUDA 13.0" + - label: "Build wheel - x86_64 - CUDA 13.0" depends_on: ~ - id: build-wheel-arm64-cuda-13-0 + id: build-wheel-x86-cuda-13-0 + agents: + queue: cpu_queue_release + commands: + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "mkdir artifacts" + - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" + - "bash .buildkite/scripts/upload-nightly-wheels.sh" + - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels' + env: + DOCKER_BUILDKIT: "1" + + - block: "Unblock to build additional Python wheels" + depends_on: ~ + key: block-build-additional-wheels + if: build.env("NIGHTLY") != "1" + + - group: "Build additional Python wheels" + key: "build-additional-wheels" + depends_on: block-build-additional-wheels + allow_dependency_failure: true + steps: + - label: "Build wheel - aarch64 - CUDA 12.9" + depends_on: ~ + id: build-wheel-arm64-cuda-12-9 agents: queue: arm64_cpu_queue_release commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." + - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda12.9 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - "mkdir artifacts" - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - "bash .buildkite/scripts/upload-nightly-wheels.sh" @@ -133,20 +157,6 @@ steps: env: DOCKER_BUILDKIT: "1" - - label: "Build wheel - x86_64 - CUDA 13.0" - depends_on: ~ - id: build-wheel-x86-cuda-13-0 - agents: - queue: cpu_queue_release - commands: - - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ." - - "mkdir artifacts" - - "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'" - - "bash .buildkite/scripts/upload-nightly-wheels.sh" - - 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels' - env: - DOCKER_BUILDKIT: "1" - - label: "Build wheel - x86_64 - CPU" depends_on: ~ id: build-wheel-x86-cpu @@ -162,12 +172,26 @@ steps: DOCKER_BUILDKIT: "1" - label: "Generate and upload wheel indices" + key: generate-wheel-indices depends_on: "build-wheels" allow_dependency_failure: true + if: build.env("NIGHTLY") != "1" + agents: + queue: cpu_queue_release + commands: + - "UPDATE_VERSION_INDEX=0 bash .buildkite/scripts/generate-and-upload-nightly-index.sh" + + - label: "Regenerate indices with additional wheels" + key: generate-additional-wheel-indices + depends_on: + - build-wheels + - build-additional-wheels + - generate-wheel-indices + allow_dependency_failure: true agents: queue: cpu_queue_release commands: - - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh" + - 'UPDATE_NIGHTLY_INDEX="$${NIGHTLY:-0}" bash .buildkite/scripts/generate-and-upload-nightly-index.sh' - block: "Unblock to build release Docker images" depends_on: ~ @@ -566,10 +590,17 @@ steps: # # ============================================================================= + - block: "Unblock ROCm wheel/image prerequisites" + depends_on: ~ + key: block-build-rocm + if: build.env("NIGHTLY") != "1" + # ROCm Job 1: Build ROCm Base Wheels (with S3 caching) - label: ":rocm: Build ROCm Base Image & Wheels" id: build-rocm-base-wheels - depends_on: ~ + depends_on: + - step: block-build-rocm + allow_failure: true agents: queue: cpu_queue_release commands: @@ -974,6 +1005,8 @@ steps: depends_on: - input-release-version - build-wheels + - build-additional-wheels + - generate-additional-wheel-indices - label: "Upload release wheels to PyPI" depends_on: diff --git a/.buildkite/scripts/generate-and-upload-nightly-index.sh b/.buildkite/scripts/generate-and-upload-nightly-index.sh index 502ed0609310..1fa75994c01a 100755 --- a/.buildkite/scripts/generate-and-upload-nightly-index.sh +++ b/.buildkite/scripts/generate-and-upload-nightly-index.sh @@ -45,8 +45,10 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --curr echo "Uploading indices to $S3_COMMIT_PREFIX" aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX" -# copy to /nightly/ only if it is on the main branch and not a PR -if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then +# copy to /nightly/ only when enabled for a main branch build that is not a PR +if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \ + "$BUILDKITE_BRANCH" == "main" && \ + "$BUILDKITE_PULL_REQUEST" == "false" ]]; then echo "Uploading indices to overwrite /nightly/" aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/" fi @@ -67,7 +69,7 @@ pure_version="${version%%+*}" echo "Pure version (without variant): $pure_version" # re-generate and copy to // only if it does not have "dev" in the version -if [[ "$version" != *"dev"* ]]; then +if [[ "${UPDATE_VERSION_INDEX:-1}" == "1" && "$version" != *"dev"* ]]; then echo "Re-generating indices for /$pure_version/" rm -rf "${INDICES_OUTPUT_DIR:?}" mkdir -p "$INDICES_OUTPUT_DIR" diff --git a/.buildkite/scripts/upload-rocm-wheels.sh b/.buildkite/scripts/upload-rocm-wheels.sh index 1f3655631204..65fca944c7d5 100755 --- a/.buildkite/scripts/upload-rocm-wheels.sh +++ b/.buildkite/scripts/upload-rocm-wheels.sh @@ -113,8 +113,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py \ echo "Uploading indices to $S3_COMMIT_PREFIX" aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX" -# Update rocm/nightly/ if on main branch and not a PR -if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] || [[ "$NIGHTLY" == "1" ]]; then +# Only scheduled nightly builds should update the moving nightly index. +if [[ "${NIGHTLY:-0}" == "1" ]]; then echo "Updating rocm/nightly/ index..." aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/" fi @@ -147,7 +147,7 @@ echo "" echo "Install command (by commit):" echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/" echo "" -if [[ "$BUILDKITE_BRANCH" == "main" ]] || [[ "$NIGHTLY" == "1" ]]; then +if [[ "${NIGHTLY:-0}" == "1" ]]; then echo "Install command (nightly):" echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/" fi diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index ce6ee3857bda..bbde36e970b2 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1444,9 +1444,9 @@ steps: optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - - vllm/model_executor/layers/fla/ops/kda.py - - vllm/model_executor/layers/fla/ops/chunk_delta_h.py - - vllm/model_executor/layers/fla/ops/l2norm.py + - vllm/third_party/flash_linear_attention/ops/kda.py + - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py + - vllm/third_party/flash_linear_attention/ops/l2norm.py - tests/kernels/test_kda.py - vllm/platforms/rocm.py commands: @@ -3234,7 +3234,7 @@ steps: - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/qwen3_next_mtp.py - - vllm/model_executor/layers/fla/ops/ + - vllm/third_party/flash_linear_attention/ops/ - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: @@ -3473,7 +3473,7 @@ steps: - vllm/model_executor/models/qwen3.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/qwen3_next_mtp.py - - vllm/model_executor/layers/fla/ops/ + - vllm/third_party/flash_linear_attention/ops/ - vllm/_aiter_ops.py - vllm/v1/attention/backends/triton_attn.py - vllm/v1/attention/backends/rocm_attn.py diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index b25884a96f54..4012f1ebd539 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -203,3 +203,25 @@ steps: commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh + +# P TP 4 - D DPEP 4 test case for DSv4-Flash +- label: DSv4-Flash Disaggregated DP EP + key: dsv4-flash-disaggregated + timeout_in_minutes: 60 + device: h200 + optional: true + working_dir: "/vllm-workspace/tests" + num_devices: 8 + env: + ENABLE_HMA_FLAG: "1" + DP_EP: "1" + GPU_MEMORY_UTILIZATION: "0.85" + PREFILLER_TP_SIZE: "4" + DECODER_TP_SIZE: "4" + PREFILL_BLOCK_SIZE: "256" + DECODE_BLOCK_SIZE: "256" + MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash" + VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8" + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index d1f3efd6f070..150c3da57a33 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -176,9 +176,9 @@ steps: timeout_in_minutes: 25 device: h200_18gb source_file_dependencies: - - vllm/model_executor/layers/fla/ops/kda.py - - vllm/model_executor/layers/fla/ops/chunk_delta_h.py - - vllm/model_executor/layers/fla/ops/l2norm.py + - vllm/third_party/flash_linear_attention/ops/kda.py + - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py + - vllm/third_party/flash_linear_attention/ops/l2norm.py - tests/kernels/test_kda.py commands: - pytest -v -s kernels/test_kda.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index d2cee7b6365a..9c08c96e4c4a 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -103,7 +103,7 @@ steps: - vllm/transformers_utils/configs/qwen3_5_moe.py - vllm/model_executor/models/qwen3_next.py - vllm/model_executor/models/qwen3_next_mtp.py - - vllm/model_executor/layers/fla/ops/ + - vllm/third_party/flash_linear_attention/ops/ commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index 83d2fa26c22c..6066c9bc7081 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -103,10 +103,10 @@ steps: commands: - pytest -v -s models/language/generation_ppl_test -- label: Language Models Test (Extended Pooling) # 36min +- label: Language Models Test (Extended Pooling) device: h200_35gb key: language-models-test-extended-pooling - timeout_in_minutes: 70 + timeout_in_minutes: 120 optional: true source_file_dependencies: - vllm/ @@ -116,7 +116,7 @@ steps: mirror: amd: device: mi325_1 - timeout_in_minutes: 100 + timeout_in_minutes: 120 depends_on: - image-build-amd diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 57166d9d9b78..4f63a75f045b 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -172,7 +172,7 @@ mkdocs.yaml @hmellor # Kernels /vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep /vllm/v1/attention/ops/triton_unified_attention.py @tdoublep -/vllm/model_executor/layers/fla @ZJY0516 @vadiklyutiy +/vllm/third_party/flash_linear_attention @ZJY0516 @vadiklyutiy # ROCm related: specify owner with write access to notify AMD folks for careful code review /vllm/**/*rocm* @tjtanaa @dllehr-amd diff --git a/.github/workflows/issue_autolabel.yml b/.github/workflows/issue_autolabel.yml index 7a98ce7cc08a..e5d1accf477c 100644 --- a/.github/workflows/issue_autolabel.yml +++ b/.github/workflows/issue_autolabel.yml @@ -323,7 +323,7 @@ jobs: // {users} will be replaced with @mentions const ccConfig = { rocm: { - users: ['hongxiayang', 'tjtanaa', 'vllmellm'], + users: ['hongxiayang', 'tjtanaa', 'vllmellm', 'giuseppegrossi'], message: 'CC {users} for ROCm-related issue', }, mistral: { diff --git a/.gitignore b/.gitignore index 26cd21a015d0..43787b8cd283 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,9 @@ vllm/third_party/deep_gemm/ # fmha_sm100 vendored package built from source vllm/third_party/fmha_sm100/ +# tml-fa4 vendored package built from source +vllm/third_party/tml_fa4/ + # triton jit .triton diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ad767c823bef..3d26a51bbacc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -30,7 +30,7 @@ repos: - id: markdownlint-cli2 language_version: lts args: [--fix] - exclude: ^CLAUDE\.md$ + exclude: (^|/)CLAUDE\.md$ - repo: https://github.com/rhysd/actionlint rev: v1.7.7 hooks: diff --git a/.yapfignore b/.yapfignore deleted file mode 100644 index 38158259032a..000000000000 --- a/.yapfignore +++ /dev/null @@ -1,2 +0,0 @@ -collect_env.py -vllm/model_executor/layers/fla/ops/*.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 6ee11e6ba269..f514ba41aa52 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1408,6 +1408,7 @@ if (VLLM_GPU_LANG STREQUAL "CUDA") include(cmake/external_projects/fmha_sm100.cmake) include(cmake/external_projects/flashmla.cmake) include(cmake/external_projects/qutlass.cmake) + include(cmake/external_projects/tml_fa4.cmake) # vllm-flash-attn should be last as it overwrites some CMake functions include(cmake/external_projects/vllm_flash_attn.cmake) diff --git a/benchmarks/kernels/bench_cp_gather_fp8.py b/benchmarks/kernels/bench_cp_gather_fp8.py index 19fc84c4df76..ca76597a1f3e 100644 --- a/benchmarks/kernels/bench_cp_gather_fp8.py +++ b/benchmarks/kernels/bench_cp_gather_fp8.py @@ -69,12 +69,11 @@ def make_inputs(total_tokens, num_reqs, block_size): # Output workspace dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda") - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") workspace_starts_t = torch.tensor( workspace_starts, dtype=torch.int32, device="cuda" ) - return cache, dst, block_table, seq_lens_t, workspace_starts_t + return cache, dst, block_table, workspace_starts_t def bench_scenario(label, num_reqs, total_tokens_list, save_path): @@ -94,7 +93,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path): ) ) def bench_fn(total_tokens, provider, num_reqs): - cache, dst, block_table, seq_lens_t, ws_starts = make_inputs( + cache, dst, block_table, ws_starts = make_inputs( total_tokens, num_reqs, BLOCK_SIZE ) @@ -102,7 +101,7 @@ def bench_fn(total_tokens, provider, num_reqs): ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( lambda: ops.cp_gather_and_upconvert_fp8_kv_cache( - cache, dst, block_table, seq_lens_t, ws_starts, num_reqs + cache, dst, block_table, ws_starts, num_reqs ), quantiles=quantiles, rep=500, diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 38d218d00acb..8f54fe58ddcd 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -28,9 +28,9 @@ if(DEEPGEMM_SRC_DIR) message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() # Keep in sync with tools/install_deepgemm.sh - set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") - # NOTE: This is currently targeting nv-dev branch due to sm120 support - set(_DEEPGEMM_UPSTREAM_TAG "a6b593d2826719dcf4892609af7b84ee23aaf32a") + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/cleonard530/DeepGEMM.git") + # TORCH_LIBRARY migration (migrate_pybind_to_torch_library); see cleonard530/DeepGEMM#2 + set(_DEEPGEMM_UPSTREAM_TAG "441c417c6cf7184593421273b7e6d79a0999a8f3") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -40,7 +40,7 @@ else() set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build") set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild") - if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp") + if(EXISTS "${_deepgemm_src}/deep_gemm/_C.py") set(deepgemm_SOURCE_DIR "${_deepgemm_src}") set(deepgemm_BINARY_DIR "${_deepgemm_bin}") else() @@ -87,33 +87,17 @@ if(DEEPGEMM_ARCHS) # # DeepGEMM integration notes # -------------------------- - # We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle a - # `_C.cpython-X.Y-*.so` for every CPython in `requires-python`. The - # per-Python build is delegated to tools/build_deepgemm_C.py. + # We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle: + # - deep_gemm/_C.py (Python shim over torch.ops.deep_gemm) + # - deep_gemm/_C_extension.abi3.so (single limited-API extension) + # The build is delegated to tools/build_deepgemm_C.py (setup.py build_ext). # - # Why per-Python: DeepGEMM's binding uses PYBIND11_MODULE, which links - # private CPython symbols — a single `_C.abi3.so` is not viable today - # (see #41476 / #41512 for the failed attempt). - # - # TODOs (tracked in vllm-project/vllm#42431): - # - Replace DeepGEMM's pybind11 binding with a TORCH_LIBRARY + shim - # binding (cf. vllm-flash-attention/csrc/common/pytorch_shim.h) to - # collapse to one `_C.abi3.so`. Needs either an upstream change or - # a maintained binding fork in vLLM. - # - AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop - # the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime - # requirement. + # TODO: AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop + # the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime + # requirement. # - # DEEPGEMM_PYTHON_INTERPRETERS: ":"-separated target Python paths. - # Empty/unset → fall back to the build interpreter (editable installs). - # (Empty-but-set env vars test as DEFINED in cmake — treat as unset.) - if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "") - string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}") - else() - set(_dg_pythons "${Python_EXECUTABLE}") - endif() - message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}") + message(STATUS "DeepGEMM extension will be built with: ${Python_EXECUTABLE}") # add_custom_command does no implicit header scanning; glob explicitly so # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild. @@ -124,40 +108,29 @@ if(DEEPGEMM_ARCHS) "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh") - set(_dg_markers) - set(_dg_seen_soabis) - foreach(_pybin IN LISTS _dg_pythons) - execute_process( - COMMAND "${_pybin}" -c - "import sysconfig; print(sysconfig.get_config_var('SOABI'))" - OUTPUT_VARIABLE _dg_soabi - OUTPUT_STRIP_TRAILING_WHITESPACE - COMMAND_ERROR_IS_FATAL ANY) - # Dedup interpreters that resolve to the same CPython. - if(_dg_soabi IN_LIST _dg_seen_soabis) - continue() - endif() - list(APPEND _dg_seen_soabis "${_dg_soabi}") - set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}") - set(_dg_marker "${_dg_dir}/.built") - add_custom_command( - OUTPUT "${_dg_marker}" - COMMAND "${Python_EXECUTABLE}" - "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py" - "${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}" - COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}" - DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py" - "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp" - ${_dg_headers} - COMMENT "Building DeepGEMM _C for ${_pybin}" - VERBATIM) - list(APPEND _dg_markers "${_dg_marker}") - install(DIRECTORY "${_dg_dir}/" - DESTINATION vllm/third_party/deep_gemm - COMPONENT _deep_gemm_C - FILES_MATCHING PATTERN "_C.cpython-*.so") - endforeach() - add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers}) + set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C") + set(_dg_marker "${_dg_dir}/.built") + add_custom_command( + OUTPUT "${_dg_marker}" + COMMAND "${Python_EXECUTABLE}" + "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py" + "${deepgemm_SOURCE_DIR}" "${_dg_dir}" + COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}" + DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py" + "${deepgemm_SOURCE_DIR}/csrc/python_api.cpp" + "${deepgemm_SOURCE_DIR}/deep_gemm/_C.py" + "${deepgemm_SOURCE_DIR}/setup.py" + ${_dg_headers} + COMMENT "Building DeepGEMM _C_extension (abi3)" + VERBATIM) + add_custom_target(_deep_gemm_C ALL DEPENDS "${_dg_marker}") + + install(DIRECTORY "${_dg_dir}/" + DESTINATION vllm/third_party/deep_gemm + COMPONENT _deep_gemm_C + FILES_MATCHING + PATTERN "_C.py" + PATTERN "_C_extension*.so") # # Vendor DeepGEMM Python package files diff --git a/cmake/external_projects/tml_fa4.cmake b/cmake/external_projects/tml_fa4.cmake new file mode 100644 index 000000000000..59e2e241c928 --- /dev/null +++ b/cmake/external_projects/tml_fa4.cmake @@ -0,0 +1,50 @@ +include(FetchContent) + +if(DEFINED ENV{TML_FA4_SRC_DIR}) + set(TML_FA4_SRC_DIR $ENV{TML_FA4_SRC_DIR}) +endif() + +if(TML_FA4_SRC_DIR) + FetchContent_Declare( + tml_fa4 + SOURCE_DIR ${TML_FA4_SRC_DIR} + CONFIGURE_COMMAND "" + BUILD_COMMAND "") +else() + FetchContent_Declare( + tml_fa4 + GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git + GIT_TAG 13374f0c855acc1add1bf30444bd67aebbc24a8e + GIT_PROGRESS TRUE + CONFIGURE_COMMAND "" + BUILD_COMMAND "") +endif() + +FetchContent_GetProperties(tml_fa4) +if(NOT tml_fa4_POPULATED) + FetchContent_Populate(tml_fa4) +endif() +message(STATUS "tml-fa4 is available at ${tml_fa4_SOURCE_DIR}") + +add_custom_target(tml_fa4) + +# Install into a private namespace so this implementation cannot shadow the +# flash_attn package used by vLLM's standard attention backends. +install(CODE " + file(GLOB_RECURSE TML_FA4_PY_FILES + \"${tml_fa4_SOURCE_DIR}/flash_attn/cute/*.py\") + foreach(SRC_FILE \${TML_FA4_PY_FILES}) + file(RELATIVE_PATH REL_PATH + \"${tml_fa4_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE}) + set(DST_FILE + \"\${CMAKE_INSTALL_PREFIX}/vllm/third_party/tml_fa4/\${REL_PATH}\") + get_filename_component(DST_DIR \${DST_FILE} DIRECTORY) + file(MAKE_DIRECTORY \${DST_DIR}) + file(READ \${SRC_FILE} FILE_CONTENTS) + string(REPLACE + \"flash_attn.cute\" + \"vllm.third_party.tml_fa4\" + FILE_CONTENTS \"\${FILE_CONTENTS}\") + file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\") + endforeach() +" COMPONENT tml_fa4) diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 97a8cfe87b7b..28a6b336e27b 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/cmake/external_projects/vllm_flash_attn.cmake @@ -39,7 +39,7 @@ else() FetchContent_Declare( vllm-flash-attn GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git - GIT_TAG bb9a72e7dde0dc614ffc663e052cd6a19ce73a42 + GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8 GIT_PROGRESS TRUE # Don't share the vllm-flash-attn build between build types BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn diff --git a/csrc/cache.h b/csrc/cache.h index a9e74b0dc2df..dbc2675fc70c 100644 --- a/csrc/cache.h +++ b/csrc/cache.h @@ -67,9 +67,8 @@ void cp_gather_and_upconvert_fp8_kv_cache( torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] torch::Tensor const& dst, // [TOT_TOKENS, 576] torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::Tensor const& seq_lens, // [BATCH] torch::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size); + int64_t batch_size, std::optional seq_starts = std::nullopt); // Indexer K quantization and cache function void indexer_k_quant_and_cache( diff --git a/csrc/libtorch_stable/cache_kernels.cu b/csrc/libtorch_stable/cache_kernels.cu index a1ac81cb10a4..2d4b47b4b43c 100644 --- a/csrc/libtorch_stable/cache_kernels.cu +++ b/csrc/libtorch_stable/cache_kernels.cu @@ -1174,7 +1174,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache( const int32_t num_reqs, const int32_t block_size, const int32_t total_tokens, const int64_t block_table_stride, const int64_t cache_block_stride, const int64_t cache_entry_stride, - const int64_t dst_entry_stride) { + const int64_t dst_entry_stride, + const int32_t* __restrict__ seq_starts) { // Optional source offsets const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; if (flat_warp_id >= total_tokens) return; const int lane_id = threadIdx.x & 31; @@ -1192,7 +1193,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache( // Compute physical token address via block table const int out_token_id = flat_warp_id; - const int token_offset = out_token_id - workspace_starts[req_id]; + int token_offset = out_token_id - workspace_starts[req_id]; + if (seq_starts != nullptr) token_offset += seq_starts[req_id]; const int cache_block_idx = token_offset / block_size; const int offset_in_block = token_offset % block_size; const int physical_block = @@ -1383,9 +1385,9 @@ void cp_gather_and_upconvert_fp8_kv_cache( torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656] torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& seq_lens, // [BATCH] torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size) { + int64_t batch_size, + std::optional seq_starts = std::nullopt) { torch::stable::accelerator::DeviceGuard device_guard( src_cache.get_device_index()); const cudaStream_t stream = get_current_cuda_stream(); @@ -1396,20 +1398,25 @@ void cp_gather_and_upconvert_fp8_kv_cache( STD_TORCH_CHECK( block_table.scalar_type() == torch::headeronly::ScalarType::Int, "block_table must be int32"); - STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int, - "seq_lens must be int32"); STD_TORCH_CHECK( workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int, "workspace_starts must be int32"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK( + seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int, + "seq_starts must be int32"); + } STD_TORCH_CHECK(src_cache.device() == dst.device(), "src_cache and dst must be on the same device"); STD_TORCH_CHECK(src_cache.device() == block_table.device(), "src_cache and block_table must be on the same device"); - STD_TORCH_CHECK(src_cache.device() == seq_lens.device(), - "src_cache and seq_lens must be on the same device"); STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(), "src_cache and workspace_starts must be on the same device"); + if (seq_starts.has_value()) { + STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(), + "src_cache and seq_starts must be on the same device"); + } auto dtype = src_cache.scalar_type(); STD_TORCH_CHECK( dtype == torch::headeronly::ScalarType::Byte || // uint8 @@ -1438,6 +1445,9 @@ void cp_gather_and_upconvert_fp8_kv_cache( constexpr int warps_per_block = 8; const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block; const int block_size_threads = warps_per_block * 32; // 256 threads + const int32_t* seq_starts_ptr = + seq_starts.has_value() ? seq_starts.value().const_data_ptr() + : nullptr; vllm::cp_gather_and_upconvert_fp8_kv_cache<<>>( @@ -1446,7 +1456,7 @@ void cp_gather_and_upconvert_fp8_kv_cache( workspace_starts.const_data_ptr(), static_cast(batch_size), block_size, total_tokens, block_table_stride, cache_block_stride, cache_entry_stride, - dst_entry_stride); + dst_entry_stride, seq_starts_ptr); } // Macro to dispatch the kernel based on the data type. diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu index 095a76678311..785bbf2f6e07 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu @@ -44,6 +44,12 @@ typedef __hip_bfloat162 __nv_bfloat162; namespace vllm { namespace moe { +template +__device__ __forceinline__ int64_t load_index_as_int64(const HashIndType* ptr, + int64_t offset) { + return static_cast(ptr[offset]); +} + /// Aligned array type template + typename HashIndType, typename InputType = float> __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ void topkGatingSoftplusSqrt( const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices, int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, - const IndType* input_ids, const IndType* tid2eid) { + const HashIndType* input_ids, const HashIndType* tid2eid) { static_assert(std::is_same_v || std::is_same_v || std::is_same_v, @@ -240,8 +246,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ // Hash MoE path: indices are predetermined from lookup table if constexpr (USE_HASH) { - const IndType token_id = input_ids[thread_row]; - const IndType* expert_indices_for_token = tid2eid + token_id * k; + const int64_t token_id = load_index_as_int64(input_ids, thread_row); + const int64_t token_expert_offset = token_id * static_cast(k); #pragma unroll for (int ii = 0; ii < VPT; ++ii) { float val = row_chunk[ii]; @@ -252,7 +258,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ float selected_sum = 0.f; #pragma unroll for (int k_idx = 0; k_idx < k; ++k_idx) { - const int expert = expert_indices_for_token[k_idx]; + const int expert = static_cast( + load_index_as_int64(tid2eid, token_expert_offset + k_idx)); const int idx = k * thread_row + k_idx; for (int ii = 0; ii < VPT; ++ii) { const int group_id = ii / ELTS_PER_LDG; @@ -261,7 +268,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id; if (expert == expert_idx) { - indices[idx] = expert; + indices[idx] = static_cast(expert); selected_sum += row_chunk[ii]; break; } @@ -285,7 +292,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ #pragma unroll for (int k_idx = 0; k_idx < k; ++k_idx) { - const int expert = expert_indices_for_token[k_idx]; + const int expert = static_cast( + load_index_as_int64(tid2eid, token_expert_offset + k_idx)); const int idx = k * thread_row + k_idx; for (int ii = 0; ii < VPT; ++ii) { const int group_id = ii / ELTS_PER_LDG; @@ -461,14 +469,15 @@ struct TopkConstants { } template + int MAX_BYTES_PER_LDG, typename IndType, typename HashIndType, + typename InputType> void topkGatingSoftplusSqrtLauncherHelper( const InputType* input, const bool* finished, float* output, IndType* indices, int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, - const bool use_hash, const IndType* input_ids, const IndType* tid2eid, - cudaStream_t stream) { + const bool use_hash, const HashIndType* input_ids, + const HashIndType* tid2eid, cudaStream_t stream) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = @@ -481,7 +490,8 @@ void topkGatingSoftplusSqrtLauncherHelper( DISPATCH_HASH(use_hash, USE_HASH, { auto* kernel = &topkGatingSoftplusSqrt; + WARP_SIZE_PARAM, USE_HASH, IndType, HashIndType, + InputType>; #ifndef USE_ROCM cudaLaunchConfig_t config = {}; config.gridDim = num_blocks; @@ -538,13 +548,14 @@ void topkGatingSoftplusSqrtLauncherHelper( } #endif -template +template void topkGatingSoftplusSqrtKernelLauncher( const InputType* gating_output, float* topk_weights, IndType* topk_indices, int* token_expert_indices, const int num_tokens, const int num_experts, const int topk, const bool renormalize, double routed_scaling_factor, - const float* correction_bias, const bool use_hash, const IndType* input_ids, - const IndType* tid2eid, cudaStream_t stream) { + const float* correction_bias, const bool use_hash, + const HashIndType* input_ids, const HashIndType* tid2eid, + cudaStream_t stream) { static constexpr int WARPS_PER_TB = 4; static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; // for bfloat16 dtype, we need 4 bytes loading to make sure num_experts @@ -644,57 +655,55 @@ void dispatch_topk_softplus_sqrt_launch( if (correction_bias.has_value()) { bias_ptr = correction_bias.value().const_data_ptr(); } - bool use_hash = false; - if (tid2eid.has_value()) { - STD_TORCH_CHECK(input_ids.has_value(), - "input_ids is required for hash MoE"); - use_hash = true; - } - if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { - const int* input_ids_ptr = nullptr; - const int* tid2eid_ptr = nullptr; + + auto launch = [&](auto* topk_indices_ptr) { + using OutIndType = + typename std::remove_pointer::type; if (tid2eid.has_value()) { - input_ids_ptr = input_ids.value().const_data_ptr(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); + STD_TORCH_CHECK(input_ids.has_value(), + "input_ids is required for hash MoE"); + STD_TORCH_CHECK( + input_ids.value().scalar_type() == tid2eid.value().scalar_type(), + "input_ids and tid2eid must have the same dtype"); + if (tid2eid.value().scalar_type() == + torch::headeronly::ScalarType::Long) { + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, true, input_ids.value().const_data_ptr(), + tid2eid.value().const_data_ptr(), stream); + } else { + STD_TORCH_CHECK(tid2eid.value().scalar_type() == + torch::headeronly::ScalarType::Int); + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, true, input_ids.value().const_data_ptr(), + tid2eid.value().const_data_ptr(), stream); + } + } else { + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, false, static_cast(nullptr), + static_cast(nullptr), stream); } + }; - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, - topk, renormalize, routed_scaling_factor, bias_ptr, use_hash, - input_ids_ptr, tid2eid_ptr, stream); + if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { + launch(topk_indices.mutable_data_ptr()); } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) { - const uint32_t* input_ids_ptr = nullptr; - const uint32_t* tid2eid_ptr = nullptr; - if (tid2eid.has_value()) { - input_ids_ptr = input_ids.value().const_data_ptr(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); - } - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, - topk, renormalize, routed_scaling_factor, bias_ptr, use_hash, - input_ids_ptr, tid2eid_ptr, stream); + launch(topk_indices.mutable_data_ptr()); } else { STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long); - - const int64_t* input_ids_ptr = nullptr; - const int64_t* tid2eid_ptr = nullptr; - if (tid2eid.has_value()) { - input_ids_ptr = input_ids.value().const_data_ptr(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); - } - - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, - topk, renormalize, routed_scaling_factor, bias_ptr, use_hash, - input_ids_ptr, tid2eid_ptr, stream); + launch(topk_indices.mutable_data_ptr()); } } @@ -738,4 +747,4 @@ void topk_softplus_sqrt( STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } -} \ No newline at end of file +} diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index ae274226fe23..3834bea58576 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -527,9 +527,9 @@ void cp_gather_and_upconvert_fp8_kv_cache( // 656] torch::stable::Tensor const& dst, // [TOT_TOKENS, 576] torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES] - torch::stable::Tensor const& seq_lens, // [BATCH] torch::stable::Tensor const& workspace_starts, // [BATCH] - int64_t batch_size); + int64_t batch_size, + std::optional seq_starts = std::nullopt); // Indexer K quantization and cache function void indexer_k_quant_and_cache( diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index fd2ffe92e089..0a475d02c6fe 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -847,8 +847,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) { ops.def( "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, " - "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int " - "batch_size) -> ()"); + "Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? " + "seq_starts) -> ()"); ops.def( "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor " diff --git a/docker/Dockerfile b/docker/Dockerfile index b47853a06c73..f008fd29e15f 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -248,10 +248,6 @@ COPY requirements/common.txt requirements/common.txt COPY requirements/cuda.txt requirements/cuda.txt COPY use_existing_torch.py use_existing_torch.py COPY pyproject.toml pyproject.toml -# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that -# share paths with different content. uv can extract them in either order, -# leaving base files that break CUDA 13 CuTe DSL JIT. -# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed. RUN --mount=type=cache,target=/opt/uv/cache \ if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \ sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \ @@ -268,13 +264,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \ else \ uv pip install --python /opt/venv/bin/python3 -r requirements/cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \ - fi \ - && if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \ - CUTLASS_DSL_VERSION=$(uv pip show --python /opt/venv/bin/python3 nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \ - if [ -n "$CUTLASS_DSL_VERSION" ]; then \ - uv pip install --python /opt/venv/bin/python3 --force-reinstall --no-deps \ - "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \ - fi; \ fi # Track PyTorch lib versions used during build and match in downstream instances. @@ -792,10 +781,6 @@ ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0 ARG PYTORCH_CUDA_INDEX_BASE_URL COPY requirements/common.txt /tmp/common.txt COPY requirements/cuda.txt /tmp/requirements-cuda.txt -# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that -# share paths with different content. uv can extract them in either order, -# leaving base files that break CUDA 13 CuTe DSL JIT. -# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed. RUN --mount=type=cache,target=/opt/uv/cache \ if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \ sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \ @@ -803,19 +788,12 @@ RUN --mount=type=cache,target=/opt/uv/cache \ fi && \ uv pip install --system -r /tmp/requirements-cuda.txt \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \ - if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \ - CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \ - if [ -n "$CUTLASS_DSL_VERSION" ]; then \ - uv pip install --system --force-reinstall --no-deps \ - "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \ - fi; \ - fi && \ rm /tmp/requirements-cuda.txt /tmp/common.txt # Install FlashInfer JIT cache (requires CUDA-version-specific index URL) # https://docs.flashinfer.ai/installation.html # From versions.json: .flashinfer.version -ARG FLASHINFER_VERSION=0.6.13 +ARG FLASHINFER_VERSION=0.6.14 RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \ --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') @@ -908,19 +886,6 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm uv pip install --system ep_kernels/dist/*.whl --verbose \ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') -# nvidia-cutlass-dsl[cu13] installs -libs-base and -libs-cu13 wheels that -# share paths with different content. Force -libs-cu13 last after runtime -# dependency installs so uv cannot leave base files behind. -# TODO(mmangkad): Remove this after NVIDIA/cutlass#3259 is fixed. -RUN --mount=type=cache,target=/opt/uv/cache \ - if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "13" ]; then \ - CUTLASS_DSL_VERSION=$(uv pip show --system nvidia-cutlass-dsl 2>/dev/null | awk '/^Version:/{print $2}') && \ - if [ -n "$CUTLASS_DSL_VERSION" ]; then \ - uv pip install --system --force-reinstall --no-deps \ - "nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}"; \ - fi; \ - fi - # CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will # return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers # consistently from the host (see https://github.com/vllm-project/vllm/issues/18859). diff --git a/docker/versions.json b/docker/versions.json index 4dffa00985c9..e6839bbb05cf 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -68,7 +68,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.13" + "default": "0.6.14" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/docs/configuration/env_vars.md b/docs/configuration/env_vars.md index f6d548a19d91..38de2a49760c 100644 --- a/docs/configuration/env_vars.md +++ b/docs/configuration/env_vars.md @@ -5,7 +5,7 @@ vLLM uses the following environment variables to configure the system: !!! warning Please note that `VLLM_PORT` and `VLLM_HOST_IP` set the port and ip for vLLM's **internal usage**. It is not the port and ip for the API server. If you use `--host $VLLM_HOST_IP` and `--port $VLLM_PORT` to start the API server, it will not work. - All environment variables used by vLLM are prefixed with `VLLM_`. **Special care should be taken for Kubernetes users**: please do not name the service as `vllm`, otherwise environment variables set by Kubernetes might conflict with vLLM's environment variables, because [Kubernetes sets environment variables for each service with the capitalized service name as the prefix](https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables). + Most vLLM-specific environment variables are prefixed with `VLLM_` (a handful of standard names — for example `CUDA_VISIBLE_DEVICES`, `MAX_JOBS`, `S3_ACCESS_KEY_ID`/`S3_SECRET_ACCESS_KEY`/`S3_ENDPOINT_URL`, `DO_NOT_TRACK`, `NO_COLOR` — are also read directly when set). **Special care should be taken for Kubernetes users**: please do not name the service as `vllm`, otherwise environment variables set by Kubernetes might conflict with vLLM's environment variables, because [Kubernetes sets environment variables for each service with the capitalized service name as the prefix](https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables). ```python --8<-- "vllm/envs.py:env-vars-definition" diff --git a/docs/contributing/ci/nightly_builds.md b/docs/contributing/ci/nightly_builds.md index 10c4a4372403..c6d83fd9fc78 100644 --- a/docs/contributing/ci/nightly_builds.md +++ b/docs/contributing/ci/nightly_builds.md @@ -6,7 +6,11 @@ vLLM maintains a per-commit wheel repository (commonly referred to as "nightly") ### Wheel Building -Wheels are built in the `Release` pipeline (`.buildkite/release-pipeline.yaml`) after a PR is merged into the main branch, with multiple variants: +Wheels are built in the `Release` pipeline +(`.buildkite/release-pipeline.yaml`) after a PR is merged into the main branch. +Regular builds produce the CUDA 13.0 wheels for x86_64 and aarch64. Additional +wheel variants and ROCm builds can be unblocked on demand and run automatically +when `NIGHTLY=1`: - **Backend variants**: `cpu` and `cuXXX` (e.g., `cu129`, `cu130`). - **Architecture variants**: `x86_64` and `aarch64`. diff --git a/docs/deployment/k8s.md b/docs/deployment/k8s.md index 7a92c99b2c4a..e7d0853e9f49 100644 --- a/docs/deployment/k8s.md +++ b/docs/deployment/k8s.md @@ -217,7 +217,7 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit) image: vllm/vllm-openai:latest command: ["/bin/sh", "-c"] args: [ - "vllm serve mistralai/Mistral-7B-Instruct-v0.3 --trust-remote-code --enable-chunked-prefill --max_num_batched_tokens 1024" + "vllm serve mistralai/Mistral-7B-Instruct-v0.3 --trust-remote-code --enable-chunked-prefill --max-num-batched-tokens 1024" ] env: - name: HF_TOKEN @@ -306,7 +306,7 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit) - SYS_PTRACE command: ["/bin/sh", "-c"] args: [ - "vllm serve mistralai/Mistral-7B-v0.3 --port 8000 --trust-remote-code --enable-chunked-prefill --max_num_batched_tokens 1024" + "vllm serve mistralai/Mistral-7B-v0.3 --port 8000 --trust-remote-code --enable-chunked-prefill --max-num-batched-tokens 1024" ] env: - name: HF_TOKEN diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 9d524a489be0..8062ce237db2 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -205,7 +205,7 @@ hardware and configuration. | Backend | Description | Dtypes | Compute Cap. | Notes | | ------- | ----------- | ------ | ------------ | ----- | -| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) | +| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) | | `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only | | `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only | | `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only | diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 37a9a7399908..5a1fe71a1eab 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -107,6 +107,7 @@ Batch invariance has been tested and verified on the following models: - **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-3B-Instruct` for example - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b` - **Mistral**: `mistralai/Mistral-7B-v0.3` +- **Phi series**: `microsoft/Phi-3.5-mini-instruct` Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose). diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 7042ef787c9f..6349b6329154 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -83,9 +83,11 @@ Each entry in `secondary_tiers` is a dict with a required `type` field plus tier The filesystem and object-store tiers can publish hash-only `BlockStored` KV events for blocks they successfully store, tagged with a stable per-tier `medium` (`FS` for the filesystem tier, `OBJ` for the object-store tier). Set `enable_kv_events: true` in the tier's entry to opt in; events are published only when KV cache events are also enabled globally via `--kv-events-config`. +Set the optional `locality` tier field to `LOCAL` or `REMOTE` to describe the tier's storage location relative to the publishing vLLM instance. `LOCAL` marks storage local to that instance, while `REMOTE` marks storage that is not local to it. When the setting is omitted, locality is unspecified. vLLM does not infer it from the tier type, so an OBJ tier is not implicitly `REMOTE`. A KV event includes `locality` only when the tier explicitly configures it. This metadata describes the tier property without implying that a consumer can already route requests to its blocks. + ### Filesystem (FS) -The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage. +The filesystem tier (`type: "fs"`) writes blocks to a filesystem directory. | Key | Required | Default | Notes | | --- | --- | --- | --- | @@ -94,6 +96,7 @@ The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage | `n_read_threads` | no | `16` | Read-priority I/O threads (load path). | | `n_write_threads` | no | `16` | Write-priority I/O threads (store path). | | `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `FS`) for successfully stored blocks. Requires KV cache events to be enabled globally. | +| `locality` | no | unspecified | `LOCAL` or `REMOTE` relative to the publishing vLLM instance. Included in the tier's KV events only when explicitly configured. | Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency. @@ -134,6 +137,7 @@ The object-store tier (`type: "obj"`) offloads blocks to an S3-compatible object | `prefix` | no | `""` | Key prefix prepended to all object keys. | | `io_threads` | no | `4` | Number of NIXL OBJ backend I/O threads. | | `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `OBJ`) for successfully stored blocks. Requires KV cache events to be enabled globally. | +| `locality` | no | unspecified | `LOCAL` or `REMOTE` relative to the publishing vLLM instance. Included in the tier's KV events only when explicitly configured; OBJ does not imply `REMOTE`. | `store_config` fields: diff --git a/docs/mkdocs/overrides/main.html b/docs/mkdocs/overrides/main.html index bdd62ebc158d..a3f3dd67f52e 100644 --- a/docs/mkdocs/overrides/main.html +++ b/docs/mkdocs/overrides/main.html @@ -1,5 +1,5 @@ {% extends "base.html" %} {% block announce %} -

You are viewing the latest developer preview docs. Click here to view docs for the latest stable release.

+

You are viewing the latest developer preview docs. Click here to view docs for the latest stable release.

{% endblock %} diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index b7c2ee873750..9a8bb68bc11f 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -220,8 +220,6 @@ For multi-node deployment, add these EPLB flags to each node's command. We recom - Use simulator flags `VLLM_MOE_ROUTING_SIMULATION_STRATEGY=uniform_random` and `VLLM_RANDOMIZE_DP_DUMMY_INPUTS=1` so token routing is balanced across EP ranks. -- Increasing `VLLM_MOE_DP_CHUNK_SIZE` may increase throughput by increasing the maximum batch size for inter-rank token transfers. This may cause DeepEP to throw `assert self.nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2`, which can be fixed by increasing environment variable `NVSHMEM_QP_DEPTH`. - ## Disaggregated Serving (Prefill/Decode Split) For production deployments requiring strict SLA guarantees for time-to-first-token and inter-token latency, disaggregated serving allows independent scaling of prefill and decode operations. diff --git a/docs/serving/online_serving/README.md b/docs/serving/online_serving/README.md index 6d984f1a62d7..5df1821f3b38 100644 --- a/docs/serving/online_serving/README.md +++ b/docs/serving/online_serving/README.md @@ -137,8 +137,12 @@ For further details on renderer APIs, please refer to [this page](renderer.md). ### Derenderer APIs -- `/v1/completions/derender` - Derenderer completion requests -- `/v1/chat/completions/derender` - Derenderer chat completion requests +For further details on derenderer APIs, please refer to [this page](derenderer.md). + +- [Chat Completions Derender API](derenderer.md) (`/v1/chat/completions/derender`) + - Derender chat completion requests +- [Completions Derender API](derenderer.md) (`/v1/completions/derender`) + - Derender completion requests ## Tokenize APIs diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md new file mode 100644 index 000000000000..e55cd3109493 --- /dev/null +++ b/docs/serving/online_serving/derenderer.md @@ -0,0 +1,98 @@ +# Derenderer APIs + +The derenderer API is the post processing counterpart to the [Renderer APIs](renderer.md). Where `/render` turns a request into token ID (preprocessing), `/derender` turns generated token IDs back into a fully formed OpenAI compatible response (detokenization, reasoning parsing, tool call parsing), all without a GPU. + +This closes the loop for a token-in / token-out engine in disaggregated serving: + +- **GPU less post processing**: Detokenization, reasoning parsing, and tool call parsing run on the same GPU less frontend that hosts `/render` +- **Parser parity**: The derenderer reuses vLLM's tool and reasoning parsers, so a disaggregated deployment produces the same `content`/`reasoning`/ `tool_calls` split as a standard `vllm serve` server +- **Non-streaming**: The endpoints expect a complete `GenerateResponse` with all token IDs present and perform one-shot parsing. Streaming derender would require a separate endpoint design and is not currently supported but is in the pipeline + +Both endpoints are hosted by the GPU less rendering server started with [`vllm launch render`](../../cli/launch/render.md), alongside the `/render` +endpoints. + +## Pipeline + +```text + render generate derender + request ───────────────▶ token_ids ─────────▶ token_ids ──────────▶ response + (chat / (GPU less) (token-in / (GPU less) (OpenAI + completion) │ token-out engine) ▲ compatible) + └─────────────── request + prompt_tokens ──┘ +``` + +The derender step needs more than the engine's `token_ids`. It also consumes the original `chat_request`/`completion_request` and `prompt_tokens` carried over from the render step (see [Request format](#request-format)) so the tool and reasoning parsers have the context they need. + +## API Reference + +- Chat Completions Derender API (`/v1/chat/completions/derender`) + - Post process a single `GenerateResponse` into a `ChatCompletionResponse` +- Completions Derender API (`/v1/completions/derender`) + - Post process a list of `GenerateResponse` objects (one per prompt) into a `CompletionResponse` + +## Request format + +Each request wraps the engine's `GenerateResponse`(s) together with the caller metadata needed to reconstruct the final response without a GPU. + +`/v1/chat/completions/derender`: + +??? code + + ```python + --8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-chat-request" + ``` + +`/v1/completions/derender`: + +??? code + + ```python + --8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-request" + ``` + +Oversized payloads are rejected with a `400` before any `tokenizer.decode()` or parser runs. + +## Example + +The example below drives the full `render → generate → derender` round trip for a chat request against a GPU less render server (`/render`, `/derender`) and a token-in / token-out engine (`/inference/v1/generate`). + +```python +import httpx + +MODEL = "meta-llama/Llama-3.2-1B-Instruct" +RENDER = "http://localhost:8100" # vllm launch render ... +ENGINE = "http://localhost:8200" # token-in / token-out engine + +chat_request = { + "model": MODEL, + "messages": [{"role": "user", "content": "What is 2+2?"}], + "max_tokens": 32, +} + +with httpx.Client(timeout=60.0) as client: + # 1. Render: request -> token IDs (GPU less) + generate_request = client.post( + f"{RENDER}/v1/chat/completions/render", json=chat_request + ).json() + prompt_tokens = len(generate_request["token_ids"]) + + # 2. Generate: token IDs -> token IDs (token-in / token-out engine) + generate_response = client.post( + f"{ENGINE}/inference/v1/generate", json=generate_request + ).json() + + # 3. Derender: token IDs -> ChatCompletionResponse (GPU less) + response = client.post( + f"{RENDER}/v1/chat/completions/derender", + json={ + "model": MODEL, + "generate_response": generate_response, + "prompt_tokens": prompt_tokens, + "chat_request": chat_request, + }, + ).json() + +print(response["choices"][0]["message"]["content"]) +``` + +Passing `chat_request` lets the derenderer run the configured tool and reasoning parsers. This means `response["choices"][0]["message"]` carries the same `content` / `reasoning` / `tool_calls` split a `vllm serve` server would produce. Omit `chat_request` for plain detokenization only. diff --git a/docs/serving/online_serving/renderer.md b/docs/serving/online_serving/renderer.md index 9ea2f369db81..517d50f2d302 100644 --- a/docs/serving/online_serving/renderer.md +++ b/docs/serving/online_serving/renderer.md @@ -12,3 +12,5 @@ Our renderer API is designed to disaggregate the render phase(preprocessing) and - Render completion requests - [Chat Completions Render API](renderer.md) (`/v1/chat/completions/render`) - Render chat completions + +For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md). diff --git a/docs/usage/security.md b/docs/usage/security.md index ee49374e7b91..2a2e1e886d64 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -155,8 +155,10 @@ When `--api-key` is configured, the following `/v1` endpoints require Bearer tok - `/v1/chat/completions` - Chat completions - `/v1/chat/completions/batch` - Batch chat completions - `/v1/chat/completions/render` - Render chat completion requests +- `/v1/chat/completions/derender` - Derender chat completion requests - `/v1/completions` - Text completions - `/v1/completions/render` - Render completion requests +- `/v1/completions/derender` - Derender completion requests - `/v1/embeddings` - Generate embeddings - `/v1/audio/transcriptions` - Audio transcription - `/v1/audio/translations` - Audio translation diff --git a/examples/features/kv_events/kv_events_subscriber.py b/examples/features/kv_events/kv_events_subscriber.py index cfe131f000d3..eaf27296c980 100644 --- a/examples/features/kv_events/kv_events_subscriber.py +++ b/examples/features/kv_events/kv_events_subscriber.py @@ -42,12 +42,16 @@ class BlockStored(KVCacheEvent): """ group_idx: int | None = None + kv_cache_spec_kind: str | None = None + kv_cache_spec_sliding_window: int | None = None + locality: str | None = None class BlockRemoved(KVCacheEvent): block_hashes: list[ExternalBlockHash] medium: str | None group_idx: int | None = None + locality: str | None = None class AllBlocksCleared(KVCacheEvent): diff --git a/pyproject.toml b/pyproject.toml index 3819ad7fc8e7..04f1df204ba8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,7 +122,8 @@ python = "./.venv" [tool.typos.files] # these files may be written in non english words extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*", - "benchmarks/sonnet.txt", "tests/lora/data/*", "build/*", + "benchmarks/sonnet.txt", "rust/src/bench/src/datasets/sonnet.txt", + "tests/lora/data/*", "build/*", "examples/pooling/token_embed/*", "tests/models/language/pooling/*", "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py", diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 91a57997684d..3b8ff816f180 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -11,9 +11,12 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor torchcodec >= 0.14 PyNvVideoCodec==2.0.4 # FlashInfer should be updated together with the Dockerfile -flashinfer-python==0.6.13 -flashinfer-cubin==0.6.13 -apache-tvm-ffi==0.1.9 +# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from +# install_requires so the published wheel does not carry an unresolvable pin +--extra-index-url https://flashinfer.ai/whl/ +flashinfer-python==0.6.14 +flashinfer-cubin==0.6.14 +apache-tvm-ffi==0.1.10 tilelang==0.1.9 nvidia-cudnn-frontend>=1.19.1 # Required for LLM_NVTX_SCOPES_FOR_PROFILING=1 @@ -22,8 +25,8 @@ nvtx==0.2.15 fastsafetensors >= 0.3.2 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) -nvidia-cutlass-dsl[cu13]==4.5.2 -quack-kernels>=0.3.3 +nvidia-cutlass-dsl[cu13]==4.6.0 +quack-kernels>=0.4.0 # Required for tml-fa4 # Tokenspeed_MLA for faster mla with spec decode tokenspeed-mla==0.1.8; platform_system == "Linux" diff --git a/requirements/kv_connectors.txt b/requirements/kv_connectors.txt index ce920816db3e..c61ac49539d1 100644 --- a/requirements/kv_connectors.txt +++ b/requirements/kv_connectors.txt @@ -2,5 +2,5 @@ lmcache >= 0.3.9 # CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0 # until a fixed newer release is verified for runtime images. cupy-cuda13x < 14.1.0 -nixl == 1.3.0 +nixl == 1.3.1 mooncake-transfer-engine >= 0.3.8 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 7fa63787433c..e4da9ac06e76 100644 --- a/requirements/test/cuda.txt +++ b/requirements/test/cuda.txt @@ -43,7 +43,7 @@ anyio==4.14.1 # sse-starlette # starlette # watchfiles -apache-tvm-ffi==0.1.9 +apache-tvm-ffi==0.1.10 # via # -c requirements/cuda.txt # xgrammar diff --git a/requirements/xpu.txt b/requirements/xpu.txt index bc960e672ca6..8ae5649b20bc 100644 --- a/requirements/xpu.txt +++ b/requirements/xpu.txt @@ -18,4 +18,4 @@ torchvision torchcodec >= 0.14 # Required for the torchcodec video decoding backend auto_round_lib==0.14.1 -vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11/vllm_xpu_kernels-0.1.11-cp38-abi3-manylinux_2_28_x86_64.whl +vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11.1/vllm_xpu_kernels-0.1.11.1-cp38-abi3-manylinux_2_28_x86_64.whl diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 2a7b365ab4f9..0709ef8fc3ac 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2220,7 +2220,7 @@ checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" [[package]] name = "llm-multimodal" version = "1.7.1" -source = "git+https://github.com/smg-project/llm-multimodal?rev=7df38e53f99aefaebe86934f010aa8084ec99b2f#7df38e53f99aefaebe86934f010aa8084ec99b2f" +source = "git+https://github.com/smg-project/llm-multimodal?rev=5390032d6dc8a3e6fdc83acd320260367eb4b9b5#5390032d6dc8a3e6fdc83acd320260367eb4b9b5" dependencies = [ "anyhow", "base64 0.22.1", @@ -2235,6 +2235,7 @@ dependencies = [ "once_cell", "pkg-config", "rayon", + "realfft", "reqwest 0.13.4", "rustfft", "serde", @@ -2584,6 +2585,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", ] [[package]] @@ -3304,6 +3306,16 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.2", +] + [[package]] name = "rawpointer" version = "0.2.1" @@ -3434,6 +3446,7 @@ dependencies = [ "base64 0.22.1", "bytes", "encoding_rs", + "futures-channel", "futures-core", "futures-util", "h2", @@ -3548,6 +3561,15 @@ dependencies = [ "rustc-hash 2.1.1", ] +[[package]] +name = "rlimit" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f35ee2729c56bb610f6dba436bf78135f728b7373bdffae2ec815b2d3eb98cc3" +dependencies = [ + "libc", +] + [[package]] name = "rmp" version = "0.8.15" @@ -5438,6 +5460,38 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "vllm-bench" +version = "0.1.0" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bytes", + "chrono", + "clap", + "dirs", + "futures", + "hf-hub", + "image", + "indicatif", + "mimalloc", + "rand 0.9.2", + "rand_distr", + "rayon", + "reqwest 0.12.28", + "rlimit", + "rustc-hash 1.1.0", + "serde", + "serde_json", + "thiserror 2.0.18", + "tiktoken-rs 0.9.1", + "tokenizers", + "tokio", + "tokio-stream", + "url", + "uuid", +] + [[package]] name = "vllm-chat" version = "0.1.0" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 7c1793b448dd..492b1c5e2607 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,5 +1,6 @@ [workspace] members = [ + "src/bench", "src/chat", "src/cmd", "src/engine-core-client", @@ -32,8 +33,10 @@ base64 = "0.22.1" bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] } byteorder = "1.5.0" bytes = "1.12.0" +chrono = "0.4.42" clap = { version = "4.5.38", features = ["derive", "env"] } criterion = "0.5.1" +dirs = "6.0.0" easy-ext = "1.0.3" educe = "0.6.0" enum-as-inner = "0.7.0" @@ -50,10 +53,12 @@ hyper-util = { version = "0.1.20", features = [ "service", "tokio", ] } +image = { version = "0.25.9", default-features = false, features = ["jpeg"] } indexmap = "2.13.0" +indicatif = "0.18.4" itertools = "0.14.0" libc = "0.2.177" -llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "7df38e53f99aefaebe86934f010aa8084ec99b2f", default-features = false, features = ["native-tls"] } +llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] } mimalloc = "0.1.52" minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } minijinja-contrib = { version = "2.0", features = ["pycompat"] } @@ -71,10 +76,13 @@ prost-types = "0.14.3" pyo3 = "0.28.3" pythonize = "0.28.0" rand = "0.9.2" +rand_distr = "0.5.1" +rayon = "1.11.0" reasoning-parser = "1.2.2" reqwest = { version = "0.12.8", default-features = false, features = ["native-tls"] } reqwest-0-13 = { package = "reqwest", version = "0.13.4", default-features = false, features = ["native-tls"] } riptoken = { version = "0.3.0", default-features = false } +rlimit = "0.11.0" rmp-serde = "1.3.1" rmpv = { version = "1.3.1", features = ["with-serde"] } rustc-hash = "1.1.0" @@ -121,6 +129,7 @@ tracing = { version = "0.1.44", features = ["release_max_level_debug"] } tracing-futures = { version = "0.2.5", features = ["futures-03"] } tracing-subscriber = { version = "0.3.20", features = ["env-filter", "fmt"] } trait-set = "0.3.0" +url = "2.5.7" uuid = { version = "1.22.0", features = ["v4"] } validator = { version = "0.20.0", features = ["derive"] } vllm-chat = { path = "src/chat" } diff --git a/rust/src/bench/AGENTS.md b/rust/src/bench/AGENTS.md new file mode 100644 index 000000000000..852138060c61 --- /dev/null +++ b/rust/src/bench/AGENTS.md @@ -0,0 +1,181 @@ +# AGENTS.md + +## Project Overview + +Rust rewrite of `vllm bench serve` — a high-performance benchmark client for vLLM serving endpoints. Standalone binary, no Python dependency at runtime. + +Member crate `vllm-bench` of the `rust/` workspace. Uses workspace dependencies and lints; the workspace `[profile.release]` (thin LTO, `panic = "abort"`) applies. Note the workspace bans rustls/ring (`rust/deny.toml`) — all HTTP must stay on native-tls, which is why HF Hub downloads go through `src/hub.rs` (async hf-hub API bridged to sync) instead of hf-hub's ureq backend. + +## Build & Test + +Run from the `rust/` workspace root: + +```bash +# Build release binary (rust/target/release/vllm-bench) +cargo build -p vllm-bench --release + +# Run all tests +cargo test -p vllm-bench + +# Run ignored integration tests (requires network for tokenizer download) +cargo test -p vllm-bench -- --ignored +``` + +## Architecture + +- `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single) +- `src/cli.rs` — clap derive CLI args (~50+ flags) +- `src/config.rs` — Validated config from CLI; `GoodputConfig`, `RampUpConfig`, sampling param merging +- `src/error.rs` — `BenchError` enum (Http, Json, Tokenizer, Config, EndpointTimeout, Backend, Io) +- `src/benchmark.rs` — Core benchmark orchestrator (spawn-per-request with tokio + Semaphore; fetches speculative decoding metrics from `/metrics`) +- `src/multi_turn.rs` — Multi-turn conversation orchestrator (channel-based worker pool, sequential turns per conversation) +- `src/sweep.rs` — Concurrency/rate parameter sweep (`--sweep-max-concurrency`, `--sweep-request-rate`) +- `src/multi_run.rs` — N-run aggregation with mean/std/min/max/CV (`--num-runs`) +- `src/compare.rs` — Side-by-side diff of two result JSON files (`--compare`) +- `src/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback +- `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source) +- `src/hub.rs` — `HubRepo`: sync facade over hf-hub's async (reqwest/native-tls) API — per-download thread with its own runtime; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans +- `src/rate_control.rs` — Gamma/Poisson request scheduling + linear/exponential ramp-up +- `src/ready_checker.rs` — Endpoint readiness with retry +- `src/backends/` — Backend implementations (enum dispatch, not trait objects) + - `mod.rs` — `Backend` enum, `RequestFuncInput`/`RequestFuncOutput` (includes `messages` field for multi-turn) + - `streaming.rs` — SSE parser (`StreamedResponseHandler`) with speculative JSON parse for split TCP segments + - `openai_completions.rs` — `/v1/completions` backend + - `openai_chat.rs` — `/v1/chat/completions` backend (uses `input.messages` when set; zero-copy raw JSON payload for multimodal) + - `pooling.rs` — Non-streaming pooling/embedding backends: `openai-embeddings`, `openai-embeddings-chat`, `vllm-pooling`, `vllm-rerank` +- `src/datasets/random.rs` — Random dataset generation with rayon parallelism +- `src/datasets/random_mm.rs` — Random multimodal dataset (synthetic JPEG images, bucket config sampling, pre-serialized JSON fragments); `--enable-multimodal-chat` pre-builds the chat `messages` array at dataset time (mirrors Python's `apply_multimodal_chat_transformation`) +- `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching +- `src/datasets/sonnet.rs` and `src/datasets/sonnet.txt` — Sonnet dataset (built-in Shakespeare sonnets via `include_str!("sonnet.txt")`; controllable token length + shared prefix; mirrors Python `SonnetDataset`) +- `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache) +- `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (datasets-server API, column auto-detection) +- `src/datasets/custom.rs` — Custom JSONL dataset (`{"prompt": ..., "output_tokens": ...}` per line; `--custom-output-len -1` uses per-line output_tokens; prompts always sent raw — no client-side chat template) +- `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`) +- `src/datasets/random_rerank.rs` — Random rerank dataset (one query + batched documents per request for `vllm-rerank`; `--no-reranker` for embedding-based scoring; mirrors Python `RandomDatasetForReranking`) +- `src/datasets/multi_turn.rs` — Multi-turn synthetic generator + ShareGPT multi-turn loader (3-tier prefix sharing: global/conversation/unique-suffix; `per_turn_input_len`) +- `src/metrics/mod.rs` — `BenchmarkMetrics` and `MultiTurnMetrics` structs +- `src/metrics/calculator.rs` — TTFT/TPOT/ITL/E2EL/throughput stats, goodput SLO checking, peak concurrency, `calculate_multi_turn_metrics` +- `src/metrics/steady_state.rs` — Steady-state window detection (in-flight concurrency plateau via two-pointer start/end merge) + plateau throughput/TTFT/TPOT; gated on `--max-concurrency` set + `--request-rate inf` (closed-loop) +- `src/output/console.rs` — Terminal output matching Python format + multi-turn per-turn breakdown +- `src/output/json.rs` — JSON result file (compatible with Python schema) + multi-turn JSON with `per_turn_metrics` + +## Key Design Decisions + +- **Enum dispatch** for backends (avoids async trait object issues with `dyn`) +- **reqwest http1_only()** to match Python aiohttp behavior +- **rayon** for parallel dataset generation (key perf win over Python) +- **mimalloc** global allocator to reduce contention at 1400+ concurrency (page-agnostic; works on aarch64 64K-page kernels where jemalloc aborts with `LG_PAGE=12` builds) +- **Arc\ prompts** zero-copy sharing across tokio tasks (~3GB savings at 100k prompts with 8k-token inputs) +- **Spawn-per-request** `tokio::spawn` + `Semaphore` (matches Python asyncio pattern) +- **Speculative JSON parse** in SSE handler — detects complete JSON before `\n\n` arrives, improving TTFT/ITL accuracy when TCP segments split +- **Tokenizer fallback chain**: Local HF → Tiktoken (`.tiktoken`/`.model` + built-in encodings) → Server-side `/tokenize`+`/detokenize`. Blocking HTTP in rayon threads for server fallback. +- **hf-hub** for downloading tokenizers and datasets from HuggingFace Hub +- **Pre-serialized mm fragments** (`Arc`) for multimodal: image content stored as JSON strings, zero-copy concatenated into payload — avoids deep-cloning ~200KB+ base64 per request +- **Steady-state metrics** (default-on in closed-loop): measure throughput/TTFT/TPOT only over the saturated plateau to cut run-to-run variance at high concurrency; `steady_state` is an `Option` in JSON (`#[serde(default)]` for backward compat), null when the scope gate fails or `--no-steady-state` +- **`--prompt-token-ids`** (random dataset only): send token-ID arrays instead of text to skip server-side tokenization; also skips the token-length verification pass (counts exact by construction) +- **`--random-range-ratio`** follows Python semantics: lengths sampled uniformly from `[len*(1-r), len*(1+r)]`, default `0.0` = fixed; accepts a float in `[0,1)` or `'{"input": r1, "output": r2}'`. (The pre-2026-07 Rust-only form `[len*r, len]` with default 1.0 is rejected with a migration hint.) +- **`prompt_list`** (`Arc<[Arc]>` on `SampleRequest`/`RequestFuncInput`): multiple inputs per request for pooling backends — embeddings batches (`--random-batch-size`) send `"input": [...]`, rerank sends `[0]` as query + `[1..]` as documents +- JSON output schema must match Python `vllm bench serve` exactly + +## Common Issues + +- **localhost vs 127.0.0.1**: Some systems resolve `localhost` to IPv6 `::1` while vLLM listens on IPv4 only. Use `127.0.0.1` or the actual hostname. +- **Models without tokenizer.json** (e.g., `nvidia/Kimi-K2.5-NVFP4`): Automatically falls back to server-side tokenization. Can also use `--tokenizer` to point to a model with `tokenizer.json`. +- **usage.completion_tokens parsing**: vLLM sends final usage chunk with `"choices":[]` (empty array). The usage `if` must be separate from the choices `if` (not `else if`). + +## Typical Usage + +```bash +# Embedding benchmark (openai-embeddings, 8 inputs batched per request) +./target/release/vllm-bench \ + --backend openai-embeddings \ + --base-url http://gb200-10:30000 \ + --model BAAI/bge-large-en-v1.5 \ + --dataset-name random \ + --random-input-len 512 \ + --random-batch-size 8 \ + --num-prompts 1000 \ + --save-result + +# vLLM rerank benchmark (one query + 8 documents per request) +./target/release/vllm-bench \ + --backend vllm-rerank \ + --base-url http://gb200-10:30000 \ + --model BAAI/bge-reranker-v2-m3 \ + --dataset-name random-rerank \ + --random-input-len 512 \ + --random-batch-size 8 \ + --num-prompts 500 \ + --save-result + +# Prefix-cache stress (10 shared prefixes, 256+256 tokens) +./target/release/vllm-bench \ + --backend vllm \ + --base-url http://gb200-10:30000 \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --dataset-name prefix_repetition \ + --prefix-repetition-prefix-len 256 \ + --prefix-repetition-suffix-len 256 \ + --prefix-repetition-num-prefixes 10 \ + --num-prompts 1000 + +# Custom JSONL workload ({"prompt": ..., "output_tokens": ...} per line) +./target/release/vllm-bench \ + --backend openai-chat \ + --base-url http://gb200-10:30000 \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --dataset-name custom \ + --dataset-path workload.jsonl \ + --custom-output-len -1 \ + --num-prompts 1000 + +# Random dataset +./target/release/vllm-bench \ + --backend vllm \ + --base-url http://gb200-10:30000 \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --dataset-name random \ + --random-input-len 8192 \ + --random-output-len 1024 \ + --ignore-eos \ + --num-prompts 4096 \ + --percentile-metrics "ttft,tpot,itl,e2el" \ + --save-result \ + --max-concurrency 1400 + +# Random multimodal dataset (VLM benchmark) +./target/release/vllm-bench \ + --backend openai-chat \ + --base-url http://gb200-10:30000 \ + --model Qwen/Qwen2.5-VL-7B-Instruct \ + --dataset-name random-mm \ + --random-input-len 512 \ + --random-output-len 128 \ + --num-prompts 100 \ + --random-mm-base-items-per-request 1 \ + --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \ + --random-mm-bucket-config '{(1024, 800, 1): 1.0}' + +# HuggingFace dataset (WildChat) +./target/release/vllm-bench \ + --backend openai-chat \ + --base-url http://gb200-10:30000 \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --dataset-name hf \ + --dataset-path allenai/WildChat-4.8M \ + --hf-split train \ + --num-prompts 1000 \ + --save-result + +# HuggingFace dataset (LongBench with subset) +./target/release/vllm-bench \ + --backend openai-chat \ + --base-url http://gb200-10:30000 \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --dataset-name hf \ + --dataset-path THUDM/LongBench \ + --hf-subset narrativeqa \ + --hf-split test \ + --hf-output-len 512 \ + --num-prompts 200 +``` diff --git a/rust/src/bench/CLAUDE.md b/rust/src/bench/CLAUDE.md new file mode 100644 index 000000000000..43c994c2d361 --- /dev/null +++ b/rust/src/bench/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml new file mode 100644 index 000000000000..22c8f84418b2 --- /dev/null +++ b/rust/src/bench/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "vllm-bench" +version.workspace = true +edition.workspace = true +description = "High-performance benchmark client for vLLM serving endpoints" +license.workspace = true + +[dependencies] +anyhow.workspace = true +base64.workspace = true +bytes.workspace = true +chrono.workspace = true +clap.workspace = true +dirs.workspace = true +futures.workspace = true +hf-hub.workspace = true +image.workspace = true +indicatif.workspace = true +mimalloc.workspace = true +rand.workspace = true +rand_distr.workspace = true +rayon.workspace = true +reqwest = { workspace = true, features = ["json", "stream", "blocking", "http2"] } +rlimit.workspace = true +rustc-hash.workspace = true +serde = { workspace = true, features = ["rc"] } +serde_json = { workspace = true, features = ["raw_value"] } +thiserror.workspace = true +tiktoken-rs.workspace = true +tokenizers.workspace = true +tokio.workspace = true +tokio-stream.workspace = true +url.workspace = true +uuid.workspace = true + +[lints] +workspace = true diff --git a/rust/src/bench/README.md b/rust/src/bench/README.md new file mode 100644 index 000000000000..62681254afb0 --- /dev/null +++ b/rust/src/bench/README.md @@ -0,0 +1,810 @@ +# vllm-bench + +High-performance Rust benchmark client for vLLM serving endpoints. A drop-in replacement for `vllm bench serve` with near-instant startup, parallel dataset generation, and a fraction of the memory overhead — and no Python at runtime. + +```bash +vllm-bench --backend vllm --base-url http://127.0.0.1:8000 \ + --model --dataset-name random \ + --random-input-len 1024 --random-output-len 128 \ + --num-prompts 1000 --max-concurrency 200 +``` + +## Highlights + +- **Fast** — ~7 ms startup, single ~7 MB static binary, no Python imports. +- **Scales** — `Arc` prompt sharing + mimalloc keep memory <100 MB at 1400+ concurrency. +- **Many datasets** — `random`, `random-mm` (VLM), `sharegpt`, `sonnet`, `speed-bench`, and any HuggingFace dataset. +- **Many backends** — completions, chat, embeddings, pooling, and rerank. +- **Beyond a single run** — concurrency/rate **sweeps**, **multi-run** stats, **multi-turn** conversations, **LoRA** multi-adapter, and result **comparison**. +- **Steady-state metrics** — throughput/latency measured over the saturated plateau, excluding ramp-up and drain. +- **Parity** — JSON output schema and timing semantics match Python `vllm bench serve` exactly. + +### Performance vs. Python + +| Metric | Python | Rust | +| -------- | -------- | ------ | +| Startup time | Multi-second (import vllm + numpy + aiohttp) | ~7 ms | +| 100k random prompts (input_len=8192) | Minutes | Seconds (rayon parallelism) | +| Binary size | — | ~7 MB | +| Peak memory at 1400 concurrency | High (GIL + per-object overhead) | <100 MB (`Arc` prompt sharing) | + +## Contents + +- [Install](#install) +- [Quick Start](#quick-start) +- [Usage Examples](#usage-examples) +- [Supported Backends](#supported-backends) +- [Supported Datasets](#supported-datasets) +- [Metrics](#metrics) +- [CLI Reference](#cli-reference) +- [Tokenizer Support](#tokenizer-support) +- [Output Format](#output-format) +- [Architecture](#architecture) +- [Environment Variables](#environment-variables) + +## Install + +### Prebuilt binaries (Linux) + +```bash +curl -fsSL https://github.com/vllm-project/vllm-bench/releases/latest/download/vllm-bench-$(uname -m)-linux-musl -o vllm-bench && chmod +x vllm-bench +``` + +### With Cargo + +Install straight from the repository (builds from source; requires [Rust](https://rustup.rs/) stable and a C compiler for the native tokenizer dependency): + +```bash +cargo install --git https://github.com/vllm-project/vllm-bench vllm-bench +``` + +The trailing `vllm-bench` selects the package — the repo also ships a `mock-llm-server` binary, so omitting it fails with `multiple packages with binaries found`. The binary is installed to `~/.cargo/bin/`. + +### Build from source + +Requires [Rust](https://rustup.rs/) (stable). + +```bash +git clone https://github.com/vllm-project/vllm-bench.git +cd vllm-bench +./install.sh # builds release and installs to ~/.local/bin +# or: ./install.sh --to ~/bin +``` + +## Quick Start + +Point it at a running vLLM server and benchmark with synthetic prompts: + +```bash +vllm-bench \ + --backend vllm \ + --base-url http://127.0.0.1:8000 \ + --model \ + --dataset-name random \ + --random-input-len 1024 \ + --random-output-len 128 \ + --num-prompts 1000 \ + --max-concurrency 200 +``` + +> **Tip:** prefer `127.0.0.1` over `localhost` — some systems resolve `localhost` to IPv6 `::1` while vLLM listens on IPv4 only. + +Add `--save-result` to write a JSON file, or `--dry-run` to generate and inspect the dataset without sending any requests. + +## Usage Examples + +
+Generation (completions / chat) + +```bash +# Full production-style run with percentile metrics and result file +vllm-bench \ + --backend vllm \ + --base-url http://127.0.0.1:8000 \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --dataset-name random \ + --random-input-len 8192 \ + --random-output-len 1024 \ + --ignore-eos \ + --num-prompts 4096 \ + --percentile-metrics "ttft,tpot,itl,e2el" \ + --save-result \ + --max-concurrency 1400 + +# Send token IDs instead of text (pure vLLM: skips server-side tokenization, +# exact token counts, faster). Random dataset only. +vllm-bench \ + --backend vllm \ + --base-url http://127.0.0.1:8000 \ + --model \ + --dataset-name random \ + --random-input-len 1024 \ + --prompt-token-ids \ + --num-prompts 1000 +``` + +
+ +
+Datasets (ShareGPT / Sonnet / HuggingFace / SPEED-Bench) + +```bash +# ShareGPT (auto-downloads from HuggingFace on first run, cached afterwards) +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name sharegpt --num-prompts 500 --save-result + +# ShareGPT with an explicit local file +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name sharegpt --dataset-path /path/to/ShareGPT_V3.json \ + --num-prompts 500 --save-result + +# Sonnet — built-in Shakespeare sonnets, no dataset file needed. +# Generates prompts of a controllable token length with a shared prefix. +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name sonnet \ + --sonnet-input-len 550 --sonnet-output-len 150 --sonnet-prefix-len 200 \ + --num-prompts 500 + +# Any public HuggingFace dataset (auto-downloads, auto-detects columns) +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name hf --dataset-path allenai/WildChat-4.8M \ + --hf-split train --num-prompts 1000 --save-result + +# HuggingFace dataset with subset + fixed output length (LongBench) +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name hf --dataset-path THUDM/LongBench \ + --hf-subset narrativeqa --hf-split test --hf-output-len 512 --num-prompts 200 + +# Gated HuggingFace dataset (requires HF_TOKEN) +HF_TOKEN=hf_xxx vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name hf --dataset-path lmsys/lmsys-chat-1m \ + --hf-split train --hf-output-len 256 --num-prompts 1000 + +# SPEED-Bench for speculative decoding evaluation (auto-downloads, cached) +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name speed-bench --speed-bench-config qualitative \ + --num-prompts 200 --output-len 256 --save-result + +# SPEED-Bench throughput split with entropy category filter + input truncation +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name speed-bench --speed-bench-config throughput_16k \ + --speed-bench-max-input-len 10240 --speed-bench-category low_entropy \ + --num-prompts 500 --output-len 256 --max-concurrency 200 --save-result +``` + +
+ +
+Multimodal (VLM with synthetic images) + +```bash +# One synthetic image per request +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 \ + --model Qwen/Qwen2.5-VL-7B-Instruct \ + --dataset-name random-mm \ + --random-input-len 512 --random-output-len 128 --num-prompts 100 \ + --random-mm-base-items-per-request 1 \ + --random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \ + --random-mm-bucket-config '{(1024, 800, 1): 1.0}' + +# Multiple images per request, mixed resolutions +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 \ + --model Qwen/Qwen2.5-VL-7B-Instruct \ + --dataset-name random-mm \ + --random-input-len 256 --random-output-len 128 --num-prompts 50 \ + --random-mm-base-items-per-request 3 \ + --random-mm-limit-mm-per-prompt '{"image": 5, "video": 0}' \ + --random-mm-bucket-config '{(256,256,1): 0.5, (720,1280,1): 0.5}' +``` + +
+ +
+Embedding / Pooling / Rerank + +```bash +# Text embedding +vllm-bench \ + --backend openai-embeddings --base-url http://127.0.0.1:8000 \ + --model BAAI/bge-large-en-v1.5 \ + --dataset-name random --random-input-len 512 --num-prompts 1000 \ + --max-concurrency 200 --save-result + +# Chat-format embedding (supports multimodal content) +vllm-bench \ + --backend openai-embeddings-chat --base-url http://127.0.0.1:8000 \ + --model BAAI/bge-large-en-v1.5 \ + --dataset-name sharegpt --num-prompts 500 --save-result + +# vLLM native pooling endpoint +vllm-bench \ + --backend vllm-pooling --base-url http://127.0.0.1:8000 \ + --model BAAI/bge-large-en-v1.5 \ + --dataset-name random --random-input-len 256 --num-prompts 1000 --save-result + +# Rerank (query from dataset, documents via --extra-body) +vllm-bench \ + --backend vllm-rerank --base-url http://127.0.0.1:8000 \ + --model BAAI/bge-reranker-v2-m3 \ + --dataset-name sharegpt --num-prompts 500 \ + --extra-body '{"documents": ["document to rerank"]}' --save-result +``` + +
+ +
+Rate control, ramp-up & goodput + +```bash +# Ramp from 10 → 100 RPS with goodput SLO tracking +vllm-bench \ + --backend vllm --base-url http://127.0.0.1:8000 --model \ + --num-prompts 2000 \ + --ramp-up-strategy linear --ramp-up-start-rps 10 --ramp-up-end-rps 100 \ + --goodput ttft:200 e2el:5000 \ + --save-result + +# Fixed Poisson arrival rate at 50 RPS +vllm-bench \ + --backend vllm --base-url http://127.0.0.1:8000 --model \ + --num-prompts 2000 --request-rate 50 --burstiness 1.0 +``` + +
+ +
+Sweep — find the optimal concurrency / rate + +```bash +# Sweep over concurrency values +vllm-bench \ + --backend vllm --base-url http://127.0.0.1:8000 --model \ + --num-prompts 500 \ + --sweep-max-concurrency 1,10,50,100,200,500,1000 + +# Sweep over request rates +vllm-bench \ + --backend vllm --base-url http://127.0.0.1:8000 --model \ + --num-prompts 500 \ + --sweep-request-rate 1,10,50,100,inf + +# Scale work with concurrency and reset the prefix cache between points +# (--sweep-num-prompts-factor sets num_prompts = concurrency * factor; +# --reset-prefix-cache requires VLLM_SERVER_DEV_MODE=1 on the server) +vllm-bench \ + --backend vllm --base-url http://127.0.0.1:8000 --model \ + --sweep-max-concurrency 1,10,50,100 \ + --sweep-num-prompts-factor 20 \ + --reset-prefix-cache +``` + +
+ +
+Multi-run & comparison + +```bash +# Run 5 times, report mean/std/min/max with coefficient of variation +vllm-bench \ + --backend vllm --base-url http://127.0.0.1:8000 --model \ + --num-prompts 1000 --max-concurrency 200 --num-runs 5 + +# Compare two saved result files side-by-side (no server needed) +vllm-bench --compare baseline.json optimized.json +``` + +
+ +
+Multi-turn conversations + +```bash +# Synthetic multi-turn (controllable per-turn token lengths) +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name random --multi-turn --multi-turn-num-turns 5 \ + --random-input-len 512 --random-output-len 256 \ + --num-prompts 50 --multi-turn-concurrency 10 \ + --percentile-metrics "ttft,tpot,itl,e2el" --save-result + +# Variable turn count per conversation + per-turn input length for turns 1+ +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name random --multi-turn \ + --multi-turn-min-turns 2 --multi-turn-max-turns 8 \ + --random-input-len 2048 --per-turn-input-len 256 --random-output-len 128 \ + --num-prompts 100 --multi-turn-concurrency 20 + +# ShareGPT conversations (loads all turns, not just the first two) +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --dataset-name sharegpt --multi-turn \ + --num-prompts 50 --multi-turn-concurrency 10 --save-result + +# Think time between turns +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 --model \ + --multi-turn --multi-turn-num-turns 3 --multi-turn-delay-ms 500 \ + --num-prompts 100 --multi-turn-concurrency 20 +``` + +
+ +
+LoRA multi-adapter + +```bash +# Distribute requests across N adapters registered on the server. +# --model stays the BASE model (tokenizer / readiness / /tokenize use it); +# the per-request `model` field is rewritten to one of --lora-modules. +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 \ + --model Qwen/Qwen3-30B-A3B \ + --lora-modules sql-lora-1 sql-lora-2 sql-lora-3 sql-lora-4 \ + sql-lora-5 sql-lora-6 sql-lora-7 sql-lora-8 \ + --lora-assignment random \ + --dataset-name random --random-input-len 1024 --random-output-len 256 \ + --num-prompts 1000 --max-concurrency 64 --save-result + +# Deterministic round-robin assignment (request i -> adapter[i % N]) +vllm-bench \ + --backend openai-chat --base-url http://127.0.0.1:8000 \ + --model Qwen/Qwen3-30B-A3B \ + --lora-modules sql-lora-1 sql-lora-2 sql-lora-3 sql-lora-4 \ + --lora-assignment round-robin \ + --dataset-name random --num-prompts 1000 +``` + +Server side — start vLLM with `--enable-lora` and one `name=path` pair per adapter: + +```bash +vllm serve \ + --enable-lora --max-loras 8 --max-lora-rank 16 \ + --lora-modules \ + sql-lora-1=jeeejeee/qwen3-moe-text2sql-spider \ + sql-lora-2=jeeejeee/qwen3-moe-text2sql-spider \ + ... +``` + +Set `--max-loras` ≥ number of adapter names to keep them all resident (clean steady-state numbers), or lower to stress the LoRA swap path. + +
+ +
+Profiling & dry-run + +```bash +# Trigger vLLM server-side profiling (start before, stop after the benchmark) +vllm-bench \ + --backend vllm --base-url http://127.0.0.1:8000 --model \ + --num-prompts 100 --profile + +# Defer profiling until the server batch is full, then capture for 10s +vllm-bench \ + --backend vllm --base-url http://127.0.0.1:8000 --model \ + --num-prompts 2000 --max-concurrency 256 \ + --profile --profile-batch-threshold 200 --profile-duration 10 + +# Dry run: generate dataset, print stats, send nothing +vllm-bench \ + --model --num-prompts 100000 --random-input-len 8192 --dry-run +``` + +
+ +## Supported Backends + +### Generation + +| Backend | API Endpoint | Description | +| --------- | ------------- | ------------- | +| `vllm` / `openai` | `/v1/completions` | OpenAI-compatible completions (streaming) | +| `openai-chat` | `/v1/chat/completions` | OpenAI-compatible chat completions (streaming, multimodal) | + +### Embedding / Pooling + +| Backend | API Endpoint | Description | +| --------- | ------------- | ------------- | +| `openai-embeddings` | `/v1/embeddings` | Text embedding (accepts text or token IDs) | +| `openai-embeddings-chat` | `/v1/embeddings` | Chat-format embedding (supports multimodal content) | +| `vllm-pooling` | `/v1/pooling` | vLLM native pooling endpoint | +| `vllm-rerank` | `/v1/rerank` | vLLM reranking (query from prompt, documents via `--extra-body`) | + +Pooling backends are non-streaming and report E2EL (end-to-end latency) only. Use `--dataset-name sharegpt`, `sonnet`, or `hf` for text-based embedding/rerank benchmarks, or `random` for token-ID-based embedding benchmarks. + +## Supported Datasets + +| Dataset | Description | +| --------- | ------------- | +| `random` | Synthetic prompts with exact token-length matching (default) | +| `random-mm` | Synthetic multimodal prompts with random JPEG images for VLM benchmarking (requires `openai-chat`) | +| `sharegpt` | Real conversations from ShareGPT (auto-downloads from HuggingFace, or use `--dataset-path`) | +| `sonnet` | Built-in Shakespeare sonnets; controllable token length + shared prefix, no dataset file needed | +| `speed-bench` | NVIDIA SPEED-Bench for speculative decoding evaluation (auto-downloads, 11 categories) | +| `hf` | Any HuggingFace dataset (auto-downloads via datasets-server API, auto-detects chat/text columns) | + +## Metrics + +### Generation backends + +- **TTFT** (Time to First Token) — latency from request send to first token received +- **TPOT** (Time per Output Token) — average time between output tokens +- **ITL** (Inter-Token Latency) — per-token latency distribution +- **E2EL** (End-to-End Latency) — total request latency +- **Throughput** — requests/sec, output tokens/sec, peak output tokens/sec, total tokens/sec +- **Concurrency** — peak concurrent requests +- **Goodput** — requests/sec meeting all specified SLOs (with `--goodput`) + +### Pooling / embedding backends + +- **E2EL** — total request latency (mean, median, std, percentiles) +- **Throughput** — requests/sec, input tokens/sec +- **Concurrency** — peak concurrent requests + +### Steady-state metrics + +When `--max-concurrency` is set and `--request-rate` is `inf` (closed-loop mode), the benchmark automatically reports an additional **Steady-State Metrics** block. It measures throughput and latency only over the window during which in-flight concurrency stays at or above a fraction of `--max-concurrency`, excluding the ramp-up and drain phases. This sharply reduces run-to-run variance at very high concurrency. + +The block reports request/input/output/total token throughput plus TTFT (mean, median, percentiles) and TPOT (mean, median, P90, P99) over the detected plateau, along with the window bounds and how many requests fell inside it. Tune it with `--steady-state-threshold` (default `0.95`) and `--steady-state-min-window`, or disable with `--no-steady-state`. The result JSON carries a `steady_state` object (null when not computed). + +## CLI Reference + +Run `vllm-bench --help` for the authoritative list. Grouped reference below. + +
+Server connection + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--backend` | `openai` | Backend type (`vllm`, `openai`, `openai-chat`, `openai-embeddings`, `openai-embeddings-chat`, `vllm-pooling`, `vllm-rerank`) | +| `--base-url` | — | Server base URL (overrides `--host`/`--port`) | +| `--host` | `127.0.0.1` | Server host | +| `--port` | `8000` | Server port | +| `--endpoint` | Auto | API endpoint path (auto-selected per backend) | +| `--insecure` | `false` | Disable SSL certificate verification | + +
+ +
+Model & tokenizer + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--model` | Auto-detect | Model name (fetched from `/v1/models` if omitted) | +| `--served-model-name` | — | Model name used in API requests | +| `--tokenizer` | Same as model | Tokenizer name or path (supports HF, tiktoken, server fallback) | +| `--tokenizer-mode` | `auto` | Tokenizer mode (`auto`, `hf`, `slow`, `mistral`) | +| `--trust-remote-code` | `false` | Trust remote code for tokenizer | +| `--skip-tokenizer-init` | `false` | Skip tokenizer initialization | + +
+ +
+Dataset + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--dataset-name` | `random` | Dataset type (`random`, `random-mm`, `sharegpt`, `sonnet`, `speed-bench`, `hf`) | +| `--dataset-path` | — | Path to dataset file (optional for `sharegpt`/`sonnet`, which auto-source) | +| `--num-prompts` | `1000` | Number of prompts to generate (conversations in multi-turn mode) | +| `--max-model-len` | — | Filter out requests where `prompt_len + output_len` exceeds this context length | +| `--input-len` | — | Override input length (general) | +| `--output-len` | — | Override output length (general) | +| `--no-oversample` | `false` | Don't oversample if dataset is smaller than `--num-prompts` | +| `--disable-shuffle` | `false` | Don't shuffle the dataset | +| `--seed` | `0` | Random seed for reproducibility | +| **Random** | | | +| `--random-input-len` | `1024` | Input token length | +| `--random-output-len` | `128` | Output token length | +| `--random-prefix-len` | `0` | Shared prefix length | +| `--random-range-ratio` | `1.0` | Length jitter, range `(0, 1]`. Lengths sampled from `[ratio × target, target]`; `1.0` = fixed length | +| `--prompt-token-ids` | `false` | Send prompts as token-ID arrays (skips server-side tokenization, exact counts). Random dataset only | +| **Random multimodal** | | | +| `--random-mm-base-items-per-request` | `1` | Base number of multimodal items (images) per request | +| `--random-mm-num-mm-items-range-ratio` | `0.0` | Range ratio for varying item count per request | +| `--random-mm-limit-mm-per-prompt` | `{"image": 255, "video": 1}` | Per-modality hard caps (JSON) | +| `--random-mm-bucket-config` | `{(256,256,1): 0.5, (720,1280,1): 0.5}` | `(height,width,frames)` → probability (Python tuple syntax; frames=1 = image) | +| **ShareGPT** | | | +| `--sharegpt-output-len` | — | Override output length | +| **Sonnet** | | | +| `--sonnet-input-len` | `550` | Input tokens per request | +| `--sonnet-output-len` | `150` | Output tokens per request | +| `--sonnet-prefix-len` | `200` | Prefix tokens shared across requests | +| **SPEED-Bench** | | | +| `--speed-bench-config` | `qualitative` | Split (`qualitative`, `throughput_1k`/`2k`/`8k`/`16k`/`32k`) | +| `--speed-bench-category` | — | Filter by category (`low_entropy`, `high_entropy`, `mixed_entropy`, `coding`, `math`, …) | +| `--speed-bench-max-input-len` | — | Truncate prompts to at most N tokens | +| **HuggingFace** | | | +| `--hf-split` | Auto | Split (`train`, `test`, `validation`); auto-detected if omitted | +| `--hf-subset` | — | Subset/config name (e.g. `narrativeqa` for LongBench) | +| `--hf-output-len` | — | Fixed output length for all requests (overrides dataset-derived length) | +| `--hf-text-column` | Auto | Column containing prompt text; auto-detected from common patterns | + +
+ +
+Rate control + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--request-rate` | `inf` | Requests per second (`inf` = all at once) | +| `--burstiness` | `1.0` | Burstiness factor (1.0 = Poisson, >1 = bursty) | +| `--max-concurrency` | `num-prompts` | Maximum concurrent requests (semaphore) | +| `--ramp-up-strategy` | — | Ramp-up mode (`linear` or `exponential`) | +| `--ramp-up-start-rps` | — | Starting request rate for ramp-up | +| `--ramp-up-end-rps` | — | Ending request rate for ramp-up | + +
+ +
+Sampling parameters + +| Flag | Description | +| ------ | ------------- | +| `--temperature` | Temperature (server default if omitted) | +| `--top-p` | Top-p (nucleus) sampling | +| `--top-k` | Top-k sampling | +| `--min-p` | Min-p sampling | +| `--frequency-penalty` | Frequency penalty | +| `--presence-penalty` | Presence penalty | +| `--repetition-penalty` | Repetition penalty | + +Merged into the request body. Only effective with generation backends (`vllm`, `openai`, `openai-chat`); ignored by pooling/embedding backends. + +
+ +
+Output & results + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--save-result` | `false` | Save results to JSON file | +| `--save-detailed` | `false` | Include per-request data in JSON (input/output lens, ITLs, texts) | +| `--append-result` | `false` | Append to existing JSON file (JSONL format) | +| `--result-dir` | — | Directory for result files | +| `--result-filename` | Auto | Custom result filename | +| `--percentile-metrics` | `ttft,tpot,itl,e2el` | Metrics for percentile reporting (pooling defaults to `e2el` only) | +| `--metric-percentiles` | `99` | Percentile values to compute | +| `--sweep-summary-percentiles` | — | Extra percentiles for sweep summary tables (auto-added to computed set) | +| `--goodput` | — | SLO pairs for goodput (`ttft:100 tpot:50 e2el:500`, values in ms) | +| `--disable-tqdm` | `false` | Disable progress bar | +| `--label` | — | Label prefix for result files | +| `--metadata` | — | Key-value metadata (`KEY=VALUE`, repeatable) | + +
+ +
+Request options + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--ignore-eos` | `false` | Ignore EOS token (force full output length) | +| `--logprobs` | — | Number of logprobs per token | +| `--num-warmups` | `0` | Warmup requests before benchmarking | +| `--ready-check-timeout-sec` | `0` | Endpoint readiness timeout (0 = skip) | +| `--request-id-prefix` | Auto (UUID) | Prefix for request IDs | +| `--header` | — | Extra headers (`KEY=VALUE`, repeatable) | +| `--extra-body` | — | Extra JSON body parameters | +| `--dry-run` | `false` | Generate dataset only, skip benchmark | + +
+ +
+Steady-state metrics + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--steady-state-threshold` | `0.95` | Fraction of `--max-concurrency` at which the steady-state window opens, range (0, 1] | +| `--steady-state-min-window` | Auto | Minimum window duration (s) below which a warning is attached. Default `max(10, 0.1 × run_duration)` | +| `--no-steady-state` | `false` | Disable steady-state metrics computation | + +Computed only when `--max-concurrency` is set and `--request-rate` is `inf`. + +
+ +
+Profiling + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--profile` | `false` | Trigger vLLM server-side profiling (`/start_profile` before, `/stop_profile` after) | +| `--profile-batch-threshold` | — | Defer profiling until `/metrics` reports ≥ N running requests, then capture. Requires `--profile` | +| `--profile-duration` | `5.0` | Seconds to capture once the batch threshold is reached. Requires `--profile-batch-threshold` | + +
+ +
+Sweep mode + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--sweep-max-concurrency` | — | Comma-separated concurrency values to sweep (e.g. `1,10,50,100,500`) | +| `--sweep-request-rate` | — | Comma-separated rate values to sweep, supports `inf` (e.g. `1,10,100,inf`) | +| `--sweep-num-prompts-factor` | — | Set `num_prompts = concurrency × factor` per concurrency sweep point | +| `--reset-prefix-cache` | `false` | Reset the server's prefix cache before each sweep iteration (requires `VLLM_SERVER_DEV_MODE=1`) | + +Runs the benchmark once per value, then prints a summary table comparing throughput and latency across all sweep points and identifies the best-throughput configuration. Works in multi-turn mode too. `--sweep-summary-percentiles` appends extra TTFT/TPOT/E2EL columns to the summary, auto-adding any missing percentiles to the computed set so they also appear in result JSON. + +
+ +
+Multi-turn conversation benchmark + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--multi-turn` | `false` | Enable multi-turn conversation mode (requires `--backend openai-chat`) | +| `--multi-turn-num-turns` | `3` | Turns per conversation (synthetic mode) | +| `--multi-turn-min-turns` | `0` | Minimum turns per conversation (0 = use `--multi-turn-num-turns`) | +| `--multi-turn-max-turns` | `0` | Maximum turns per conversation (0 = `--multi-turn-num-turns` synthetic / uncapped ShareGPT) | +| `--multi-turn-concurrency` | — | Concurrent conversations (defaults to `--max-concurrency` or `--num-prompts`) | +| `--multi-turn-delay-ms` | `0` | Delay between turns in ms (simulates user think time) | +| `--per-turn-input-len` | `0` | Input token length for turns 1+ (0 = use `--random-input-len` for all turns) | +| `--multi-turn-prefix-global-ratio` | `0.0` | Fraction of per-turn input shared across all conversations (random dataset only) | +| `--multi-turn-prefix-conversation-ratio` | `0.0` | Fraction shared within each conversation (random dataset only) | + +With `--multi-turn`, `--num-prompts` controls the number of **conversations**, not individual requests. + +**How it works:** + +- Turn 1: send `[user_1]`, get `assistant_1` +- Turn 2: send `[user_1, assistant_1, user_2]`, get `assistant_2` +- Turn N: send full history + `user_N` — measures growing-context performance + +**Data sources:** + +- `--dataset-name random` — synthetic conversations with controllable per-turn token lengths. Auto-sets `min_tokens` to enforce output length without `ignore_eos`. +- `--dataset-name sharegpt` — loads all turns (not just the first two); filters for entries with ≥ 2 real turns. + +**Prefix sharing** (random dataset): when `--multi-turn-prefix-global-ratio` or `--multi-turn-prefix-conversation-ratio` is > 0, each turn sends a fixed-length message (no history accumulation) composed of a global prefix + per-conversation prefix + unique suffix. The two ratios must sum to < 1.0. + +**Router affinity:** every turn sends `X-Session-ID: {conversation_id}` for KV-cache reuse behind a vLLM router. + +**Output:** overall metrics plus a per-turn breakdown (TTFT/TPOT/ITL/E2EL by turn index). Expect TTFT to climb across turns due to growing context. JSON includes a `per_turn_metrics` array. + +
+ +
+LoRA multi-adapter + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--lora-modules` | — | Adapter names registered on the server (`vllm serve --lora-modules name=path`). Each request's `model` field is rewritten to one of these. Repeatable | +| `--lora-assignment` | `random` | Distribution: `random` (uniform, seeded by `--seed`) or `round-robin` (deterministic `i % N`) | + +`--model` must stay the **base** model — its tokenizer builds prompts, and `/v1/models`, `/tokenize`, ready check, and warmup all use it. Only the per-request `model` field in completions/chat payloads is rewritten to the assigned adapter (vLLM routes by name). + +**Assignment scope:** per request in single-shot mode; **per conversation** (sticky across all turns) in multi-turn mode, to avoid breaking prefix-cache reuse mid-dialog. + +**Reproducibility:** with `--lora-assignment random`, the same `--seed` + same `--lora-modules` list yields identical request-to-adapter mappings. Pooling/embedding backends are rejected — LoRA routing applies to generative paths only. + +
+ +
+Multi-run & comparison + +| Flag | Default | Description | +| ------ | --------- | ------------- | +| `--num-runs` | `1` | Run benchmark N times; report mean/std/min/max with CV | +| `--compare` | — | Compare two result JSON files side-by-side (skips benchmarking) | + +`--num-runs` aggregates metrics across runs and reports the coefficient of variation (CV) for throughput stability. `--compare` reads two previously-saved result files and prints a diff with delta, % change, and improvement/regression markers. + +
+ +## Tokenizer Support + +Tokenizers are loaded with a three-tier fallback chain: + +1. **Local HuggingFace** — `tokenizer.json` from a local path or the Hub (fastest) +2. **Tiktoken** — `.tiktoken` / `.model` format for Kimi, Qwen, etc. (auto-extracts `pat_str` from Python source) +3. **Server-side** — falls back to vLLM's `/tokenize` + `/detokenize` endpoints + +For the `random` dataset, prompt token lengths are verified against the server on the first run and cached; subsequent runs with the same model+server skip verification. Verification is also skipped when `--prompt-token-ids` is set (token counts are exact by construction). + +Models without `tokenizer.json` (e.g. `nvidia/Kimi-K2.5-NVFP4`) fall back to server-side tokenization automatically; you can also point `--tokenizer` at a model that ships `tokenizer.json`. + +## Output Format + +JSON output is compatible with the `vllm bench serve` Python schema. Result files are named: + +```text +{label}-{rate}qps-concurrency{max_concurrency}-{model}-{timestamp}.json +``` + +Use `--append-result` to append multiple runs to the same file in JSONL format. `--save-detailed` adds per-request arrays (input/output lengths, ITLs, generated text). + +## Architecture + +
+Source layout + +```text +src/ +├── main.rs # Entry point, mimalloc, tokio runtime, mode dispatch +├── cli.rs # clap CLI argument definitions +├── config.rs # Validated config, goodput/ramp-up parsing +├── benchmark.rs # Core orchestrator (schedule, spawn, collect, verify, profile) +├── multi_turn.rs # Multi-turn conversation orchestrator (channel workers) +├── compare.rs # Result diff (--compare file_a.json file_b.json) +├── sweep.rs # Parameter sweep (--sweep-max-concurrency, --sweep-request-rate) +├── multi_run.rs # Multi-run statistics (--num-runs N) +├── rate_control.rs # Gamma/Poisson scheduling + linear/exponential ramp-up +├── ready_checker.rs # Endpoint readiness with retry +├── tokenizer.rs # Tokenizer abstraction (HF, tiktoken, server) +├── tiktoken.rs # Tiktoken BPE loader with pat_str extraction +├── error.rs # Error types +├── backends/ +│ ├── mod.rs # Backend enum dispatch, typed SSE structs +│ ├── streaming.rs # SSE stream parser with speculative JSON parse +│ ├── openai_completions.rs # /v1/completions backend +│ ├── openai_chat.rs # /v1/chat/completions backend +│ └── pooling.rs # Embedding/pooling/rerank backends (non-streaming) +├── datasets/ +│ ├── mod.rs # SampleRequest, ConversationTurn, MultiTurnConversation types +│ ├── random.rs # Random dataset with rayon parallelism +│ ├── random_mm.rs # Random multimodal dataset (JPEG generation, bucket sampling) +│ ├── multi_turn.rs # Multi-turn synthetic + ShareGPT conversation generators +│ ├── sharegpt.rs # ShareGPT JSON dataset loader +│ ├── sonnet.rs # Sonnet dataset (built-in Shakespeare sonnets) +│ ├── speed_bench.rs # NVIDIA SPEED-Bench loader (auto-download + cache) +│ └── hf_dataset.rs # Generic HuggingFace dataset (auto-download, column detection) +├── metrics/ +│ ├── mod.rs # BenchmarkMetrics, MultiTurnMetrics structs +│ ├── calculator.rs # Percentile/throughput/goodput/peak/multi-turn computation +│ └── steady_state.rs # Steady-state window detection + plateau metrics +└── output/ + ├── mod.rs + ├── console.rs # Terminal output (matches Python format) + └── json.rs # JSON result serialization (Python-compatible schema) +``` + +
+ +### Key design decisions + +- **reqwest + tokio** — HTTP client with connection pooling, forced HTTP/1.1, TCP_NODELAY to match Python's aiohttp and avoid Nagle latency inflation on TTFT +- **mimalloc** — global allocator to reduce contention under high concurrency (1400+ tasks); page-agnostic, runs on aarch64 4K- and 64K-page kernels +- **`Arc` prompts** — zero-copy prompt sharing across tokio tasks, eliminating ~3 GB peak memory at 100k requests with 8k-token prompts +- **Spawn-per-request** — `tokio::spawn` per request with a `Semaphore` for concurrency control (matches Python's asyncio pattern) +- **rayon** — parallel dataset generation across CPU cores (200–500× faster than Python for 100k+ prompts) +- **Enum dispatch** — backend variants instead of trait objects (avoids async trait-object limitations) +- **Typed SSE deserialization** — `CompletionChunk`/`ChatChunk` structs skip unused JSON fields (cheaper than `serde_json::Value`) +- **Speculative JSON parse** — SSE handler uses `serde_json::value::RawValue` to detect complete JSON before `\n\n` arrives, improving TTFT/ITL accuracy when TCP segments split +- **Connection error retry** — automatic retry with backoff on connection reset/timeout/refused (up to 3 attempts) +- **Tokenizer verification cache** — server-side token-length verification is cached per model+server pair + +### Behavioral parity with Python + +The Rust implementation matches Python `vllm bench serve` in: + +- SSE streaming protocol handling (including speculative parse for split TCP segments) +- Timing semantics (monotonic `Instant` matching Python's `time.perf_counter()`) +- Chat vs. completions differences (`max_completion_tokens` vs. `max_tokens`, Content-Type, timestamp placement) +- JSON output schema (all fields, key naming, `request_rate` as the string `"inf"`) +- Rate control (Gamma distribution, normalization, burstiness, linear/exponential ramp-up) +- Metrics (TTFT/TPOT/ITL/E2EL percentiles, peak tokens/sec, peak concurrency, goodput) +- Sampling parameters merged into the request body via `extra_body` (same precedence rules) + +## Environment Variables + +| Variable | Description | +| ---------- | ------------- | +| `OPENAI_API_KEY` | API key for authenticated endpoints (cached, not read per-request) | +| `HF_TOKEN` | HuggingFace token for gated model tokenizers and gated datasets | +| `TOKIO_WORKER_THREADS` | Override tokio worker thread count (default: physical cores) | + +## License + +Apache-2.0 + + diff --git a/rust/src/bench/src/backends/mod.rs b/rust/src/bench/src/backends/mod.rs new file mode 100644 index 000000000000..698a20e49a29 --- /dev/null +++ b/rust/src/bench/src/backends/mod.rs @@ -0,0 +1,213 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +pub mod openai_chat; +pub mod openai_completions; +pub mod pooling; +pub mod streaming; + +use std::collections::HashMap; +use std::sync::Arc; + +use serde::{Deserialize, Serialize}; + +// --- Typed SSE chunk structs for zero-alloc deserialization --- +// Using typed deserialization avoids building a full serde_json::Value tree. +// Only the fields we need are extracted; everything else is skipped by serde. + +/// Completions API streaming chunk (minimal fields). +#[derive(Deserialize)] +pub struct CompletionChunk { + #[serde(default)] + pub choices: Vec, + pub usage: Option, +} + +#[derive(Deserialize)] +pub struct CompletionChoice { + pub text: Option, +} + +/// Chat API streaming chunk (minimal fields). +#[derive(Deserialize)] +pub struct ChatChunk { + #[serde(default)] + pub choices: Vec, + pub usage: Option, +} + +#[derive(Deserialize)] +pub struct ChatChoice { + pub delta: Option, +} + +#[derive(Deserialize)] +pub struct ChatDelta { + pub content: Option, +} + +#[derive(Deserialize)] +pub struct ChunkUsage { + pub completion_tokens: Option, +} + +use crate::cli::BackendKind; +use crate::error::Result; + +/// Input for a single benchmark request. +#[derive(Debug, Clone)] +pub struct RequestFuncInput { + pub prompt: Arc, + pub api_url: String, + pub prompt_len: usize, + pub output_len: usize, + pub model: String, + pub model_name: Option, + pub logprobs: Option, + pub extra_headers: Option>, + pub extra_body: Option, + pub ignore_eos: bool, + pub request_id: Option, + /// Pre-built messages array for multi-turn conversations. + /// When set, the chat backend uses this instead of building from `prompt`. + pub messages: Option, + /// Pre-computed token IDs for this prompt. + /// When set, the completions backend sends these directly via `prompt_token_ids` + /// instead of the text `prompt`, skipping server-side tokenization. + pub prompt_token_ids: Option>, + /// Multimodal content as pre-serialized JSON fragments. + /// When set, the chat backend concatenates these directly into the payload bytes, + /// avoiding any parsing or deep-cloning of base64 image data. + pub multi_modal_content: Option]>>, + /// Complete pre-serialized chat `messages` array (--enable-multimodal-chat). + /// When set, the chat backend splices it verbatim into the payload bytes, + /// taking precedence over `messages`, `prompt`, and `multi_modal_content`. + pub chat_messages_json: Option>, + /// Multiple text inputs for one request (pooling backends only): + /// embeddings batch (`"input": [...]`) or rerank query+documents. + pub prompt_list: Option]>>, +} + +/// Output from a single benchmark request including timing metrics. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RequestFuncOutput { + pub generated_text: String, + pub success: bool, + pub latency: f64, + pub output_tokens: usize, + pub ttft: f64, + pub itl: Vec, + pub tpot: f64, + pub prompt_len: usize, + pub error: String, + pub start_time: f64, +} + +impl Default for RequestFuncOutput { + fn default() -> Self { + Self { + generated_text: String::new(), + success: false, + latency: 0.0, + output_tokens: 0, + ttft: 0.0, + itl: Vec::new(), + tpot: 0.0, + prompt_len: 0, + error: String::new(), + start_time: 0.0, + } + } +} + +impl Default for RequestFuncInput { + fn default() -> Self { + Self { + prompt: Arc::from(""), + api_url: String::new(), + prompt_len: 0, + output_len: 0, + model: String::new(), + model_name: None, + logprobs: None, + extra_headers: None, + extra_body: None, + ignore_eos: false, + request_id: None, + messages: None, + prompt_token_ids: None, + multi_modal_content: None, + chat_messages_json: None, + prompt_list: None, + } + } +} + +/// Enum dispatch for backend implementations (avoids async trait object issues). +#[derive(Clone)] +pub enum Backend { + OpenAICompletions(openai_completions::OpenAICompletionsBackend), + OpenAIChat(openai_chat::OpenAIChatBackend), + Pooling(pooling::PoolingBackend), +} + +impl Backend { + /// Send a single request and collect timing metrics. + pub async fn send_request( + &self, + input: &RequestFuncInput, + client: &reqwest::Client, + ) -> Result { + match self { + Backend::OpenAICompletions(b) => b.send_request(input, client).await, + Backend::OpenAIChat(b) => b.send_request(input, client).await, + Backend::Pooling(b) => b.send_request(input, client).await, + } + } +} + +/// Get a backend by kind. +pub fn get_backend(kind: BackendKind) -> Result { + match kind { + BackendKind::Vllm | BackendKind::Openai => Ok(Backend::OpenAICompletions( + openai_completions::OpenAICompletionsBackend, + )), + BackendKind::OpenaiChat => Ok(Backend::OpenAIChat(openai_chat::OpenAIChatBackend)), + kind if kind.is_pooling() => Ok(Backend::Pooling(pooling::PoolingBackend { kind })), + _ => unreachable!(), + } +} + +/// Cached API key to avoid per-request env var syscall. +static API_KEY: std::sync::OnceLock> = std::sync::OnceLock::new(); + +fn cached_api_key() -> &'static Option { + API_KEY.get_or_init(|| std::env::var("OPENAI_API_KEY").ok()) +} + +/// Build common headers including auth and extras. +pub fn build_headers( + content_type: Option<&str>, + extra_headers: &Option>, + request_id: &Option, +) -> HashMap { + let mut headers = HashMap::new(); + + if let Some(ct) = content_type { + headers.insert("Content-Type".to_string(), ct.to_string()); + } + + if let Some(api_key) = cached_api_key() { + headers.insert("Authorization".to_string(), format!("Bearer {api_key}")); + } + + if let Some(extra) = extra_headers { + headers.extend(extra.iter().map(|(k, v)| (k.clone(), v.clone()))); + } + + if let Some(rid) = request_id { + headers.insert("x-request-id".to_string(), rid.clone()); + } + + headers +} diff --git a/rust/src/bench/src/backends/openai_chat.rs b/rust/src/bench/src/backends/openai_chat.rs new file mode 100644 index 000000000000..0842d23e6e5c --- /dev/null +++ b/rust/src/bench/src/backends/openai_chat.rs @@ -0,0 +1,376 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::time::Instant; + +use futures::StreamExt; + +use super::streaming::{StreamedResponseHandler, trim_bytes}; +use super::{ChatChunk, RequestFuncInput, RequestFuncOutput, build_headers}; +use crate::error::Result; + +/// Backend for OpenAI Chat Completions API (/v1/chat/completions). +#[derive(Clone)] +pub struct OpenAIChatBackend; + +impl OpenAIChatBackend { + pub async fn send_request( + &self, + input: &RequestFuncInput, + client: &reqwest::Client, + ) -> Result { + // Content-Type is set below by `.json()` / `.header()`; keep it out of + // headers_map to avoid a duplicate that strict gateways reject. + let headers_map = build_headers(None, &input.extra_headers, &input.request_id); + + let mut output = RequestFuncOutput { + prompt_len: input.prompt_len, + itl: Vec::with_capacity(input.output_len.max(1)), + ..Default::default() + }; + + let st = Instant::now(); + + let mut most_recent_timestamp = st; + let mut generated_text = String::new(); + let mut first_token_received = false; + + // Build request: use zero-copy raw JSON for multimodal, serde_json for text-only + let mut request = + if input.multi_modal_content.is_some() || input.chat_messages_json.is_some() { + let payload_bytes = build_mm_payload(input); + client + .post(&input.api_url) + .header("content-type", "application/json") + .body(payload_bytes) + } else { + let payload = build_text_payload(input); + client.post(&input.api_url).json(&payload) + }; + for (k, v) in &headers_map { + request = request.header(k, v); + } + + match request.send().await { + Ok(response) => { + if response.status().is_success() { + let mut handler = StreamedResponseHandler::new(); + let mut stream = response.bytes_stream(); + + while let Some(chunk_result) = stream.next().await { + let chunk_bytes = match chunk_result { + Ok(b) => b, + Err(e) => { + output.success = false; + output.error = format!("Stream error: {e}"); + return Ok(output); + } + }; + + let trimmed_bytes = trim_bytes(&chunk_bytes); + if trimmed_bytes.is_empty() { + continue; + } + + let messages = handler.add_chunk(trimmed_bytes); + for message in messages { + // Skip SSE comments + if message.starts_with(':') { + continue; + } + + // Handle multi-field SSE events (e.g., Dynamo sends + // "event: message\ndata: {...}"). Extract the data: line. + let raw = if message.contains('\n') { + match message.lines().find(|l| l.starts_with("data: ")) { + Some(l) => l, + None => continue, + } + } else { + message.as_str() + }; + + let chunk = raw.strip_prefix("data: ").unwrap_or(raw); + + if chunk == "[DONE]" { + continue; + } + + // Python chat backend: timestamp is captured for ALL + // non-DONE messages, and most_recent_timestamp is updated + // unconditionally (outside `if choices:`). This differs from + // completions which only timestamps content chunks. + let timestamp = Instant::now(); + + let data: ChatChunk = match serde_json::from_str(chunk) { + Ok(d) => d, + Err(_) => continue, + }; + + if !data.choices.is_empty() { + let content = data.choices[0] + .delta + .as_ref() + .and_then(|d| d.content.as_deref()) + .unwrap_or(""); + + if !first_token_received { + first_token_received = true; + output.ttft = timestamp.duration_since(st).as_secs_f64(); + } else { + output.itl.push( + timestamp + .duration_since(most_recent_timestamp) + .as_secs_f64(), + ); + } + + generated_text.push_str(content); + } + // Separate `if` (not `else if`) — Dynamo may send + // both choices and usage in the same chunk. + if let Some(ref usage) = data.usage + && let Some(ct) = usage.completion_tokens + { + output.output_tokens = ct as usize; + } + + most_recent_timestamp = timestamp; + } + } + + output.generated_text = generated_text; + output.success = true; + output.latency = most_recent_timestamp.duration_since(st).as_secs_f64(); + } else { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + output.error = if body.is_empty() { + format!("HTTP {status}") + } else { + format!("HTTP {status}: {body}") + }; + output.success = false; + } + } + Err(e) => { + output.success = false; + output.error = format!("{e:#}"); + } + } + + Ok(output) + } +} + +/// Build a JSON payload for text-only (non-multimodal) requests using serde_json. +fn build_text_payload(input: &RequestFuncInput) -> serde_json::Value { + let model = input.model_name.as_deref().unwrap_or(&input.model); + + let messages = if let Some(ref msgs) = input.messages { + msgs.clone() + } else { + let content = serde_json::json!([ + {"type": "text", "text": input.prompt} + ]); + serde_json::json!([{"role": "user", "content": content}]) + }; + + let mut payload = serde_json::json!({ + "model": model, + "messages": messages, + "max_completion_tokens": input.output_len, + "stream": true, + "stream_options": { + "include_usage": true, + }, + }); + + if input.ignore_eos { + payload["ignore_eos"] = serde_json::json!(true); + } + if let Some(serde_json::Value::Object(map)) = input.extra_body.as_ref() { + for (k, v) in map { + payload[k] = v.clone(); + } + } + + payload +} + +/// Build the JSON payload as raw bytes for multimodal requests. +/// +/// This is the zero-copy fast path: pre-serialized mm content fragments +/// (each ~200KB+ of base64 image data) are concatenated directly into the +/// output buffer without being parsed, cloned, or re-serialized. +/// +/// Saves ~200KB of allocation + copy per image per request compared to +/// the serde_json::Value approach. +fn build_mm_payload(input: &RequestFuncInput) -> Vec { + let model = input.model_name.as_deref().unwrap_or(&input.model); + + // Estimate total size: JSON overhead (~300 bytes) + prompt + mm fragments + let mm_total: usize = input + .multi_modal_content + .as_ref() + .map(|mm| mm.iter().map(|f| f.len() + 1).sum()) + .unwrap_or(0) + + input.chat_messages_json.as_ref().map_or(0, |m| m.len()); + let estimated = 512 + input.prompt.len() * 2 + mm_total; + let mut json = String::with_capacity(estimated); + + // {"model": + json.push_str(r#"{"model":"#); + // serde_json::to_string on &str produces a JSON-escaped quoted string + json.push_str(&serde_json::to_string(model).unwrap()); + + json.push_str(r#","messages":"#); + if let Some(ref msgs) = input.chat_messages_json { + // --enable-multimodal-chat: the dataset pre-built the full messages + // array (text + mm parts); splice it verbatim. + json.push_str(msgs); + } else { + let mm = input.multi_modal_content.as_ref().unwrap(); + + // [{"role":"user","content":[ + json.push_str(r#"[{"role":"user","content":[{"type":"text","text":""#); + // JSON-escape the prompt text (handles \n, \t, unicode, quotes) + push_json_escaped_str(&mut json, &input.prompt); + json.push_str(r#""}"#); + + // ,,,... + for fragment in mm.iter() { + json.push(','); + json.push_str(fragment); + } + + // Close content, message, messages + json.push_str(r#"]}]"#); + } + + // ,"max_completion_tokens": N, "stream": true, ... + json.push_str(r##","max_completion_tokens":"##); + json.push_str(&input.output_len.to_string()); + json.push_str(r##","stream":true,"stream_options":{"include_usage":true}"##); + + if input.ignore_eos { + json.push_str(r#","ignore_eos":true"#); + } + + // Merge extra_body key-value pairs, skipping keys already set above + if let Some(serde_json::Value::Object(map)) = input.extra_body.as_ref() { + for (k, v) in map { + match k.as_str() { + "model" + | "messages" + | "max_completion_tokens" + | "stream" + | "stream_options" + | "ignore_eos" => continue, + _ => { + json.push(','); + json.push_str(&serde_json::to_string(k).unwrap()); + json.push(':'); + json.push_str(&serde_json::to_string(v).unwrap()); + } + } + } + } + + json.push('}'); + json.into_bytes() +} + +/// Write a JSON-escaped string (without surrounding quotes) into the buffer. +/// +/// Handles: `\n`, `\r`, `\t`, `\\`, `\"`, and control characters. +/// This avoids the allocation of `serde_json::to_string` which produces +/// a new String with surrounding quotes. +fn push_json_escaped_str(buf: &mut String, s: &str) { + use std::fmt::Write; + for ch in s.chars() { + match ch { + '"' => buf.push_str(r#"\""#), + '\\' => buf.push_str(r"\\"), + '\n' => buf.push_str(r"\n"), + '\r' => buf.push_str(r"\r"), + '\t' => buf.push_str(r"\t"), + c if c.is_control() => { + // \uXXXX escape for control characters + for unit in c.encode_utf16(&mut [0; 2]) { + write!(buf, "\\u{unit:04x}").unwrap(); + } + } + c => buf.push(c), + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + + fn mm_input() -> RequestFuncInput { + let frag: Arc = + Arc::from(r#"{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,AAAA"}}"#); + RequestFuncInput { + prompt: Arc::from("hello \"world\"\nline2"), + model: "test-model".to_string(), + output_len: 128, + multi_modal_content: Some(Arc::from(vec![frag])), + ..Default::default() + } + } + + /// Regression test: the assembled multimodal payload must be valid JSON + /// (the text part once shipped without its opening quote — a raw-string + /// delimiter eating the trailing `"` in `"text":"`). + #[test] + fn test_build_mm_payload_is_valid_json() { + let payload = build_mm_payload(&mm_input()); + let v: serde_json::Value = + serde_json::from_slice(&payload).expect("mm payload must be valid JSON"); + assert_eq!(v["model"], "test-model"); + assert_eq!(v["messages"][0]["role"], "user"); + let content = v["messages"][0]["content"].as_array().unwrap(); + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[0]["text"], "hello \"world\"\nline2"); + assert_eq!(content[1]["type"], "image_url"); + assert_eq!(v["max_completion_tokens"], 128); + assert_eq!(v["stream"], true); + assert_eq!(v["stream_options"]["include_usage"], true); + } + + /// --enable-multimodal-chat (dataset pre-built messages) must produce a + /// payload semantically identical to the fragment-assembly path. + #[test] + fn test_chat_messages_json_path_equivalent_to_fragment_path() { + let base = mm_input(); + let fragment_payload = build_mm_payload(&base); + + let mut chat = base.clone(); + let mm = chat.multi_modal_content.take().unwrap(); + let msgs = crate::datasets::random_mm::build_chat_messages_json(&chat.prompt, Some(&mm)); + chat.chat_messages_json = Some(Arc::from(msgs.as_str())); + let chat_payload = build_mm_payload(&chat); + + let a: serde_json::Value = serde_json::from_slice(&fragment_payload).unwrap(); + let b: serde_json::Value = serde_json::from_slice(&chat_payload).unwrap(); + assert_eq!(a, b); + } + + /// ignore_eos and extra_body must survive the raw-splice path. + #[test] + fn test_mm_payload_tail_fields() { + let mut input = mm_input(); + input.ignore_eos = true; + input.extra_body = Some(serde_json::json!({"temperature": 0.5, "stream": false})); + let v: serde_json::Value = serde_json::from_slice(&build_mm_payload(&input)).unwrap(); + assert_eq!(v["ignore_eos"], true); + assert_eq!(v["temperature"], 0.5); + // keys already set above must not be overridden by extra_body + assert_eq!(v["stream"], true); + } +} diff --git a/rust/src/bench/src/backends/openai_completions.rs b/rust/src/bench/src/backends/openai_completions.rs new file mode 100644 index 000000000000..e678da891414 --- /dev/null +++ b/rust/src/bench/src/backends/openai_completions.rs @@ -0,0 +1,199 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::error::Error as StdError; +use std::time::Instant; + +use futures::StreamExt; + +use super::streaming::{StreamedResponseHandler, trim_bytes}; +use super::{CompletionChunk, RequestFuncInput, RequestFuncOutput, build_headers}; +use crate::error::Result; + +/// Backend for OpenAI-compatible Completions API (/v1/completions). +/// Used by "vllm" and "openai" backends. +#[derive(Clone)] +pub struct OpenAICompletionsBackend; + +impl OpenAICompletionsBackend { + pub async fn send_request( + &self, + input: &RequestFuncInput, + client: &reqwest::Client, + ) -> Result { + let model = input.model_name.as_deref().unwrap_or(&input.model); + + // When prompt_token_ids are available, send them as the `prompt` value + // (JSON array of integers). vLLM's completions API accepts both string + // and token ID array as `prompt`, skipping server-side tokenization. + let prompt_value = if let Some(ref token_ids) = input.prompt_token_ids { + serde_json::json!(token_ids.as_ref()) + } else { + serde_json::json!(input.prompt) + }; + + let mut payload = serde_json::json!({ + "model": model, + "prompt": prompt_value, + "max_tokens": input.output_len, + "stream": true, + "stream_options": { + "include_usage": true, + }, + }); + + // Always include logprobs (null when not set) — matches Python which + // sends logprobs=None explicitly rather than omitting the key. + payload["logprobs"] = match input.logprobs { + Some(n) => serde_json::json!(n), + None => serde_json::Value::Null, + }; + + // Apply ignore_eos and extra_body + if input.ignore_eos { + payload["ignore_eos"] = serde_json::json!(true); + } + if let Some(serde_json::Value::Object(map)) = input.extra_body.as_ref() { + for (k, v) in map { + payload[k] = v.clone(); + } + } + + let headers_map = build_headers(None, &input.extra_headers, &input.request_id); + + let mut output = RequestFuncOutput { + prompt_len: input.prompt_len, + itl: Vec::with_capacity(input.output_len.max(1)), + ..Default::default() + }; + + let st = Instant::now(); + // start_time is overwritten by benchmark.rs with monotonic offset + + let mut most_recent_timestamp = st; + let mut generated_text = String::new(); + let mut first_chunk_received = false; + + let mut request = client.post(&input.api_url).json(&payload); + for (k, v) in &headers_map { + request = request.header(k, v); + } + + match request.send().await { + Ok(response) => { + if response.status().is_success() { + let mut handler = StreamedResponseHandler::new(); + let mut stream = response.bytes_stream(); + + while let Some(chunk_result) = stream.next().await { + let chunk_bytes = match chunk_result { + Ok(b) => b, + Err(e) => { + output.success = false; + output.error = format!("Stream error: {e}"); + return Ok(output); + } + }; + + let trimmed_bytes = trim_bytes(&chunk_bytes); + if trimmed_bytes.is_empty() { + continue; + } + + let messages = handler.add_chunk(trimmed_bytes); + for message in messages { + // Skip SSE comments + if message.starts_with(':') { + continue; + } + + // Handle multi-field SSE events (e.g., Dynamo sends + // "event: message\ndata: {...}"). Extract the data: line. + let raw = if message.contains('\n') { + match message.lines().find(|l| l.starts_with("data: ")) { + Some(l) => l, + None => continue, + } + } else { + message.as_str() + }; + + let chunk = raw.strip_prefix("data: ").unwrap_or(raw); + + if chunk == "[DONE]" { + continue; + } + + // Typed deserialization — avoids allocating a full + // serde_json::Value tree; only extracts needed fields. + let data: CompletionChunk = match serde_json::from_str(chunk) { + Ok(d) => d, + Err(_) => continue, + }; + + if !data.choices.is_empty() { + let text = data.choices[0].text.as_deref().unwrap_or(""); + + let timestamp = Instant::now(); + + if !first_chunk_received { + first_chunk_received = true; + output.ttft = timestamp.duration_since(st).as_secs_f64(); + } else { + output.itl.push( + timestamp + .duration_since(most_recent_timestamp) + .as_secs_f64(), + ); + } + + most_recent_timestamp = timestamp; + generated_text.push_str(text); + } + // Separate `if` (not `else if`) — Dynamo may send + // both choices and usage in the same chunk. + if let Some(ref usage) = data.usage + && let Some(ct) = usage.completion_tokens + { + output.output_tokens = ct as usize; + } + } + } + + if first_chunk_received { + output.success = true; + } else { + output.success = false; + output.error = "Never received a valid chunk to calculate TTFT. \ + This response will be marked as failed!" + .to_string(); + } + output.generated_text = generated_text; + output.latency = most_recent_timestamp.duration_since(st).as_secs_f64(); + } else { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + output.error = if body.is_empty() { + format!("HTTP {status}") + } else { + format!("HTTP {status}: {body}") + }; + output.success = false; + } + } + Err(e) => { + output.success = false; + // Capture full error chain for debugging + let mut error_msg = format!("{e}"); + let mut source = e.source(); + while let Some(cause) = source { + error_msg.push_str(&format!("\n Caused by: {cause}")); + source = cause.source(); + } + output.error = error_msg; + } + } + + Ok(output) + } +} diff --git a/rust/src/bench/src/backends/pooling.rs b/rust/src/bench/src/backends/pooling.rs new file mode 100644 index 000000000000..6a1cdf4fe160 --- /dev/null +++ b/rust/src/bench/src/backends/pooling.rs @@ -0,0 +1,325 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Pooling/embedding backends: non-streaming HTTP POST for embedding, pooling, and rerank +//! endpoints. +//! +//! Supported variants: +//! - `openai-embeddings`: Standard OpenAI `/v1/embeddings` with text input +//! - `openai-embeddings-chat`: OpenAI `/v1/embeddings` with chat message format (supports +//! multimodal) +//! - `vllm-pooling`: vLLM `/v1/pooling` endpoint +//! - `vllm-rerank`: vLLM `/v1/rerank` endpoint (query + documents) + +use std::time::Instant; + +use crate::backends::{RequestFuncInput, RequestFuncOutput, build_headers}; +use crate::cli::BackendKind; +use crate::error::Result; + +/// Response from embedding/pooling endpoints (minimal fields for usage extraction). +#[derive(serde::Deserialize)] +struct PoolingResponse { + usage: Option, +} + +#[derive(serde::Deserialize)] +struct PoolingUsage { + prompt_tokens: Option, +} + +#[derive(Clone)] +pub struct PoolingBackend { + pub kind: BackendKind, +} + +impl PoolingBackend { + pub async fn send_request( + &self, + input: &RequestFuncInput, + client: &reqwest::Client, + ) -> Result { + // Preserve client-side prompt_len as fallback if server doesn't report usage. + let mut output = RequestFuncOutput { + prompt_len: input.prompt_len, + ..Default::default() + }; + + let headers = build_headers( + Some("application/json"), + &input.extra_headers, + &input.request_id, + ); + + let payload = self.build_payload(input); + + let mut request = client.post(&input.api_url); + for (k, v) in &headers { + request = request.header(k, v); + } + + let st = Instant::now(); + + let response = match request.json(&payload).send().await { + Ok(r) => r, + Err(e) => { + output.error = format!("Request failed: {e}"); + return Ok(output); + } + }; + + if response.status().is_success() { + let latency = st.elapsed().as_secs_f64(); + output.latency = latency; + output.ttft = latency; + output.success = true; + + // Parse usage from response; keep client-side prompt_len as fallback. + match response.json::().await { + Ok(data) => { + if let Some(usage) = data.usage + && let Some(tokens) = usage.prompt_tokens + { + output.prompt_len = tokens as usize; + } + } + Err(_) => { + // Response parsed but no usage — keep client-side prompt_len + } + } + } else { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + output.error = format!("HTTP {status}: {body}"); + } + + Ok(output) + } + + fn build_payload(&self, input: &RequestFuncInput) -> serde_json::Value { + let model = input.model_name.as_deref().unwrap_or(&input.model); + + // For "input" field (openai-embeddings, vllm-pooling): a batched request + // (--random-batch-size) sends the text list; otherwise prefer prompt_token_ids + // when available. The random dataset sets prompt="" and relies on token IDs; + // the OpenAI embeddings API accepts both text strings and token ID arrays. + // Note: embeddings-chat uses text in messages; vllm-rerank uses text as query. + let input_value = if let Some(ref list) = input.prompt_list { + serde_json::json!(list.iter().map(|s| s.as_ref()).collect::>()) + } else if let Some(ref token_ids) = input.prompt_token_ids { + serde_json::json!(token_ids.as_ref()) + } else { + serde_json::json!(input.prompt.as_ref()) + }; + + let is_vllm_backend = matches!( + self.kind, + BackendKind::VllmPooling | BackendKind::VllmRerank + ); + + let mut payload = match self.kind { + BackendKind::OpenaiEmbeddings => { + let mut p = serde_json::json!({ + "model": model, + "input": input_value, + }); + // truncate_prompt_tokens is vLLM-specific; only include for vLLM backends + // to avoid breaking standard OpenAI providers. + if is_vllm_backend { + p["truncate_prompt_tokens"] = serde_json::json!(-1); + } + p + } + BackendKind::OpenaiEmbeddingsChat => { + // Chat format: uses text prompt in messages array (for multimodal support). + // Python's _get_chat_content always returns a content array. + // Use raw string concatenation for multimodal fragments (zero-copy, + // avoids re-parsing ~200KB+ base64 per image). + let content_json = build_chat_content_json(input); + + let mut p = serde_json::json!({ + "model": model, + "messages": [{"role": "user", "content": content_json}], + }); + if is_vllm_backend { + p["truncate_prompt_tokens"] = serde_json::json!(-1); + } + p + } + BackendKind::VllmPooling => { + serde_json::json!({ + "model": model, + "input": input_value, + "truncate_prompt_tokens": -1, + }) + } + BackendKind::VllmRerank => { + // random-rerank dataset: prompt_list = [query, doc1, doc2, ...] + // (mirrors Python async_request_vllm_rerank). + if let Some(ref list) = input.prompt_list { + if list.len() < 2 { + eprintln!( + "WARNING: vllm-rerank request has no documents \ + (prompt_list needs [query, doc, ...])" + ); + } + let query = list.first().map(|s| s.as_ref()).unwrap_or(""); + let documents: Vec<&str> = list.iter().skip(1).map(|s| s.as_ref()).collect(); + serde_json::json!({ + "model": model, + "query": query, + "documents": documents, + "truncate_prompt_tokens": -1, + }) + } else { + // Legacy path: text prompt as query, documents via --extra-body. + let query = input.prompt.as_ref(); + if query.is_empty() && input.prompt_token_ids.is_some() { + eprintln!( + "WARNING: vllm-rerank received empty query (random dataset uses \ + token IDs only). Use --dataset-name random-rerank for meaningful \ + rerank benchmarks." + ); + } + serde_json::json!({ + "model": model, + "query": query, + "truncate_prompt_tokens": -1, + }) + } + } + _ => unreachable!("PoolingBackend with non-pooling kind"), + }; + + // Merge extra_body fields into payload + if let Some(ref extra) = input.extra_body + && let (Some(base), Some(extra_obj)) = (payload.as_object_mut(), extra.as_object()) + { + for (k, v) in extra_obj { + base.insert(k.clone(), v.clone()); + } + } + + payload + } +} + +/// Build the chat content JSON array for embeddings-chat. +/// Uses raw string concatenation for multimodal fragments to avoid +/// re-parsing large base64 image data (matching openai_chat.rs approach). +fn build_chat_content_json(input: &RequestFuncInput) -> serde_json::Value { + if input.multi_modal_content.is_none() { + // Text-only: return content array with single text element + return serde_json::json!([{ + "type": "text", + "text": input.prompt.as_ref(), + }]); + } + + // Multimodal: build JSON string manually for zero-copy fragment embedding + let mm = input.multi_modal_content.as_ref().unwrap(); + let prompt = input.prompt.as_ref(); + + let mm_total: usize = mm.iter().map(|f| f.len() + 1).sum(); + let mut json = String::with_capacity(64 + prompt.len() * 2 + mm_total); + + // [{"type":"text","text":""} + json.push_str(r#"[{"type":"text","text":""#); + push_json_escaped_str(&mut json, prompt); + json.push_str(r#""}"#); + + // ,,,... + for fragment in mm.iter() { + json.push(','); + json.push_str(fragment); + } + + json.push(']'); + + // Parse the assembled string into a Value for embedding in the payload. + // This parse is O(n) but operates on the pre-built string once, not per-fragment. + serde_json::from_str(&json).unwrap_or_else(|_| { + serde_json::json!([{ + "type": "text", + "text": input.prompt.as_ref(), + }]) + }) +} + +/// Escape a string for safe JSON embedding (matching openai_chat.rs). +fn push_json_escaped_str(buf: &mut String, s: &str) { + use std::fmt::Write; + for ch in s.chars() { + match ch { + '"' => buf.push_str(r#"\""#), + '\\' => buf.push_str(r"\\"), + '\n' => buf.push_str(r"\n"), + '\r' => buf.push_str(r"\r"), + '\t' => buf.push_str(r"\t"), + c if c < '\x20' => { + let _ = write!(buf, "\\u{:04x}", c as u32); + } + c => buf.push(c), + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + + fn list(items: &[&str]) -> Option]>> { + Some(items.iter().map(|s| Arc::from(*s)).collect()) + } + + #[test] + fn test_embeddings_payload_batched_input() { + let backend = PoolingBackend { + kind: BackendKind::OpenaiEmbeddings, + }; + let input = RequestFuncInput { + model: "bge".to_string(), + prompt_list: list(&["t1", "t2", "t3"]), + ..Default::default() + }; + let payload = backend.build_payload(&input); + assert_eq!(payload["input"], serde_json::json!(["t1", "t2", "t3"])); + assert_eq!(payload["model"], "bge"); + // truncate_prompt_tokens is vLLM-specific and deliberately omitted for + // the plain OpenAI embeddings backend. + assert!(payload.get("truncate_prompt_tokens").is_none()); + } + + #[test] + fn test_rerank_payload_query_and_documents() { + let backend = PoolingBackend { + kind: BackendKind::VllmRerank, + }; + let input = RequestFuncInput { + model: "reranker".to_string(), + prompt_list: list(&["the query", "doc a", "doc b"]), + ..Default::default() + }; + let payload = backend.build_payload(&input); + assert_eq!(payload["query"], "the query"); + assert_eq!(payload["documents"], serde_json::json!(["doc a", "doc b"])); + assert_eq!(payload["truncate_prompt_tokens"], -1); + } + + #[test] + fn test_rerank_payload_legacy_single_prompt() { + let backend = PoolingBackend { + kind: BackendKind::VllmRerank, + }; + let input = RequestFuncInput { + model: "reranker".to_string(), + prompt: Arc::from("query text"), + ..Default::default() + }; + let payload = backend.build_payload(&input); + assert_eq!(payload["query"], "query text"); + assert!(payload.get("documents").is_none()); + } +} diff --git a/rust/src/bench/src/backends/streaming.rs b/rust/src/bench/src/backends/streaming.rs new file mode 100644 index 000000000000..16837a274a50 --- /dev/null +++ b/rust/src/bench/src/backends/streaming.rs @@ -0,0 +1,151 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +/// SSE streaming response handler. +/// +/// Accumulates incoming byte chunks and extracts complete SSE messages. +/// Mirrors Python's `StreamedResponseHandler` from endpoint_request_func.py:22-60. +pub struct StreamedResponseHandler { + buffer: String, + /// Reusable message buffer — avoids allocating a new Vec per `add_chunk` call. + messages: Vec, +} + +impl StreamedResponseHandler { + pub fn new() -> Self { + Self { + buffer: String::with_capacity(4096), + messages: Vec::with_capacity(4), + } + } + + /// Add a chunk of bytes and return any complete SSE messages. + /// + /// The returned slice borrows from the handler and is valid until the next + /// `add_chunk` call. + pub fn add_chunk(&mut self, chunk_bytes: &[u8]) -> &[String] { + self.messages.clear(); + + let chunk_str = String::from_utf8_lossy(chunk_bytes); + self.buffer.push_str(&chunk_str); + + // Split by double newlines (SSE message separator) + while let Some(pos) = self.buffer.find("\n\n") { + let message = self.buffer[..pos].trim().to_string(); + // Efficiently remove consumed bytes by shifting remaining data + self.buffer.drain(..pos + 2); + if !message.is_empty() { + self.messages.push(message); + } + } + + // Handle buffered data without trailing `\n\n`. + // Matches Python's speculative json.loads() in StreamedResponseHandler. + // This matters for TTFT/ITL accuracy: when a data message and its `\n\n` + // arrive in separate TCP segments, we want to emit the message at the + // first segment's arrival time, not the second. + // + // Also handles multi-field SSE events where the buffer may start with + // "event: ...\ndata: ..." (Dynamo frontend). + let data_start = if self.buffer.starts_with("data: ") { + Some(0) + } else { + // Look for a "data: " line in multi-field events + self.buffer.find("\ndata: ").map(|p| p + 1) + }; + if let Some(offset) = data_start { + let content = self.buffer[offset + 6..].trim(); + if content == "[DONE]" + || (!content.is_empty() + && serde_json::from_str::<&serde_json::value::RawValue>(content).is_ok()) + { + self.messages.push(self.buffer.trim().to_string()); + self.buffer.clear(); + } + } + + &self.messages + } +} + +/// Trim leading/trailing ASCII whitespace from a byte slice. +pub fn trim_bytes(bytes: &[u8]) -> &[u8] { + let start = bytes.iter().position(|b| !b.is_ascii_whitespace()).unwrap_or(bytes.len()); + let end = bytes + .iter() + .rposition(|b| !b.is_ascii_whitespace()) + .map(|p| p + 1) + .unwrap_or(start); + &bytes[start..end] +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basic_sse() { + let mut handler = StreamedResponseHandler::new(); + let msgs = + handler.add_chunk(b"data: {\"choices\":[{\"text\":\"hi\"}]}\n\ndata: [DONE]\n\n"); + assert_eq!(msgs.len(), 2); + assert!(msgs[0].contains("choices")); + assert!(msgs[1].contains("[DONE]")); + } + + #[test] + fn test_split_chunks() { + let mut handler = StreamedResponseHandler::new(); + let msgs1 = handler.add_chunk(b"data: {\"cho"); + assert!(msgs1.is_empty()); + let msgs2 = handler.add_chunk(b"ices\":[{\"text\":\"a\"}]}\n\n"); + assert_eq!(msgs2.len(), 1); + } + + #[test] + fn test_comment_lines() { + let mut handler = StreamedResponseHandler::new(); + let msgs = handler.add_chunk(b": ping\n\ndata: {\"test\":1}\n\n"); + assert_eq!(msgs.len(), 2); + assert!(msgs[0].starts_with(":")); + } + + #[test] + fn test_done_without_newlines() { + let mut handler = StreamedResponseHandler::new(); + let msgs = handler.add_chunk(b"data: [DONE]"); + assert_eq!(msgs.len(), 1); + assert!(msgs[0].contains("[DONE]")); + } + + #[test] + fn test_incomplete_json_in_buffer() { + let mut handler = StreamedResponseHandler::new(); + let msgs = handler.add_chunk(b"data: {\"partial\":"); + assert!(msgs.is_empty()); + // Complete JSON without \n\n — speculative parse emits it + let msgs2 = handler.add_chunk(b"true}"); + assert_eq!(msgs2.len(), 1); + assert!(msgs2[0].contains("partial")); + } + + #[test] + fn test_multi_field_sse_event() { + // Dynamo frontend sends "event: message\ndata: {...}\n\n" + let mut handler = StreamedResponseHandler::new(); + let msgs = + handler.add_chunk(b"event: message\ndata: {\"choices\":[{\"text\":\"hi\"}]}\n\n"); + assert_eq!(msgs.len(), 1); + assert!(msgs[0].contains("choices")); + assert!(msgs[0].contains("event: message")); + } + + #[test] + fn test_multi_field_sse_speculative_parse() { + // Multi-field event without trailing \n\n — speculative parse should emit it + let mut handler = StreamedResponseHandler::new(); + let msgs = handler.add_chunk(b"event: message\ndata: {\"choices\":[{\"text\":\"hi\"}]}"); + assert_eq!(msgs.len(), 1); + assert!(msgs[0].contains("choices")); + } +} diff --git a/rust/src/bench/src/benchmark.rs b/rust/src/bench/src/benchmark.rs new file mode 100644 index 000000000000..8c834641acb0 --- /dev/null +++ b/rust/src/bench/src/benchmark.rs @@ -0,0 +1,1977 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Instant; + +use indicatif::{ProgressBar, ProgressStyle}; +use tokio::sync::Semaphore; + +use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend}; +use crate::cli::{BackendKind, DatasetName, LoraAssignment}; +use crate::config::BenchConfig; +use crate::error::{BenchError, Result}; +use crate::metrics::calculator::{calculate_embedding_metrics, calculate_metrics}; +use crate::metrics::steady_state; +use crate::output::console::print_results; +use crate::output::json::{append_result, build_result_json, compute_result_filename, save_result}; +use crate::rate_control::compute_schedule; +use crate::ready_checker::wait_for_endpoint; + +/// Pre-resolve the hostname in `base_url` and pin all resolved IPs on the +/// client builder via [`reqwest::ClientBuilder::resolve_to_addrs`]. This +/// avoids repeated DNS lookups under high concurrency which can cause +/// transient "Temporary failure in name resolution" errors while preserving +/// happy-eyeballs and multi-A failover. +/// +/// Skipped when: URL parse fails, host is already an IP, host is a loopback +/// name (resolved from `/etc/hosts`, no DNS pressure and dual-stack ambiguity +/// between `127.0.0.1` and `::1` breaks IPv4-only servers like vLLM), or +/// resolution fails. +pub fn pre_resolve_dns( + base_url: &str, + mut builder: reqwest::ClientBuilder, +) -> reqwest::ClientBuilder { + let parsed = match url::Url::parse(base_url) { + Ok(u) => u, + Err(_) => return builder, + }; + + let host = match parsed.host_str() { + Some(h) => h, + None => return builder, + }; + + if host.parse::().is_ok() { + return builder; + } + + let host_lower = host.to_ascii_lowercase(); + if host_lower == "localhost" + || host_lower == "ip6-localhost" + || host_lower.ends_with(".localhost") + { + return builder; + } + + let port = parsed.port_or_known_default().unwrap_or(80); + let addr_str = format!("{host}:{port}"); + + match std::net::ToSocketAddrs::to_socket_addrs(&addr_str) { + Ok(addrs) => { + let mut v4 = Vec::new(); + let mut v6 = Vec::new(); + for addr in addrs { + if addr.is_ipv4() { + v4.push(addr); + } else { + v6.push(addr); + } + } + v4.extend(v6); + if !v4.is_empty() { + let ips: Vec<_> = v4.iter().map(|a| a.ip()).collect(); + println!("Pre-resolved {host} -> {ips:?}"); + builder = builder.resolve_to_addrs(host, &v4); + } + } + Err(e) => { + eprintln!("Warning: DNS pre-resolution for '{host}' failed: {e}"); + } + } + + builder +} + +/// Raw speculative decoding metrics from the server's Prometheus endpoint. +#[derive(Debug, Clone)] +pub(crate) struct SpecDecodeMetrics { + num_drafts: u64, + num_draft_tokens: u64, + num_accepted_tokens: u64, + accepted_per_pos: HashMap, +} + +/// Computed speculative decoding statistics (delta between before/after benchmark). +#[derive(Debug, Clone)] +pub struct SpecDecodeStats { + pub num_drafts: u64, + pub draft_tokens: u64, + pub accepted_tokens: u64, + pub acceptance_rate: f64, + pub acceptance_length: f64, + pub per_position_acceptance_rates: Vec, +} + +/// Fetch speculative decoding metrics from the server's Prometheus `/metrics` endpoint. +/// +/// Returns None if speculative decoding is not enabled or metrics are not available. +pub(crate) async fn fetch_spec_decode_metrics( + base_url: &str, + client: &reqwest::Client, + extra_headers: &Option>, +) -> Option { + let metrics_url = format!("{base_url}/metrics"); + let mut request = client.get(&metrics_url); + if let Some(headers) = extra_headers { + for (k, v) in headers { + request = request.header(k, v); + } + } + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") { + request = request.header("Authorization", format!("Bearer {api_key}")); + } + + let response = match request.send().await { + Ok(r) if r.status().is_success() => r, + _ => return None, + }; + + let text = match response.text().await { + Ok(t) => t, + Err(_) => return None, + }; + + let mut num_drafts: u64 = 0; + let mut num_draft_tokens: u64 = 0; + let mut num_accepted_tokens: u64 = 0; + let mut accepted_per_pos: HashMap = HashMap::new(); + let mut found_spec_decode = false; + + for line in text.lines() { + let line = line.trim(); + if line.is_empty() || line.starts_with('#') { + continue; + } + if !line.starts_with("vllm:spec_decode") { + continue; + } + found_spec_decode = true; + let parts: Vec<&str> = line.split_whitespace().collect(); + if parts.is_empty() { + continue; + } + let val = match parts.last().and_then(|s| s.parse::().ok()) { + Some(v) => v as u64, + None => continue, + }; + + if line.contains("num_drafts") { + num_drafts += val; + } else if line.contains("num_draft_tokens") { + num_draft_tokens += val; + } else if line.contains("num_accepted_tokens_per_pos") { + // Parse position label: position="N" + if let Some(start) = line.find("position=\"") { + let start = start + "position=\"".len(); + if let Some(end) = line[start..].find('"') + && let Ok(pos) = line[start..start + end].parse::() + { + *accepted_per_pos.entry(pos).or_insert(0) += val; + } + } + } else if line.contains("num_accepted_tokens") { + num_accepted_tokens += val; + } + } + + if !found_spec_decode { + return None; + } + + Some(SpecDecodeMetrics { + num_drafts, + num_draft_tokens, + num_accepted_tokens, + accepted_per_pos, + }) +} + +/// Compute speculative decoding stats from before/after metrics snapshots. +pub(crate) fn compute_spec_decode_stats( + before: &SpecDecodeMetrics, + after: &SpecDecodeMetrics, +) -> Option { + let delta_drafts = after.num_drafts.saturating_sub(before.num_drafts); + let delta_draft_tokens = after.num_draft_tokens.saturating_sub(before.num_draft_tokens); + let delta_accepted = after.num_accepted_tokens.saturating_sub(before.num_accepted_tokens); + + if delta_draft_tokens == 0 { + return None; + } + + let mut per_pos_rates = Vec::new(); + if delta_drafts > 0 { + let mut positions: Vec = before + .accepted_per_pos + .keys() + .chain(after.accepted_per_pos.keys()) + .copied() + .collect::>() + .into_iter() + .collect(); + positions.sort(); + + for pos in positions { + let before_val = before.accepted_per_pos.get(&pos).copied().unwrap_or(0); + let after_val = after.accepted_per_pos.get(&pos).copied().unwrap_or(before_val); + let delta_pos = after_val.saturating_sub(before_val); + per_pos_rates.push(delta_pos as f64 / delta_drafts as f64); + } + } + + let acceptance_rate = (delta_accepted as f64 / delta_draft_tokens as f64) * 100.0; + let acceptance_length = if delta_drafts > 0 { + 1.0 + delta_accepted as f64 / delta_drafts as f64 + } else { + 0.0 + }; + + Some(SpecDecodeStats { + num_drafts: delta_drafts, + draft_tokens: delta_draft_tokens, + accepted_tokens: delta_accepted, + acceptance_rate, + acceptance_length, + per_position_acceptance_rates: per_pos_rates, + }) +} + +/// Pre-assign a LoRA adapter name per item based on the configured strategy. +/// +/// Returns `None` when LoRA is not configured. With LoRA enabled, returns a +/// `Vec>` of length `n` whose i-th entry is the adapter name to use +/// for the i-th item. +/// +/// - `RoundRobin` cycles deterministically: `lora_modules[i % N]`. +/// - `Random` uses `StdRng::seed_from_u64(seed)` so the assignment is fully reproducible across +/// runs with the same seed. +/// +/// "Item" is a request in single-shot mode and a conversation in multi-turn +/// mode (sticky across all turns of that conversation). +pub(crate) fn assign_lora_modules( + lora_modules: &Option>>, + assignment: LoraAssignment, + n: usize, + seed: u64, +) -> Option>> { + let modules = lora_modules.as_ref()?; + if modules.is_empty() { + return None; + } + let m = modules.len(); + let mut out = Vec::with_capacity(n); + match assignment { + LoraAssignment::RoundRobin => { + for i in 0..n { + out.push(modules[i % m].clone()); + } + } + LoraAssignment::Random => { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + let mut rng = StdRng::seed_from_u64(seed); + for _ in 0..n { + out.push(modules[rng.random_range(0..m)].clone()); + } + } + } + Some(out) +} + +/// Fetch the first model from the server's /v1/models endpoint. +async fn get_first_model_from_server( + base_url: &str, + client: &reqwest::Client, + extra_headers: &Option>, +) -> Result<(String, String)> { + let url = format!("{base_url}/v1/models"); + let mut request = client.get(&url); + if let Some(headers) = extra_headers { + for (k, v) in headers { + request = request.header(k, v); + } + } + // Add API key from environment + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") { + request = request.header("Authorization", format!("Bearer {api_key}")); + } + + let response = request.send().await?; + let data: serde_json::Value = response.json().await?; + + if let Some(models) = data.get("data").and_then(|d| d.as_array()) + && let Some(first) = models.first() + { + let id = first.get("id").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + let root = first.get("root").and_then(|v| v.as_str()).unwrap_or(&id).to_string(); + return Ok((id, root)); + } + + Err(BenchError::Config(format!( + "No models found on the server at {base_url}" + ))) +} + +/// Run the complete benchmark. +/// +/// This is the core orchestrator matching Python's benchmark() + main_async(). +pub async fn run_benchmark(config: &BenchConfig) -> Result { + // Validate backend name early + let _ = get_backend(config.backend)?; + + // Build HTTP client with connection pool settings matching Python's aiohttp.TCPConnector + // Force HTTP/1.1 to match aiohttp behavior (avoids HTTP/2 negotiation issues) + let mut client_builder = reqwest::Client::builder() + .pool_max_idle_per_host(config.max_concurrency.unwrap_or(2048).max(256)) + .timeout(std::time::Duration::from_secs(6 * 60 * 60)) + .connect_timeout(std::time::Duration::from_secs(30)) + .tcp_keepalive(std::time::Duration::from_secs(60)) + .tcp_nodelay(true) + .http1_only() + .no_proxy(); + + if config.insecure { + client_builder = client_builder.danger_accept_invalid_certs(true); + } + + client_builder = pre_resolve_dns(&config.base_url, client_builder); + + let client = client_builder + .build() + .map_err(|e| BenchError::Backend(format!("Failed to build HTTP client: {e}")))?; + + // Resolve model + let (model_id, model_name) = if let Some(ref m) = config.model { + (m.clone(), config.model_name.clone()) + } else { + println!("Model not specified, fetching first model from server..."); + let (name, id) = + get_first_model_from_server(&config.base_url, &client, &config.extra_headers).await?; + println!("First model name: {name}, first model id: {id}"); + (id, Some(name)) + }; + + // Load tokenizer (if needed) + let tokenizer = if config.skip_tokenizer_init { + None + } else { + let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id); + println!("Loading tokenizer: {tid}"); + let server_info = Some((config.base_url.as_str(), model_id.as_str())); + let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?; + println!("Tokenizer loaded successfully."); + Some(t) + }; + let has_tokenizer = tokenizer.is_some(); + + // Generate dataset + let dataset_label = match config.dataset_name { + DatasetName::Random => format!("{} random prompts", config.num_prompts), + DatasetName::RandomMm => format!( + "{} random multimodal prompts ({} items/req)", + config.num_prompts, config.random_mm_base_items_per_request + ), + DatasetName::ShareGpt => format!( + "{} prompts from ShareGPT ({})", + config.num_prompts, + config.dataset_path.as_deref().unwrap_or("auto-download") + ), + DatasetName::Sonnet => format!( + "{} prompts from Sonnet ({}, isl={}, osl={}, prefix={})", + config.num_prompts, + config.dataset_path.as_deref().unwrap_or("built-in"), + config.sonnet_input_len, + config.sonnet_output_len, + config.sonnet_prefix_len, + ), + DatasetName::SpeedBench => { + let truncate_info = config + .speed_bench_max_input_len + .map(|n| format!(", truncated to {n} tokens")) + .unwrap_or_default(); + format!( + "{} prompts from SPEED-Bench ({}/{}{})", + config.num_prompts, + config.speed_bench_config, + config.speed_bench_category.as_deref().unwrap_or("all"), + truncate_info, + ) + } + DatasetName::Hf => format!( + "{} prompts from HF dataset ({})", + config.num_prompts, + config.dataset_path.as_deref().unwrap_or("unknown"), + ), + DatasetName::Custom => format!( + "{} prompts from custom JSONL ({})", + config.num_prompts, + config.dataset_path.as_deref().unwrap_or("unknown"), + ), + DatasetName::PrefixRepetition => format!( + "{} prefix-repetition prompts ({} prefixes, prefix={}, suffix={})", + config.num_prompts, + config.prefix_repetition_num_prefixes, + config.prefix_repetition_prefix_len, + config.prefix_repetition_suffix_len, + ), + DatasetName::RandomRerank => format!( + "{} random rerank requests (batch={}, reranker={})", + config.num_prompts, config.random_batch_size, config.is_reranker, + ), + }; + println!("Generating {dataset_label}..."); + let gen_start = Instant::now(); + + let mut input_requests = match config.dataset_name { + DatasetName::Random => { + let tok = tokenizer + .as_ref() + .ok_or_else(|| BenchError::Config("Random dataset requires a tokenizer".into()))?; + crate::datasets::random::generate_random_dataset( + tok, + config.num_prompts, + config.random_input_len, + config.random_output_len, + config.random_prefix_len, + config.random_range_ratio, + config.random_cache_hit_fraction, + config.random_cache_ratio, + config.seed, + &config.request_id_prefix, + config.prompt_token_ids, + config.random_batch_size, + )? + } + DatasetName::RandomMm => { + let tok = tokenizer.as_ref().ok_or_else(|| { + BenchError::Config("Random-MM dataset requires a tokenizer".into()) + })?; + crate::datasets::random_mm::generate_random_mm_dataset( + tok, + config.num_prompts, + config.random_input_len, + config.random_output_len, + config.random_prefix_len, + config.random_range_ratio, + config.seed, + &config.request_id_prefix, + config.random_mm_base_items_per_request, + config.random_mm_num_mm_items_range_ratio, + &config.random_mm_limit, + &config.random_mm_buckets, + config.enable_multimodal_chat, + )? + } + DatasetName::ShareGpt => { + let tok = tokenizer.as_ref().ok_or_else(|| { + BenchError::Config("ShareGPT dataset requires a tokenizer".into()) + })?; + let downloaded; + let path = match config.dataset_path.as_deref() { + Some(p) => p, + None => { + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?; + downloaded.as_str() + } + }; + crate::datasets::sharegpt::load_sharegpt_dataset( + tok, + path, + config.num_prompts, + config.sharegpt_output_len, + config.seed, + &config.request_id_prefix, + config.no_oversample, + config.disable_shuffle, + )? + } + DatasetName::Sonnet => { + let tok = tokenizer + .as_ref() + .ok_or_else(|| BenchError::Config("Sonnet dataset requires a tokenizer".into()))?; + crate::datasets::sonnet::load_sonnet_dataset( + tok, + config.dataset_path.as_deref(), + config.num_prompts, + config.sonnet_input_len, + config.sonnet_output_len, + config.sonnet_prefix_len, + config.seed, + &config.request_id_prefix, + )? + } + DatasetName::SpeedBench => { + let tok = tokenizer.as_ref().ok_or_else(|| { + BenchError::Config("SPEED-Bench dataset requires a tokenizer".into()) + })?; + let downloaded; + let path = match config.dataset_path.as_deref() { + Some(p) => p, + None => { + downloaded = crate::datasets::speed_bench::download_speed_bench( + config.speed_bench_config, + )?; + downloaded.as_str() + } + }; + let output_len = config.sharegpt_output_len.unwrap_or(config.random_output_len); + crate::datasets::speed_bench::load_speed_bench_dataset( + tok, + path, + config.num_prompts, + output_len, + config.seed, + &config.request_id_prefix, + config.speed_bench_category.as_deref(), + config.no_oversample, + config.disable_shuffle, + config.speed_bench_max_input_len, + )? + } + DatasetName::Hf => { + let tok = tokenizer + .as_ref() + .ok_or_else(|| BenchError::Config("HF dataset requires a tokenizer".into()))?; + let dataset_id = config.dataset_path.as_deref().ok_or_else(|| { + BenchError::Config("--dataset-path is required for --dataset-name hf".into()) + })?; + let (downloaded_path, _config, _split) = + crate::datasets::hf_dataset::download_hf_dataset( + dataset_id, + config.hf_subset.as_deref(), + config.hf_split.as_deref(), + config.num_prompts, + )?; + crate::datasets::hf_dataset::load_hf_dataset( + tok, + &downloaded_path, + config.num_prompts, + config.hf_output_len, + config.seed, + &config.request_id_prefix, + config.hf_text_column.as_deref(), + config.no_oversample, + config.disable_shuffle, + )? + } + DatasetName::Custom => { + let tok = tokenizer + .as_ref() + .ok_or_else(|| BenchError::Config("Custom dataset requires a tokenizer".into()))?; + let path = config.dataset_path.as_deref().ok_or_else(|| { + BenchError::Config("--dataset-path is required for --dataset-name custom".into()) + })?; + crate::datasets::custom::load_custom_dataset( + tok, + path, + config.num_prompts, + config.custom_output_len, + config.seed, + &config.request_id_prefix, + config.no_oversample, + config.disable_shuffle, + )? + } + DatasetName::PrefixRepetition => { + let tok = tokenizer.as_ref().ok_or_else(|| { + BenchError::Config("Prefix repetition dataset requires a tokenizer".into()) + })?; + crate::datasets::prefix_repetition::generate_prefix_repetition_dataset( + tok, + config.num_prompts, + config.prefix_repetition_prefix_len, + config.prefix_repetition_suffix_len, + config.prefix_repetition_num_prefixes, + config.prefix_repetition_output_len, + config.seed, + &config.request_id_prefix, + config.disable_shuffle, + )? + } + DatasetName::RandomRerank => { + let tok = tokenizer.as_ref().ok_or_else(|| { + BenchError::Config("Random rerank dataset requires a tokenizer".into()) + })?; + crate::datasets::random_rerank::generate_random_rerank_dataset( + tok, + config.num_prompts, + config.random_input_len, + config.random_range_ratio, + config.seed, + &config.request_id_prefix, + config.random_batch_size, + config.is_reranker, + )? + } + }; + + let gen_elapsed = gen_start.elapsed(); + println!( + "Generated {} prompts in {:.2}s", + input_requests.len(), + gen_elapsed.as_secs_f64() + ); + + let filtered_count = + filter_requests_by_max_model_len(&mut input_requests, config.max_model_len); + if filtered_count > 0 { + println!( + "Filtered {filtered_count} prompt(s) above --max-model-len {}.", + config.max_model_len.unwrap() + ); + } + if input_requests.is_empty() { + return Err(BenchError::Config( + "No requests remain after applying --max-model-len".into(), + )); + } + + // Dry run: just print stats and exit + if config.dry_run { + let total_input_tokens: usize = input_requests.iter().map(|r| r.prompt_len).sum(); + let total_output_tokens: usize = input_requests.iter().map(|r| r.expected_output_len).sum(); + println!("Dry run stats:"); + println!(" Total prompts: {}", input_requests.len()); + println!(" Total input tokens: {total_input_tokens}"); + println!(" Total expected output tokens: {total_output_tokens}"); + println!( + " Avg input tokens: {:.1}", + total_input_tokens as f64 / input_requests.len() as f64 + ); + println!( + " Avg output tokens: {:.1}", + total_output_tokens as f64 / input_requests.len() as f64 + ); + return Ok(serde_json::json!({"dry_run": true})); + } + + // Build test input from first request + let first = &input_requests[0]; + let test_input = RequestFuncInput { + prompt: first.prompt.clone(), + api_url: config.api_url.clone(), + prompt_len: first.prompt_len, + output_len: first.expected_output_len, + model: model_id.clone(), + model_name: model_name.clone(), + logprobs: config.logprobs, + extra_headers: config.extra_headers.clone(), + extra_body: config.extra_body.clone(), + ignore_eos: config.ignore_eos, + request_id: first.request_id.clone(), + messages: None, + prompt_token_ids: first.prompt_token_ids.clone(), + multi_modal_content: first.multi_modal_content.clone(), + chat_messages_json: first.chat_messages_json.clone(), + prompt_list: first.prompt_list.clone(), + }; + + // Ready check + if config.ready_check_timeout_sec > 0 { + println!("Starting initial single prompt test run..."); + let test_output = wait_for_endpoint( + config.backend, + &client, + &test_input, + config.ready_check_timeout_sec, + 5, + ) + .await?; + if !test_output.success { + return Err(BenchError::Backend(format!( + "Initial test run failed: {}", + test_output.error + ))); + } + println!("Initial test run completed."); + } + + // Verify and fix prompt token lengths against the server's /tokenize endpoint. + // Runs after the ready check so a still-starting server isn't mistaken for a + // tokenize failure, and after the dry-run exit so dry runs stay offline. + // Uses a cache: if a previous run with the same model+server already verified OK, + // skip entirely. Otherwise sample 10 prompts first — if all match, cache and skip. + // If any mismatch, do full verify+fix for all prompts. + // Skip verification when prompt_token_ids are set (token counts are exact by construction). + let has_token_ids = input_requests.first().is_some_and(|r| r.prompt_token_ids.is_some()); + // Python aligns prompts to the server tokenizer for random AND prefix_repetition + // (both are synthetic exact-length datasets). + let verifiable_dataset = matches!( + config.dataset_name, + DatasetName::Random | DatasetName::PrefixRepetition + ); + if verifiable_dataset && has_token_ids && !config.backend.is_pooling() { + println!("Using prompt_token_ids, skipping server-side tokenizer verification."); + } + if verifiable_dataset && !has_token_ids && !config.backend.is_pooling() { + let cache_key = tokenizer_verify_cache_key(&config.base_url, &model_id); + if is_tokenizer_verified(&cache_key) { + println!("Tokenizer verified in previous run (cached), skipping verification."); + } else { + let num_special = + tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0); + match sample_verify_prompts( + &client, + &config.base_url, + &model_id, + &input_requests, + num_special, + &config.extra_headers, + ) + .await? + { + SampleVerifyOutcome::Passed => { + println!("Sample verification passed, skipping full verification."); + mark_tokenizer_verified(&cache_key); + } + SampleVerifyOutcome::Skipped(reason) => { + println!("Server /tokenize unavailable ({reason}), skipping verification."); + } + SampleVerifyOutcome::Mismatch => { + println!("Sample verification found mismatch, running full verify+fix..."); + match verify_and_fix_prompt_lengths( + &client, + &config.base_url, + &model_id, + &mut input_requests, + num_special, + &config.extra_headers, + ) + .await + { + Ok(()) => { + println!( + "All {} prompts verified: exact token length match.", + input_requests.len() + ); + mark_tokenizer_verified(&cache_key); + } + Err(BenchError::TokenizeUnavailable(reason)) => { + println!( + "Server /tokenize became unavailable during verification \ + ({reason}); proceeding with client-side token counts." + ); + } + Err(e) => return Err(e), + } + } + } + } + } + + // Warmup + if config.num_warmups > 0 { + println!("Warming up with {} requests...", config.num_warmups); + run_warmup( + config.backend, + &client, + &test_input, + config.num_warmups, + config.max_concurrency, + config.request_rate, + config.burstiness, + config.seed, + config.disable_tqdm, + ) + .await; + println!("Warmup run completed."); + } + + // Start profiler if requested (immediate mode — no batch threshold) + if config.profile && config.profile_batch_threshold.is_none() { + start_profiler_immediate(&client, &config.base_url, &config.extra_headers).await; + } + + // Threshold-based profiling: spawn background task that polls /metrics + // and triggers start/stop profile when batch size is reached. + // A oneshot channel lets us cancel the polling loop when the benchmark ends + // (e.g. if the threshold is never reached). + let profile_task = if let Some(threshold) = config.profile_batch_threshold { + let poll_client = client.clone(); + let base_url = config.base_url.clone(); + let extra_headers = config.extra_headers.clone(); + let duration_secs = config.profile_duration; + let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel::<()>(); + let handle = tokio::spawn(async move { + profile_on_batch_threshold( + &poll_client, + &base_url, + &extra_headers, + threshold, + duration_secs, + cancel_rx, + ) + .await; + }); + Some((cancel_tx, handle)) + } else { + None + }; + + // Fetch speculative decoding metrics before benchmark + let spec_decode_before = + fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await; + if spec_decode_before.is_some() { + println!("Speculative decoding detected, will collect metrics."); + } + + // Main benchmark + println!("Starting main benchmark run..."); + let distribution = if config.burstiness == 1.0 { + "Poisson process" + } else { + "Gamma distribution" + }; + println!( + "Traffic request rate: {}", + if config.request_rate.is_infinite() { + "inf".to_string() + } else { + format!("{}", config.request_rate) + } + ); + println!("Burstiness factor: {} ({distribution})", config.burstiness); + println!( + "Maximum request concurrency: {}", + config.max_concurrency.unwrap_or(config.num_prompts) + ); + + // Pre-assign LoRA adapters to each request (None when --lora-modules not set). + let lora_assignments = assign_lora_modules( + &config.lora_modules, + config.lora_assignment, + input_requests.len(), + config.seed, + ); + if let (Some(modules), Some(_)) = (config.lora_modules.as_ref(), lora_assignments.as_ref()) { + let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect(); + println!( + "LoRA adapters ({}): {:?} [assignment={:?}]", + modules.len(), + names, + config.lora_assignment + ); + } + + // Compute request schedule + let schedule = compute_schedule( + input_requests.len(), + config.request_rate, + config.burstiness, + config.seed, + config.ramp_up.as_ref(), + ); + + // Progress bar + let pb = if config.disable_tqdm { + None + } else { + let bar = ProgressBar::new(input_requests.len() as u64); + bar.set_style( + ProgressStyle::with_template( + "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({eta})", + ) + .unwrap() + .progress_chars("#>-"), + ); + Some(bar) + }; + + // Semaphore for concurrency control + let semaphore = config.max_concurrency.map(|mc| Arc::new(Semaphore::new(mc))); + + let benchmark_start = Instant::now(); + + // Arc-wrap shared data to avoid cloning per request + let shared_api_url = Arc::new(config.api_url.clone()); + let shared_model_id = Arc::new(model_id.clone()); + let shared_model_name = Arc::new(model_name.clone()); + let shared_extra_headers = Arc::new(config.extra_headers.clone()); + let shared_extra_body = Arc::new(config.extra_body.clone()); + let shared_backend = get_backend(config.backend)?; + let shared_logprobs = config.logprobs; + let shared_ignore_eos = config.ignore_eos; + + // Spawn all request tasks. Store prompt_len alongside handle so + // we can preserve it in the panic recovery path. + let mut handles: Vec<(usize, tokio::task::JoinHandle)> = + Vec::with_capacity(input_requests.len()); + + for (i, (request, delay)) in input_requests.iter().zip(schedule.delays.iter()).enumerate() { + let client = client.clone(); + let backend = shared_backend.clone(); + let sem = semaphore.clone(); + let pb = pb.clone(); + let api_url = shared_api_url.clone(); + let model = shared_model_id.clone(); + let model_name = shared_model_name.clone(); + let extra_headers = shared_extra_headers.clone(); + let extra_body = shared_extra_body.clone(); + + // Per-request LoRA override: the adapter name replaces both `model` and + // `model_name` in the outgoing payload (vLLM routes by name). Tokenizer, + // /v1/models, /tokenize, etc. continue using the base model unchanged. + let lora_name = lora_assignments.as_ref().map(|v| v[i].clone()); + + let prompt = request.prompt.clone(); + let prompt_len = request.prompt_len; + let output_len = request.expected_output_len; + let request_id = request.request_id.clone(); + let prompt_token_ids = request.prompt_token_ids.clone(); + let multi_modal_content = request.multi_modal_content.clone(); + let chat_messages_json = request.chat_messages_json.clone(); + let prompt_list = request.prompt_list.clone(); + + let delay_dur = std::time::Duration::from_secs_f64(*delay); + let bench_start = benchmark_start; + + handles.push(( + prompt_len, + tokio::spawn(async move { + // Sleep until scheduled time + let target = bench_start + delay_dur; + let now = Instant::now(); + if target > now { + tokio::time::sleep(target - now).await; + } + + // Acquire semaphore permit + let _permit = if let Some(ref s) = sem { + Some(s.acquire().await.unwrap()) + } else { + None + }; + + let (req_model, req_model_name) = match lora_name.as_ref() { + Some(name) => (name.to_string(), Some(name.to_string())), + None => ((*model).clone(), (*model_name).clone()), + }; + + let input = RequestFuncInput { + prompt, + api_url: (*api_url).clone(), + prompt_len, + output_len, + model: req_model, + model_name: req_model_name, + logprobs: shared_logprobs, + extra_headers: (*extra_headers).clone(), + extra_body: (*extra_body).clone(), + ignore_eos: shared_ignore_eos, + request_id, + messages: None, + prompt_token_ids, + multi_modal_content, + chat_messages_json, + prompt_list, + }; + + // Send request, retry on connection errors + let max_retries = 3; + let mut output = None; + + for attempt in 0..=max_retries { + // Capture monotonic start time right before sending, + // relative to benchmark_start. Python uses perf_counter() + // for both start_time and ttft/itl, keeping them on the + // same clock. We do the same with Instant. + let request_instant = Instant::now(); + let result = backend.send_request(&input, &client).await; + + match result { + Ok(mut o) => { + // Override SystemTime-based start_time with monotonic offset + o.start_time = + request_instant.duration_since(bench_start).as_secs_f64(); + + // Retry on connection-level failures reported as !success + if !o.success && attempt < max_retries && is_connection_error(&o.error) + { + tokio::time::sleep(std::time::Duration::from_millis( + 500 * (attempt as u64 + 1), + )) + .await; + continue; + } + output = Some(o); + break; + } + Err(e) => { + if attempt < max_retries { + tokio::time::sleep(std::time::Duration::from_millis( + 500 * (attempt as u64 + 1), + )) + .await; + continue; + } + output = Some(RequestFuncOutput { + success: false, + error: e.to_string(), + prompt_len: input.prompt_len, + start_time: request_instant + .duration_since(bench_start) + .as_secs_f64(), + ..Default::default() + }); + break; + } + } + } + + if let Some(pb) = pb { + pb.inc(1); + } + + output.unwrap() + }), + )); + } + + // Collect all results, handling task panics gracefully + let mut outputs = Vec::with_capacity(handles.len()); + for (prompt_len, handle) in handles { + match handle.await { + Ok(output) => outputs.push(output), + Err(e) => { + outputs.push(RequestFuncOutput { + success: false, + error: format!("Task panicked: {e}"), + prompt_len, + ..Default::default() + }); + } + } + } + + if let Some(ref pb) = pb { + pb.finish_and_clear(); + } + + let benchmark_duration = benchmark_start.elapsed().as_secs_f64(); + + // Fetch speculative decoding metrics after benchmark and compute stats + let spec_decode_stats = if spec_decode_before.is_some() { + let spec_decode_after = + fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await; + match (spec_decode_before.as_ref(), spec_decode_after.as_ref()) { + (Some(before), Some(after)) => compute_spec_decode_stats(before, after), + _ => None, + } + } else { + None + }; + + // Calculate metrics — pooling uses a dedicated path matching Python's + // calculate_metrics_for_embeddings (uses server-reported prompt_len, e2el only). + let (mut metrics, actual_output_lens) = if config.backend.is_pooling() { + let m = + calculate_embedding_metrics(&outputs, benchmark_duration, &config.selected_percentiles); + (m, Vec::new()) + } else { + calculate_metrics( + &input_requests, + &outputs, + benchmark_duration, + &config.selected_percentiles, + has_tokenizer, + &config.goodput, + ) + }; + + // Attach steady-state metrics when the closed-loop scope gate passes. + let scope_ok = !config.no_steady_state + && config.max_concurrency.is_some() + && config.request_rate.is_infinite(); + if scope_ok { + let target = config.max_concurrency; + let min_window = config + .steady_state_min_window + .unwrap_or_else(|| (0.1 * benchmark_duration).max(10.0)); + if let Some(window) = steady_state::detect_window( + &outputs, + target, + config.steady_state_threshold, + min_window, + benchmark_duration, + ) { + let ss = steady_state::compute( + &outputs, + &input_requests, + &window, + &config.selected_percentiles, + config.backend.is_pooling(), + ); + metrics.steady_state = Some(ss); + } + } + + // Print console output + print_results( + &metrics, + benchmark_duration, + config, + has_tokenizer, + spec_decode_stats.as_ref(), + ); + + // Stop profiler if requested (immediate mode — no batch threshold) + if config.profile && config.profile_batch_threshold.is_none() { + stop_profiler_immediate(&client, &config.base_url, &config.extra_headers).await; + } + + // Signal the threshold-based profile task that the benchmark is done, then wait + if let Some((cancel_tx, task)) = profile_task { + let _ = cancel_tx.send(()); + if let Err(e) = task.await { + eprintln!("WARNING: Profile background task failed: {e}"); + } + } + + // Build result JSON + let date_iso = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string(); + let dt_filename = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string(); + let result_json = build_result_json( + config, + &metrics, + &actual_output_lens, + &outputs, + benchmark_duration, + &date_iso, + spec_decode_stats.as_ref(), + ); + + // Save if requested + if config.save_result || config.append_result { + let model_for_filename = config.model.as_deref().unwrap_or(&model_id); + let file_name = compute_result_filename(config, model_for_filename, &dt_filename); + + // Create result directory if needed + if let Some(ref dir) = config.result_dir { + std::fs::create_dir_all(dir)?; + } + + if config.append_result { + append_result(&result_json, &file_name)?; + } else { + save_result(&result_json, &file_name)?; + } + } + + Ok(result_json) +} + +fn filter_requests_by_max_model_len( + requests: &mut Vec, + max_model_len: Option, +) -> usize { + let Some(max_model_len) = max_model_len else { + return 0; + }; + + let before = requests.len(); + requests.retain(|request| { + request.prompt_len.saturating_add(request.expected_output_len) <= max_model_len + }); + before - requests.len() +} + +#[cfg(test)] +mod max_model_len_tests { + use std::sync::Arc; + + use super::filter_requests_by_max_model_len; + use crate::datasets::SampleRequest; + + fn sample(prompt_len: usize, expected_output_len: usize) -> SampleRequest { + SampleRequest { + prompt: Arc::from("prompt"), + prompt_len, + expected_output_len, + request_id: None, + ..Default::default() + } + } + + #[test] + fn test_filter_requests_by_max_model_len_keeps_boundary() { + let mut requests = vec![sample(80, 20), sample(81, 20), sample(30, 10)]; + + let filtered = filter_requests_by_max_model_len(&mut requests, Some(100)); + + assert_eq!(filtered, 1); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].prompt_len, 80); + assert_eq!(requests[1].prompt_len, 30); + } + + #[test] + fn test_filter_requests_by_max_model_len_noop_without_limit() { + let mut requests = vec![sample(80, 20), sample(81, 20)]; + + let filtered = filter_requests_by_max_model_len(&mut requests, None); + + assert_eq!(filtered, 0); + assert_eq!(requests.len(), 2); + } +} + +async fn run_warmup( + backend: BackendKind, + client: &reqwest::Client, + test_input: &RequestFuncInput, + num_warmups: usize, + max_concurrency: Option, + request_rate: f64, + burstiness: f64, + seed: u64, + disable_tqdm: bool, +) { + let pb = if disable_tqdm { + None + } else { + let bar = ProgressBar::new(num_warmups as u64); + bar.set_style( + ProgressStyle::with_template("{spinner:.green} Warmup [{bar:30}] {pos}/{len}") + .unwrap() + .progress_chars("#>-"), + ); + Some(bar) + }; + + let semaphore = max_concurrency.map(|mc| Arc::new(Semaphore::new(mc))); + + // Use the same rate-limited scheduling as the main benchmark run + let schedule = compute_schedule(num_warmups, request_rate, burstiness, seed, None); + let start = Instant::now(); + + let mut handles = Vec::with_capacity(num_warmups); + for i in 0..num_warmups { + // Wait until scheduled time + let target = std::time::Duration::from_secs_f64(schedule.delays[i]); + let elapsed = start.elapsed(); + if target > elapsed { + tokio::time::sleep(target - elapsed).await; + } + + let client = client.clone(); + let input = test_input.clone(); + let sem = semaphore.clone(); + let pb = pb.clone(); + handles.push(tokio::spawn(async move { + let _permit = if let Some(ref s) = sem { + Some(s.acquire().await.unwrap()) + } else { + None + }; + if let Ok(b) = get_backend(backend) { + let _ = b.send_request(&input, &client).await; + } + if let Some(pb) = pb { + pb.inc(1); + } + })); + } + + for handle in handles { + let _ = handle.await; + } + + if let Some(pb) = pb { + pb.finish_and_clear(); + } +} + +/// Start the profiler (immediate mode — no batch threshold). +pub(crate) async fn start_profiler_immediate( + client: &reqwest::Client, + base_url: &str, + extra_headers: &Option>, +) { + println!("Starting profiler..."); + let profile_url = format!("{base_url}/start_profile"); + match send_profile_request(client, &profile_url, extra_headers).await { + Ok(true) => println!("Profiler started"), + Ok(false) => eprintln!("WARNING: Profiler start request returned non-success"), + Err(e) => eprintln!("WARNING: Failed to start profiler: {e}"), + } +} + +/// Stop the profiler (immediate mode — no batch threshold). +pub(crate) async fn stop_profiler_immediate( + client: &reqwest::Client, + base_url: &str, + extra_headers: &Option>, +) { + println!("Stopping profiler..."); + let profile_url = format!("{base_url}/stop_profile"); + match send_profile_request(client, &profile_url, extra_headers).await { + Ok(true) => println!("Profiler stopped"), + Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"), + Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"), + } +} + +/// Send a request to the vLLM profiler endpoint (start_profile / stop_profile). +/// Returns Ok(true) if the server responded with 200. +pub(crate) async fn send_profile_request( + client: &reqwest::Client, + url: &str, + extra_headers: &Option>, +) -> Result { + let mut request = client.post(url); + if let Some(headers) = extra_headers { + for (k, v) in headers { + request = request.header(k, v); + } + } + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") { + request = request.header("Authorization", format!("Bearer {api_key}")); + } + + let resp = request + .send() + .await + .map_err(|e| BenchError::Backend(format!("Profile request to {url} failed: {e}")))?; + + Ok(resp.status().is_success()) +} + +/// Parse `vllm:num_requests_running` from the server's `/metrics` Prometheus endpoint. +pub(crate) async fn fetch_num_requests_running( + client: &reqwest::Client, + base_url: &str, +) -> Option { + let url = format!("{base_url}/metrics"); + let resp = client.get(&url).send().await.ok()?; + let text = resp.text().await.ok()?; + for line in text.lines() { + if let Some(rest) = line.strip_prefix("vllm:num_requests_running") { + // Line format: `vllm:num_requests_running{model_name="..."} 42` + // The value is after the last space. + if let Some(val_str) = rest.rsplit(' ').next() + && let Ok(v) = val_str.parse::() + { + return Some(v as usize); + } + } + } + None +} + +/// Poll `/metrics` until `num_requests_running >= threshold`, then start the +/// profiler, wait `duration_secs`, and stop it. If `cancel_rx` fires before +/// the threshold is reached (i.e. the benchmark finished), exits early. +pub(crate) async fn profile_on_batch_threshold( + client: &reqwest::Client, + base_url: &str, + extra_headers: &Option>, + threshold: usize, + duration_secs: f64, + mut cancel_rx: tokio::sync::oneshot::Receiver<()>, +) { + println!( + "Waiting for batch size >= {threshold} before starting profiler \ + (will capture {duration_secs}s)..." + ); + + loop { + if let Some(running) = fetch_num_requests_running(client, base_url).await + && running >= threshold + { + println!("Batch size {running} >= {threshold}, starting profiler..."); + break; + } + // Wait 500ms or until the benchmark signals cancellation + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_millis(500)) => {} + _ = &mut cancel_rx => { + eprintln!( + "NOTE: Benchmark finished before batch threshold {threshold} was reached; \ + profiling skipped." + ); + return; + } + } + } + + let start_url = format!("{base_url}/start_profile"); + match send_profile_request(client, &start_url, extra_headers).await { + Ok(true) => println!("Profiler started"), + Ok(false) => { + eprintln!("WARNING: Profiler start request returned non-success"); + return; + } + Err(e) => { + eprintln!("WARNING: Failed to start profiler: {e}"); + return; + } + } + + // Wait for the capture duration, but exit early if the benchmark finishes + tokio::select! { + _ = tokio::time::sleep(std::time::Duration::from_secs_f64(duration_secs)) => {} + _ = &mut cancel_rx => { + println!("Benchmark finished, stopping profiler early..."); + } + } + + let stop_url = format!("{base_url}/stop_profile"); + match send_profile_request(client, &stop_url, extra_headers).await { + Ok(true) => println!("Profiler stopped after capturing"), + Ok(false) => eprintln!("WARNING: Profiler stop request returned non-success"), + Err(e) => eprintln!("WARNING: Failed to stop profiler: {e}"), + } +} + +/// Verify and fix prompt token lengths against the server's /tokenize + /detokenize. +/// For each prompt, if the server's token count doesn't match that prompt's own +/// target (its generated length plus the special tokens the server adds), adjust +/// the token sequence (pad/truncate) and detokenize, repeating until exact match. +/// Runs up to 64 prompts concurrently. +async fn verify_and_fix_prompt_lengths( + client: &reqwest::Client, + base_url: &str, + model: &str, + requests: &mut [crate::datasets::SampleRequest], + num_special: usize, + extra_headers: &Option>, +) -> Result<()> { + let tokenize_url = Arc::new(format!("{base_url}/tokenize")); + let detokenize_url = Arc::new(format!("{base_url}/detokenize")); + let api_key = Arc::new(std::env::var("OPENAI_API_KEY").ok()); + let extra_headers = Arc::new(extra_headers.clone()); + let concurrency = 64; + let sem = Arc::new(Semaphore::new(concurrency)); + + let pb = Arc::new(ProgressBar::new(requests.len() as u64)); + pb.set_style( + ProgressStyle::with_template( + "{spinner:.green} Verifying [{bar:40.cyan/blue}] {pos}/{len} ({per_sec}, {eta})", + ) + .unwrap() + .progress_chars("#>-"), + ); + + let fixed_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + + // Spawn a task per prompt + let mut handles: Vec>> = + Vec::with_capacity(requests.len()); + for (i, req) in requests.iter().enumerate() { + let client = client.clone(); + let tok_url = tokenize_url.clone(); + let detok_url = detokenize_url.clone(); + let model = model.to_string(); + let prompt = req.prompt.to_string(); // Convert Arc to String for mutation + let expected_input_len = req.prompt_len + num_special; + let api_key = api_key.clone(); + let headers = extra_headers.clone(); + let sem = sem.clone(); + let pb = pb.clone(); + let fixed_count = fixed_count.clone(); + + handles.push(tokio::spawn(async move { + let _permit = sem.acquire().await.unwrap(); + let mut prompt = prompt; + let max_iterations = 20; + let mut needed_fix = false; + // Track systematic offset (e.g., server auto-adds BOS token). + // If the same excess appears twice consecutively, compensate. + let mut last_excess: Option = None; + + for _iter in 0..max_iterations { + let tokens = server_tokenize( + &client, &tok_url, &model, &prompt, &api_key, &headers, i, + ).await?; + + if tokens.len() == expected_input_len { + if needed_fix { + fixed_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + pb.inc(1); + return Ok((i, prompt)); + } + + needed_fix = true; + + // Detect systematic offset: if the server consistently returns + // expected + N tokens (e.g., auto-prepended BOS), reduce target + // by N so the server's re-tokenization lands on the expected count. + let excess = tokens.len().saturating_sub(expected_input_len); + let compensate = if excess > 0 && last_excess == Some(excess) { + if _iter == 1 { + eprintln!( + "Prompt {i}: server consistently adds {excess} extra token(s) \ + (likely BOS), compensating target to {}.", + expected_input_len.saturating_sub(excess), + ); + } + excess + } else { + 0 + }; + last_excess = if excess > 0 { Some(excess) } else { None }; + + let target = expected_input_len.saturating_sub(compensate); + + let mut adjusted = tokens; + if adjusted.is_empty() { + return Err(BenchError::Tokenizer(format!( + "Prompt {i}: server returned no tokens, cannot fix prompt length" + ))); + } + if adjusted.len() < target { + let pad_needed = target - adjusted.len(); + let original_len = adjusted.len(); + for j in 0..pad_needed { + adjusted.push(adjusted[j % original_len]); + } + } else { + adjusted.truncate(target); + } + + prompt = server_detokenize( + &client, &detok_url, &model, &adjusted, &api_key, &headers, i, + ).await?; + } + + Err(BenchError::Tokenizer(format!( + "Prompt {i}: server verification failed to converge after {max_iterations} iterations" + ))) + })); + } + + // Collect results and write back + for handle in handles { + let result = match handle.await { + Ok(r) => r, + Err(e) => { + return Err(BenchError::Tokenizer(format!( + "Verification task panicked: {e}" + ))); + } + }; + let (i, prompt) = result?; + requests[i].prompt = Arc::from(prompt); + // Server-side count the fix loop converged on for this prompt. + requests[i].prompt_len += num_special; + } + + pb.finish_and_clear(); + + let fc = fixed_count.load(std::sync::atomic::Ordering::Relaxed); + if fc > 0 { + println!("Fixed {fc} prompt(s) via server tokenize/detokenize convergence."); + } + + Ok(()) +} + +/// Call the server's /tokenize endpoint and return the token ID list. +/// Retries up to 3 times on transient errors (5xx, connection errors). +async fn server_tokenize( + client: &reqwest::Client, + url: &str, + model: &str, + prompt: &str, + api_key: &Option, + extra_headers: &Option>, + prompt_idx: usize, +) -> Result> { + let max_retries = 3; + + for attempt in 0..=max_retries { + let payload = serde_json::json!({ + "model": model, + "prompt": prompt, + }); + + let mut request = client.post(url).json(&payload); + if let Some(headers) = extra_headers { + for (k, v) in headers { + request = request.header(k, v); + } + } + if let Some(key) = api_key { + request = request.header("Authorization", format!("Bearer {key}")); + } + + let resp = match request.send().await { + Ok(r) => r, + Err(e) => { + if attempt < max_retries { + tokio::time::sleep(std::time::Duration::from_millis( + 500 * (attempt as u64 + 1), + )) + .await; + continue; + } + return Err(BenchError::Tokenizer(format!( + "Server /tokenize failed for prompt {prompt_idx} after {max_retries} retries: {e}" + ))); + } + }; + + if resp.status().is_server_error() && attempt < max_retries { + tokio::time::sleep(std::time::Duration::from_millis(500 * (attempt as u64 + 1))).await; + continue; + } + + if resp.status().is_client_error() { + // Server doesn't expose /tokenize (404, e.g. Dynamo) or a gateway + // rejects it with another 4xx (LLM-d/EPP returns 400) — caller + // should skip verification. 5xx and connection errors stay fatal. + return Err(BenchError::TokenizeUnavailable(format!( + "HTTP {}", + resp.status() + ))); + } + + if !resp.status().is_success() { + return Err(BenchError::Tokenizer(format!( + "Server /tokenize returned HTTP {} for prompt {prompt_idx}", + resp.status() + ))); + } + + let data: serde_json::Value = resp.json().await.map_err(|e| { + BenchError::Tokenizer(format!( + "Failed to parse /tokenize response for prompt {prompt_idx}: {e}" + )) + })?; + + let tokens = data.get("tokens").and_then(|t| t.as_array()).ok_or_else(|| { + BenchError::Tokenizer(format!( + "No 'tokens' array in /tokenize response for prompt {prompt_idx}" + )) + })?; + + return tokens + .iter() + .map(|v| { + v.as_u64().ok_or_else(|| { + BenchError::Tokenizer("Invalid token ID in /tokenize response".to_string()) + }) + }) + .collect(); + } + + unreachable!() +} + +/// Call the server's /detokenize endpoint and return the prompt text. +/// Retries up to 3 times on transient errors (5xx, connection errors). +async fn server_detokenize( + client: &reqwest::Client, + url: &str, + model: &str, + tokens: &[u64], + api_key: &Option, + extra_headers: &Option>, + prompt_idx: usize, +) -> Result { + let max_retries = 3; + + for attempt in 0..=max_retries { + let payload = serde_json::json!({ + "model": model, + "tokens": tokens, + }); + + let mut request = client.post(url).json(&payload); + if let Some(headers) = extra_headers { + for (k, v) in headers { + request = request.header(k, v); + } + } + if let Some(key) = api_key { + request = request.header("Authorization", format!("Bearer {key}")); + } + + let resp = match request.send().await { + Ok(r) => r, + Err(e) => { + if attempt < max_retries { + tokio::time::sleep(std::time::Duration::from_millis( + 500 * (attempt as u64 + 1), + )) + .await; + continue; + } + return Err(BenchError::Tokenizer(format!( + "Server /detokenize failed for prompt {prompt_idx} after {max_retries} retries: {e}" + ))); + } + }; + + if resp.status().is_server_error() && attempt < max_retries { + tokio::time::sleep(std::time::Duration::from_millis(500 * (attempt as u64 + 1))).await; + continue; + } + + if resp.status().is_client_error() { + return Err(BenchError::TokenizeUnavailable(format!( + "HTTP {}", + resp.status() + ))); + } + + if !resp.status().is_success() { + return Err(BenchError::Tokenizer(format!( + "Server /detokenize returned HTTP {} for prompt {prompt_idx}", + resp.status() + ))); + } + + let data: serde_json::Value = resp.json().await.map_err(|e| { + BenchError::Tokenizer(format!( + "Failed to parse /detokenize response for prompt {prompt_idx}: {e}" + )) + })?; + + return data.get("prompt").and_then(|p| p.as_str()).map(|s| s.to_string()).ok_or_else( + || { + BenchError::Tokenizer(format!( + "No 'prompt' in /detokenize response for prompt {prompt_idx}" + )) + }, + ); + } + + unreachable!() +} + +/// Check if an error message indicates a transient connection error worth retrying. +fn is_connection_error(error: &str) -> bool { + let patterns = [ + "connection reset", + "Connection reset", + "connection error", + "broken pipe", + "connect error", + "timed out", + "connection refused", + "Connection refused", + ]; + patterns.iter().any(|p| error.contains(p)) +} + +// --- Tokenizer verification cache --- + +/// Build a cache key from base_url + model_id. +fn tokenizer_verify_cache_key(base_url: &str, model_id: &str) -> String { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut h = DefaultHasher::new(); + base_url.hash(&mut h); + model_id.hash(&mut h); + format!("{:016x}", h.finish()) +} + +/// Cache directory for tokenizer verification markers. +fn verify_cache_dir() -> std::path::PathBuf { + dirs::cache_dir() + .unwrap_or_else(|| std::path::PathBuf::from("/tmp")) + .join("vllm-bench") + .join("verified_tokenizers") +} + +/// Check if tokenizer was previously verified for this model+server. +fn is_tokenizer_verified(cache_key: &str) -> bool { + verify_cache_dir().join(cache_key).exists() +} + +/// Mark tokenizer as verified for this model+server. +fn mark_tokenizer_verified(cache_key: &str) { + let dir = verify_cache_dir(); + let _ = std::fs::create_dir_all(&dir); + let _ = std::fs::write(dir.join(cache_key), ""); +} + +/// Outcome of sample verification against the server's /tokenize. +enum SampleVerifyOutcome { + /// All sampled prompts matched their expected token counts. + Passed, + /// At least one sampled prompt did not match; full verify+fix is needed. + Mismatch, + /// The server cannot tokenize (4xx from /tokenize); verification skipped. + /// Unlike Passed, this must NOT be cached as "verified". + Skipped(String), +} + +/// Sample-verify a small number of prompts against the server's /tokenize. +/// Each prompt is checked against its own generated length (plus the special +/// tokens the server adds), so variable-length datasets +/// (--random-range-ratio < 1.0) and shared prefixes (--random-prefix-len) +/// verify correctly. +async fn sample_verify_prompts( + client: &reqwest::Client, + base_url: &str, + model: &str, + requests: &[crate::datasets::SampleRequest], + num_special: usize, + extra_headers: &Option>, +) -> Result { + let sample_size = 10.min(requests.len()); + let tokenize_url = format!("{base_url}/tokenize"); + let api_key = std::env::var("OPENAI_API_KEY").ok(); + + println!("Sampling {sample_size} prompts for verification..."); + + for (i, request) in requests.iter().enumerate().take(sample_size) { + let tokens = match server_tokenize( + client, + &tokenize_url, + model, + &request.prompt, + &api_key, + extra_headers, + i, + ) + .await + { + Ok(t) => t, + Err(BenchError::TokenizeUnavailable(reason)) => { + return Ok(SampleVerifyOutcome::Skipped(reason)); + } + Err(e) => return Err(e), + }; + + let expected = request.prompt_len + num_special; + if tokens.len() != expected { + println!( + "Prompt {i}: expected {expected} tokens, server returned {}", + tokens.len() + ); + return Ok(SampleVerifyOutcome::Mismatch); + } + } + + Ok(SampleVerifyOutcome::Passed) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_metrics( + drafts: u64, + draft_tokens: u64, + accepted_tokens: u64, + per_pos: &[(u64, u64)], + ) -> SpecDecodeMetrics { + SpecDecodeMetrics { + num_drafts: drafts, + num_draft_tokens: draft_tokens, + num_accepted_tokens: accepted_tokens, + accepted_per_pos: per_pos.iter().copied().collect(), + } + } + + #[test] + fn test_compute_spec_decode_stats_basic() { + let before = make_metrics(100, 300, 200, &[(0, 90), (1, 70), (2, 50)]); + let after = make_metrics(200, 600, 440, &[(0, 180), (1, 150), (2, 120)]); + + let stats = compute_spec_decode_stats(&before, &after).unwrap(); + + assert_eq!(stats.num_drafts, 100); + assert_eq!(stats.draft_tokens, 300); + assert_eq!(stats.accepted_tokens, 240); + assert!((stats.acceptance_rate - 80.0).abs() < 0.01); + assert!((stats.acceptance_length - 3.4).abs() < 0.01); + assert_eq!(stats.per_position_acceptance_rates.len(), 3); + assert!((stats.per_position_acceptance_rates[0] - 0.9).abs() < 0.01); + assert!((stats.per_position_acceptance_rates[1] - 0.8).abs() < 0.01); + assert!((stats.per_position_acceptance_rates[2] - 0.7).abs() < 0.01); + } + + #[test] + fn test_compute_spec_decode_stats_zero_draft_tokens_returns_none() { + let before = make_metrics(0, 100, 80, &[]); + let after = make_metrics(0, 100, 80, &[]); + assert!(compute_spec_decode_stats(&before, &after).is_none()); + } + + #[test] + fn test_compute_spec_decode_stats_zero_drafts() { + let before = make_metrics(0, 0, 0, &[]); + let after = make_metrics(0, 100, 80, &[]); + + let stats = compute_spec_decode_stats(&before, &after).unwrap(); + assert_eq!(stats.num_drafts, 0); + assert_eq!(stats.draft_tokens, 100); + assert_eq!(stats.accepted_tokens, 80); + assert!((stats.acceptance_rate - 80.0).abs() < 0.01); + assert!((stats.acceptance_length - 0.0).abs() < 0.01); + assert!(stats.per_position_acceptance_rates.is_empty()); + } + + #[test] + fn test_compute_spec_decode_stats_new_positions_in_after() { + let before = make_metrics(50, 150, 100, &[(0, 40)]); + let after = make_metrics(150, 450, 320, &[(0, 130), (1, 80)]); + + let stats = compute_spec_decode_stats(&before, &after).unwrap(); + assert_eq!(stats.num_drafts, 100); + assert_eq!(stats.per_position_acceptance_rates.len(), 2); + assert!((stats.per_position_acceptance_rates[0] - 0.9).abs() < 0.01); + assert!((stats.per_position_acceptance_rates[1] - 0.8).abs() < 0.01); + } + + fn lora_names(v: &[Arc]) -> Vec<&str> { + v.iter().map(|s| s.as_ref()).collect() + } + + #[test] + fn test_assign_lora_modules_none_when_unset() { + assert!(assign_lora_modules(&None, LoraAssignment::Random, 5, 0).is_none()); + } + + #[test] + fn test_assign_lora_modules_round_robin_cycles() { + let modules = Some(vec![Arc::::from("a"), Arc::from("b"), Arc::from("c")]); + let out = assign_lora_modules(&modules, LoraAssignment::RoundRobin, 7, 42).unwrap(); + assert_eq!(lora_names(&out), vec!["a", "b", "c", "a", "b", "c", "a"]); + } + + #[test] + fn test_assign_lora_modules_random_is_seed_reproducible() { + let modules = Some(vec![Arc::::from("x"), Arc::from("y"), Arc::from("z")]); + let a = assign_lora_modules(&modules, LoraAssignment::Random, 100, 7).unwrap(); + let b = assign_lora_modules(&modules, LoraAssignment::Random, 100, 7).unwrap(); + assert_eq!(lora_names(&a), lora_names(&b)); + + // Different seed should (almost certainly) produce a different sequence. + let c = assign_lora_modules(&modules, LoraAssignment::Random, 100, 8).unwrap(); + assert_ne!(lora_names(&a), lora_names(&c)); + } + + #[test] + fn test_assign_lora_modules_random_covers_all_modules() { + let modules = Some(vec![ + Arc::::from("a"), + Arc::from("b"), + Arc::from("c"), + Arc::from("d"), + ]); + let out = assign_lora_modules(&modules, LoraAssignment::Random, 1000, 0).unwrap(); + let unique: std::collections::HashSet<&str> = out.iter().map(|s| s.as_ref()).collect(); + assert_eq!( + unique.len(), + 4, + "all 4 adapters should be sampled in 1000 draws" + ); + } + + #[test] + fn test_assign_lora_modules_single_adapter() { + let modules = Some(vec![Arc::::from("only")]); + for assignment in [LoraAssignment::Random, LoraAssignment::RoundRobin] { + let out = assign_lora_modules(&modules, assignment, 5, 0).unwrap(); + assert!(out.iter().all(|s| s.as_ref() == "only")); + } + } +} diff --git a/rust/src/bench/src/cli.rs b/rust/src/bench/src/cli.rs new file mode 100644 index 000000000000..9aa8f388633b --- /dev/null +++ b/rust/src/bench/src/cli.rs @@ -0,0 +1,750 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::fmt; + +use clap::Parser; + +/// Backend type for the benchmark endpoint. +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)] +pub enum BackendKind { + #[value(name = "vllm")] + Vllm, + #[value(name = "openai")] + Openai, + #[value(name = "openai-chat")] + OpenaiChat, + #[value(name = "openai-embeddings")] + OpenaiEmbeddings, + #[value(name = "openai-embeddings-chat")] + OpenaiEmbeddingsChat, + #[value(name = "vllm-pooling")] + VllmPooling, + #[value(name = "vllm-rerank")] + VllmRerank, +} + +impl BackendKind { + pub fn as_str(self) -> &'static str { + match self { + Self::Vllm => "vllm", + Self::Openai => "openai", + Self::OpenaiChat => "openai-chat", + Self::OpenaiEmbeddings => "openai-embeddings", + Self::OpenaiEmbeddingsChat => "openai-embeddings-chat", + Self::VllmPooling => "vllm-pooling", + Self::VllmRerank => "vllm-rerank", + } + } + + /// Return true if the backend is compatible with OpenAI-style API and sampling parameters. + pub fn is_openai_compatible(self) -> bool { + match self { + Self::Vllm | Self::Openai | Self::OpenaiChat => true, + Self::OpenaiEmbeddings + | Self::OpenaiEmbeddingsChat + | Self::VllmPooling + | Self::VllmRerank => false, + } + } + + /// Return true if the backend is a pooling/embedding backend (non-generative). + pub fn is_pooling(self) -> bool { + matches!( + self, + Self::OpenaiEmbeddings + | Self::OpenaiEmbeddingsChat + | Self::VllmPooling + | Self::VllmRerank + ) + } +} + +impl fmt::Display for BackendKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +/// Dataset to benchmark with. +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)] +pub enum DatasetName { + #[value(name = "random")] + Random, + #[value(name = "random-mm")] + RandomMm, + #[value(name = "sharegpt")] + ShareGpt, + #[value(name = "sonnet")] + Sonnet, + #[value(name = "speed-bench")] + SpeedBench, + #[value(name = "hf")] + Hf, + #[value(name = "custom")] + Custom, + #[value(name = "prefix_repetition", alias = "prefix-repetition")] + PrefixRepetition, + #[value(name = "random-rerank")] + RandomRerank, +} + +/// Ramp-up strategy for request rate. +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)] +pub enum RampUpStrategy { + #[value(name = "linear")] + Linear, + #[value(name = "exponential")] + Exponential, +} + +/// Strategy for assigning LoRA modules to requests. +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)] +pub enum LoraAssignment { + #[value(name = "random")] + Random, + #[value(name = "round-robin")] + RoundRobin, +} + +/// SPEED-Bench dataset split/config. +#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)] +pub enum SpeedBenchConfig { + #[value(name = "qualitative")] + Qualitative, + #[value(name = "throughput_1k")] + Throughput1k, + #[value(name = "throughput_2k")] + Throughput2k, + #[value(name = "throughput_8k")] + Throughput8k, + #[value(name = "throughput_16k")] + Throughput16k, + #[value(name = "throughput_32k")] + Throughput32k, +} + +impl SpeedBenchConfig { + pub fn as_str(self) -> &'static str { + match self { + Self::Qualitative => "qualitative", + Self::Throughput1k => "throughput_1k", + Self::Throughput2k => "throughput_2k", + Self::Throughput8k => "throughput_8k", + Self::Throughput16k => "throughput_16k", + Self::Throughput32k => "throughput_32k", + } + } +} + +impl fmt::Display for SpeedBenchConfig { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +/// High-performance benchmark client for vLLM serving endpoints. +#[derive(Parser, Debug, Clone)] +#[command( + name = "vllm-bench", + about = "Benchmark online serving throughput", + version +)] +pub struct Cli { + /// The type of backend or endpoint to use for the benchmark. + #[arg(long, default_value = "openai")] + pub backend: BackendKind, + + /// Server or API base url if not using http host and port. + #[arg(long)] + pub base_url: Option, + + /// Server host. + #[arg(long, default_value = "127.0.0.1")] + pub host: String, + + /// Server port. + #[arg(long, default_value_t = 8000)] + pub port: u16, + + /// API endpoint. Auto-selected based on --backend if not specified. + #[arg(long)] + pub endpoint: Option, + + /// Name of the model. If not specified, will fetch from server. + #[arg(long)] + pub model: Option, + + /// The model name used in the API (for --served-model-name). + #[arg(long)] + pub served_model_name: Option, + + /// Name or path of the tokenizer. + #[arg(long)] + pub tokenizer: Option, + + /// Tokenizer mode (auto, hf, slow, mistral). + #[arg(long, default_value = "auto")] + pub tokenizer_mode: String, + + /// Skip initialization of tokenizer. + #[arg(long, default_value_t = false)] + pub skip_tokenizer_init: bool, + + /// Trust remote code for tokenizer. + #[arg(long, default_value_t = false)] + pub trust_remote_code: bool, + + /// Dataset name. + #[arg(long, default_value = "random")] + pub dataset_name: DatasetName, + + /// General input length for datasets. + #[arg(long)] + pub input_len: Option, + + /// General output length for datasets. + #[arg(long)] + pub output_len: Option, + + /// Maximum model context length. Requests with prompt_len + output_len above this are filtered + /// out. + #[arg(long)] + pub max_model_len: Option, + + /// Random dataset input length. + #[arg(long, default_value_t = 1024)] + pub random_input_len: usize, + + /// Random dataset output length. + #[arg(long, default_value_t = 128)] + pub random_output_len: usize, + + /// Random dataset prefix length. + #[arg(long, default_value_t = 0)] + pub random_prefix_len: usize, + + /// Per-turn input length for turns 1+ in multi-turn mode. + /// 0 = fallback to --random-input-len for all turns. + /// Mirrors sglang bench_multiturn.py --sub-question-input-length. + #[arg(long, default_value_t = 0)] + pub per_turn_input_len: usize, + + /// Range ratio for sampling input/output lengths, matching Python + /// `vllm bench serve`: lengths are drawn uniformly from + /// [len*(1-r), len*(1+r)]. 0.0 (the default) = exact target lengths. + /// Accepts a single float in [0, 1) or a JSON object + /// '{"input": r1, "output": r2}' for independent control. + /// NOTE: semantics changed — the old Rust-only form sampled [len*r, len] + /// with default 1.0; old values like 1.0 are now rejected. + #[arg(long, default_value = "0.0")] + pub random_range_ratio: String, + + /// Batch multiple generated inputs into one request (embeddings/pooling + /// backends only). E.g. 8 sends "input": [t1..t8] per request. Mirrors + /// Python --random-batch-size. Default 1 = no batching. + #[arg(long, default_value_t = 1)] + pub random_batch_size: usize, + + /// random-rerank: the served model is NOT a reranker (embedding-based + /// scoring). Changes query/document length accounting to mirror Python + /// --no-reranker. + #[arg(long, default_value_t = false)] + pub no_reranker: bool, + + /// Bimodal prefix-cache (random dataset): fraction of prompts that are "warm" + /// and reuse a shared cached prefix. 0.0 = off (default). E.g. 0.8 = 80% warm + /// (prefix-cache hit), 20% cold (full prefill). Requires --random-cache-ratio > 0 + /// and --prompt-token-ids. In this mode --random-input-len is the TOTAL length. + #[arg(long, default_value_t = 0.0)] + pub random_cache_hit_fraction: f64, + + /// Bimodal prefix-cache (random dataset): fraction of each WARM prompt's length + /// that is the shared cached prefix. 0.0 = off (default). E.g. 0.95 = 95% cached, + /// 5% unique suffix. Used with --random-cache-hit-fraction. + #[arg(long, default_value_t = 0.0)] + pub random_cache_ratio: f64, + + /// Send prompt as token ID arrays instead of text strings. + /// By default, prompts are decoded to text for maximum + /// compatibility. Enable this for pure vLLM deployments to skip server-side + /// tokenization (faster, exact token counts). + #[arg(long, default_value_t = false)] + pub prompt_token_ids: bool, + + // --- Random multimodal dataset --- + /// Base number of multimodal items (images/videos) per request. + #[arg(long, default_value_t = 1)] + pub random_mm_base_items_per_request: usize, + + /// Range ratio for varying the number of multimodal items per request. + /// Items sampled from [floor(n*(1-r)), ceil(n*(1+r))]. + #[arg(long, default_value_t = 0.0)] + pub random_mm_num_mm_items_range_ratio: f64, + + /// Per-modality hard caps as JSON, e.g. '{"image": 3, "video": 0}'. + #[arg(long, default_value = "{\"image\": 255, \"video\": 1}")] + pub random_mm_limit_mm_per_prompt: String, + + /// Bucket config mapping (height,width,num_frames) to probability. + /// Uses Python-style syntax: '{(256,256,1): 0.5, (720,1280,1): 0.5}'. + /// num_frames=1 means image, num_frames>1 means video. + #[arg(long, default_value = "{(256,256,1): 0.5, (720,1280,1): 0.5}")] + pub random_mm_bucket_config: String, + + /// Enable multimodal chat transformation for datasets that support it. + /// The dataset pre-builds the OpenAI chat `messages` array (text part + + /// multimodal items) at generation time, and the request sends it verbatim. + /// Mirrors Python's --enable-multimodal-chat. Currently applies to random-mm. + #[arg(long, default_value_t = false)] + pub enable_multimodal_chat: bool, + + // --- Custom dataset (JSONL) --- + /// Output tokens per request for the custom dataset. Set to -1 to use the + /// per-line "output_tokens" field from the JSONL file instead. + #[arg(long, default_value_t = 256, allow_negative_numbers = true)] + pub custom_output_len: i64, + + /// Skip applying a chat template to custom dataset prompts. + /// NOTE: the Rust client never renders chat templates client-side, so this + /// is always effectively on; passing it silences the informational notice. + #[arg(long, default_value_t = false)] + pub skip_chat_template: bool, + + // --- Prefix repetition dataset --- + /// Shared-prefix token length for the prefix_repetition dataset. + #[arg(long, default_value_t = 256)] + pub prefix_repetition_prefix_len: usize, + + /// Per-request random suffix token length for the prefix_repetition dataset. + #[arg(long, default_value_t = 256)] + pub prefix_repetition_suffix_len: usize, + + /// Number of distinct shared prefixes for the prefix_repetition dataset. + /// Requests are split evenly across prefixes (num-prompts / num-prefixes each). + #[arg(long, default_value_t = 10)] + pub prefix_repetition_num_prefixes: usize, + + /// Output tokens per request for the prefix_repetition dataset. + #[arg(long, default_value_t = 128)] + pub prefix_repetition_output_len: usize, + + /// Number of prompts to generate. + #[arg(long, default_value_t = 1000)] + pub num_prompts: usize, + + /// Number of requests per second. Use "inf" for all at once. + #[arg(long, default_value_t = f64::INFINITY)] + pub request_rate: f64, + + /// Burstiness factor of request generation. + #[arg(long, default_value_t = 1.0)] + pub burstiness: f64, + + /// Maximum number of concurrent requests. + #[arg(long)] + pub max_concurrency: Option, + + /// Fraction of --max-concurrency at which the steady-state window opens. + /// Range: (0.0, 1.0]. Used only when --max-concurrency is set and + /// --request-rate is inf. + #[arg(long, default_value_t = 0.95)] + pub steady_state_threshold: f64, + + /// Minimum steady-state window duration in seconds. Below this, a warning + /// is attached. If unset, computed as max(10.0, 0.1 * run_duration). + #[arg(long)] + pub steady_state_min_window: Option, + + /// Disable steady-state metrics computation entirely. + #[arg(long, default_value_t = false)] + pub no_steady_state: bool, + + /// Disable tqdm progress bar. + #[arg(long, default_value_t = false)] + pub disable_tqdm: bool, + + /// Number of warmup requests. + #[arg(long, default_value_t = 0)] + pub num_warmups: usize, + + /// Use vLLM profiling. --profiler-config must be provided on the server. + #[arg(long, default_value_t = false)] + pub profile: bool, + + /// Minimum server batch size (num_requests_running) before starting the + /// profiler. When set, profiling is deferred until the /metrics endpoint + /// reports at least this many running requests, then captures for + /// --profile-duration seconds. Requires --profile. + #[arg(long)] + pub profile_batch_threshold: Option, + + /// How many seconds to capture once the batch threshold is reached. + /// Defaults to 5. Requires --profile and --profile-batch-threshold. + #[arg(long, default_value_t = 5.0)] + pub profile_duration: f64, + + /// Save benchmark results to a JSON file. + #[arg(long, default_value_t = false)] + pub save_result: bool, + + /// Save detailed per-request results. + #[arg(long, default_value_t = false)] + pub save_detailed: bool, + + /// Directory to save benchmark JSON results. + #[arg(long)] + pub result_dir: Option, + + /// Filename to save benchmark JSON results. + #[arg(long)] + pub result_filename: Option, + + /// Random seed. + #[arg(long, default_value_t = 0)] + pub seed: u64, + + /// Set ignore_eos flag when sending the benchmark request. + #[arg(long, default_value_t = false)] + pub ignore_eos: bool, + + /// Comma-separated list of metrics to report percentiles for. + #[arg(long)] + pub percentile_metrics: Option, + + /// Comma-separated list of percentiles for selected metrics. + #[arg(long, default_value = "99")] + pub metric_percentiles: String, + + /// Comma-separated list of extra percentiles to show in sweep summaries. + #[arg(long)] + pub sweep_summary_percentiles: Option, + + /// The label (prefix) of the benchmark results. + #[arg(long)] + pub label: Option, + + /// Number of logprobs-per-token to compute. + #[arg(long)] + pub logprobs: Option, + + /// Prefix for request IDs. + #[arg(long)] + pub request_id_prefix: Option, + + /// Maximum time to wait for endpoint readiness in seconds. + #[arg(long, default_value_t = 0)] + pub ready_check_timeout_sec: u64, + + /// Key-value pairs for extra headers (KEY=VALUE). + #[arg(long = "header", num_args = 1..)] + pub headers: Option>, + + /// JSON string for extra body parameters. + #[arg(long)] + pub extra_body: Option, + + /// Key-value pairs for metadata (KEY=VALUE). + #[arg(long = "metadata", num_args = 1..)] + pub metadata: Option>, + + /// Dry run: only generate dataset and print stats, don't benchmark. + #[arg(long, default_value_t = false)] + pub dry_run: bool, + + // --- Sampling parameters --- + /// Top-p sampling parameter. Only affects openai-compatible backends. + #[arg(long)] + pub top_p: Option, + + /// Top-k sampling parameter. Only affects openai-compatible backends. + #[arg(long)] + pub top_k: Option, + + /// Min-p sampling parameter. Only affects openai-compatible backends. + #[arg(long)] + pub min_p: Option, + + /// Temperature sampling parameter. Only affects openai-compatible backends. + #[arg(long)] + pub temperature: Option, + + /// Frequency penalty sampling parameter. Only affects openai-compatible backends. + #[arg(long)] + pub frequency_penalty: Option, + + /// Presence penalty sampling parameter. Only affects openai-compatible backends. + #[arg(long)] + pub presence_penalty: Option, + + /// Repetition penalty sampling parameter. Only affects openai-compatible backends. + #[arg(long)] + pub repetition_penalty: Option, + + // --- SSL --- + /// Disable SSL certificate verification. + #[arg(long, default_value_t = false)] + pub insecure: bool, + + // --- Ramp-up --- + /// Ramp-up strategy for request rate (linear or exponential). + #[arg(long)] + pub ramp_up_strategy: Option, + + /// Starting request rate for ramp-up (RPS). + #[arg(long)] + pub ramp_up_start_rps: Option, + + /// Ending request rate for ramp-up (RPS). + #[arg(long)] + pub ramp_up_end_rps: Option, + + // --- Goodput --- + /// Service level objectives for goodput as "KEY:VALUE" pairs (e.g. ttft:100 tpot:50 e2el:500). + /// Values are in milliseconds. + #[arg(long = "goodput", num_args = 1..)] + pub goodput: Option>, + + // --- Result --- + /// Append the benchmark result to the existing JSON file. + #[arg(long, default_value_t = false)] + pub append_result: bool, + + // --- ShareGPT dataset --- + /// Path to dataset file (required for sharegpt dataset). + #[arg(long)] + pub dataset_path: Option, + + /// Override output length for ShareGPT dataset. + #[arg(long)] + pub sharegpt_output_len: Option, + + /// Do not oversample if dataset is smaller than num_prompts. + #[arg(long, default_value_t = false)] + pub no_oversample: bool, + + /// Do not shuffle the dataset. + #[arg(long, default_value_t = false)] + pub disable_shuffle: bool, + + // --- Sonnet dataset --- + /// Number of input tokens per request (sonnet dataset). + #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_INPUT_LEN)] + pub sonnet_input_len: usize, + + /// Number of output tokens per request (sonnet dataset). + #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_OUTPUT_LEN)] + pub sonnet_output_len: usize, + + /// Number of prefix tokens shared across requests (sonnet dataset). + #[arg(long, default_value_t = crate::datasets::sonnet::DEFAULT_PREFIX_LEN)] + pub sonnet_prefix_len: usize, + + /// SPEED-Bench config/split (qualitative, throughput_1k, throughput_2k, throughput_8k, + /// throughput_16k, throughput_32k). + #[arg(long, default_value = "qualitative")] + pub speed_bench_config: SpeedBenchConfig, + + /// Filter SPEED-Bench by category (e.g. low_entropy, high_entropy, coding, math). + #[arg(long)] + pub speed_bench_category: Option, + + /// Truncate SPEED-Bench prompts to at most this many tokens. + /// Useful for creating custom input lengths from larger splits (e.g. --speed-bench-config + /// throughput_16k --speed-bench-max-input-len 10240). + #[arg(long)] + pub speed_bench_max_input_len: Option, + + // --- HuggingFace dataset --- + /// HuggingFace dataset split (e.g. train, test, validation). + #[arg(long)] + pub hf_split: Option, + + /// HuggingFace dataset subset/config name. + #[arg(long)] + pub hf_subset: Option, + + /// Fixed output length for HF dataset requests (overrides dataset-derived length). + #[arg(long)] + pub hf_output_len: Option, + + /// Column name containing the prompt text. Auto-detected if not specified. + #[arg(long)] + pub hf_text_column: Option, + + // --- Compare mode --- + /// Compare two benchmark result JSON files (e.g. --compare a.json b.json). + /// Prints side-by-side metrics with delta and % change. Skips benchmarking. + #[arg(long = "compare", num_args = 2, value_names = ["FILE_A", "FILE_B"])] + pub compare: Option>, + + // --- Sweep mode --- + /// Sweep over max-concurrency values (comma-separated, e.g. --sweep-max-concurrency + /// 1,10,50,100,500). + #[arg(long)] + pub sweep_max_concurrency: Option, + + /// When sweeping concurrency, set num_prompts = concurrency * this factor for each sweep + /// point. + #[arg(long)] + pub sweep_num_prompts_factor: Option, + + /// Sweep over request-rate values (comma-separated, supports "inf", e.g. --sweep-request-rate + /// 1,10,100,inf). + #[arg(long)] + pub sweep_request_rate: Option, + + /// Reset the server's prefix cache before each sweep iteration. + /// Requires VLLM_SERVER_DEV_MODE=1 on the vLLM server. + #[arg(long, default_value_t = false)] + pub reset_prefix_cache: bool, + + // --- Multi-run --- + /// Number of benchmark runs for statistical aggregation. + #[arg(long, default_value_t = 1)] + pub num_runs: usize, + + // --- Multi-turn conversation benchmark --- + /// Enable multi-turn conversation benchmark mode. + #[arg(long, default_value_t = false)] + pub multi_turn: bool, + + /// Number of turns per conversation in synthetic multi-turn mode. + #[arg(long, default_value_t = 3)] + pub multi_turn_num_turns: usize, + + /// Minimum turns per conversation. 0 = use --multi-turn-num-turns. + #[arg(long, default_value_t = 0)] + pub multi_turn_min_turns: usize, + + /// Maximum turns per conversation. + /// For synthetic multi-turn, 0 = use --multi-turn-num-turns. + /// For ShareGPT multi-turn, 0 = uncapped. + #[arg(long, default_value_t = 0)] + pub multi_turn_max_turns: usize, + + /// Number of concurrent conversations (defaults to max-concurrency or num-prompts). + #[arg(long)] + pub multi_turn_concurrency: Option, + + /// Delay between turns in milliseconds (simulates user think time). + #[arg(long, default_value_t = 0)] + pub multi_turn_delay_ms: u64, + + /// Fraction of per-turn input tokens shared across ALL conversations (0.0–1.0). + /// When > 0, enables prefix sharing mode: each turn sends a fixed-length message + /// (no history accumulation). Only works with --dataset-name random. + #[arg(long, default_value_t = 0.0)] + pub multi_turn_prefix_global_ratio: f64, + + /// Fraction of per-turn input tokens shared within each conversation (0.0–1.0). + /// When > 0, enables prefix sharing mode: each turn sends a fixed-length message + /// (no history accumulation). Only works with --dataset-name random. + #[arg(long, default_value_t = 0.0)] + pub multi_turn_prefix_conversation_ratio: f64, + + // --- LoRA --- + /// LoRA adapter names registered on the server (server-side + /// `--lora-modules name=path`). Each request's `model` field is rewritten + /// to one of these names; tokenizer and other endpoints keep using --model. + /// In multi-turn mode, one adapter is assigned per conversation (sticky + /// across turns). + #[arg(long = "lora-modules", num_args = 1..)] + pub lora_modules: Option>, + + /// Strategy for assigning LoRA adapters to requests. + /// 'random' (default) picks uniformly at random; 'round-robin' cycles + /// through `--lora-modules` deterministically (i % N). + #[arg(long = "lora-assignment", default_value = "random")] + pub lora_assignment: LoraAssignment, +} + +impl Cli { + /// Resolve the base URL from explicit --base-url or from --host/--port. + pub fn resolve_base_url(&self) -> String { + if let Some(ref base) = self.base_url { + base.clone() + } else { + format!("http://{}:{}", self.host, self.port) + } + } + + /// Resolve the API endpoint, auto-selecting based on backend if not explicit. + pub fn resolve_endpoint(&self) -> String { + if let Some(ref ep) = self.endpoint { + return ep.clone(); + } + match self.backend { + BackendKind::OpenaiChat => "/v1/chat/completions".to_string(), + BackendKind::Vllm | BackendKind::Openai => "/v1/completions".to_string(), + BackendKind::OpenaiEmbeddings | BackendKind::OpenaiEmbeddingsChat => { + "/v1/embeddings".to_string() + } + BackendKind::VllmPooling => "/v1/pooling".to_string(), + BackendKind::VllmRerank => "/v1/rerank".to_string(), + } + } + + /// Resolve the full API URL. + pub fn resolve_api_url(&self) -> String { + format!("{}{}", self.resolve_base_url(), self.resolve_endpoint()) + } + + /// Parse extra headers from KEY=VALUE pairs. + pub fn parse_headers( + &self, + ) -> crate::error::Result>> { + match &self.headers { + None => Ok(None), + Some(items) => { + let mut map = std::collections::HashMap::new(); + for item in items { + let (k, v) = item.split_once('=').ok_or_else(|| { + crate::error::BenchError::Config( + "Invalid header format. Use KEY=VALUE".into(), + ) + })?; + map.insert(k.trim().to_string(), v.trim().to_string()); + } + Ok(Some(map)) + } + } + } + + /// Parse extra body JSON. + pub fn parse_extra_body(&self) -> crate::error::Result> { + match &self.extra_body { + None => Ok(None), + Some(s) => { + let v: serde_json::Value = serde_json::from_str(s).map_err(|e| { + crate::error::BenchError::Config(format!("Invalid --extra-body JSON: {e}")) + })?; + Ok(Some(v)) + } + } + } + + /// Generate the request ID prefix (auto-generate if not provided). + pub fn get_request_id_prefix(&self) -> String { + self.request_id_prefix + .clone() + .unwrap_or_else(|| format!("bench-{}-", &uuid::Uuid::new_v4().to_string()[..8])) + } + + /// Resolve input/output lengths, applying --input-len/--output-len overrides. + pub fn resolved_random_input_len(&self) -> usize { + self.input_len.unwrap_or(self.random_input_len) + } + + pub fn resolved_random_output_len(&self) -> usize { + self.output_len.unwrap_or(self.random_output_len) + } + + pub fn resolved_per_turn_input_len(&self) -> usize { + if self.per_turn_input_len > 0 { + self.per_turn_input_len + } else { + self.resolved_random_input_len() + } + } +} diff --git a/rust/src/bench/src/compare.rs b/rust/src/bench/src/compare.rs new file mode 100644 index 000000000000..6d41b72b85cb --- /dev/null +++ b/rust/src/bench/src/compare.rs @@ -0,0 +1,302 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use crate::error::{BenchError, Result}; + +/// Metric definition for comparison: name, JSON key, and whether lower is better. +struct MetricDef { + label: &'static str, + key: &'static str, + lower_is_better: bool, +} + +const METRICS: &[MetricDef] = &[ + MetricDef { + label: "Request throughput (req/s)", + key: "request_throughput", + lower_is_better: false, + }, + MetricDef { + label: "Output throughput (tok/s)", + key: "output_throughput", + lower_is_better: false, + }, + MetricDef { + label: "Total token throughput (tok/s)", + key: "total_token_throughput", + lower_is_better: false, + }, + MetricDef { + label: "Peak output tokens/s", + key: "max_output_tokens_per_s", + lower_is_better: false, + }, + MetricDef { + label: "Peak concurrent requests", + key: "max_concurrent_requests", + lower_is_better: false, + }, + MetricDef { + label: "Mean TTFT (ms)", + key: "mean_ttft_ms", + lower_is_better: true, + }, + MetricDef { + label: "Median TTFT (ms)", + key: "median_ttft_ms", + lower_is_better: true, + }, + MetricDef { + label: "P99 TTFT (ms)", + key: "p99_ttft_ms", + lower_is_better: true, + }, + MetricDef { + label: "Mean TPOT (ms)", + key: "mean_tpot_ms", + lower_is_better: true, + }, + MetricDef { + label: "Median TPOT (ms)", + key: "median_tpot_ms", + lower_is_better: true, + }, + MetricDef { + label: "P99 TPOT (ms)", + key: "p99_tpot_ms", + lower_is_better: true, + }, + MetricDef { + label: "Mean ITL (ms)", + key: "mean_itl_ms", + lower_is_better: true, + }, + MetricDef { + label: "Median ITL (ms)", + key: "median_itl_ms", + lower_is_better: true, + }, + MetricDef { + label: "P99 ITL (ms)", + key: "p99_itl_ms", + lower_is_better: true, + }, + MetricDef { + label: "Mean E2EL (ms)", + key: "mean_e2el_ms", + lower_is_better: true, + }, + MetricDef { + label: "Median E2EL (ms)", + key: "median_e2el_ms", + lower_is_better: true, + }, + MetricDef { + label: "P99 E2EL (ms)", + key: "p99_e2el_ms", + lower_is_better: true, + }, + MetricDef { + label: "Completed requests", + key: "completed", + lower_is_better: false, + }, + MetricDef { + label: "Failed requests", + key: "failed", + lower_is_better: true, + }, + MetricDef { + label: "Duration (s)", + key: "duration", + lower_is_better: true, + }, +]; + +const STEADY_STATE_METRICS: &[MetricDef] = &[ + MetricDef { + label: "SS Request throughput (req/s)", + key: "request_throughput", + lower_is_better: false, + }, + MetricDef { + label: "SS Output throughput (tok/s)", + key: "output_throughput", + lower_is_better: false, + }, + MetricDef { + label: "SS Input throughput (tok/s)", + key: "input_throughput", + lower_is_better: false, + }, + MetricDef { + label: "SS Total token throughput (tok/s)", + key: "total_token_throughput", + lower_is_better: false, + }, + MetricDef { + label: "SS Mean TTFT (ms)", + key: "mean_ttft_ms", + lower_is_better: true, + }, + MetricDef { + label: "SS Median TTFT (ms)", + key: "median_ttft_ms", + lower_is_better: true, + }, + MetricDef { + label: "SS Mean TPOT (ms)", + key: "mean_tpot_ms", + lower_is_better: true, + }, + MetricDef { + label: "SS Median TPOT (ms)", + key: "median_tpot_ms", + lower_is_better: true, + }, + MetricDef { + label: "SS P90 TPOT (ms)", + key: "p90_tpot_ms", + lower_is_better: true, + }, + MetricDef { + label: "SS P99 TPOT (ms)", + key: "p99_tpot_ms", + lower_is_better: true, + }, +]; + +/// Compare two benchmark result JSON files and print a side-by-side table. +pub fn compare_results(file_a: &str, file_b: &str) -> Result<()> { + let json_a = load_result_json(file_a)?; + let json_b = load_result_json(file_b)?; + + // Print header with file context + let model_a = json_a.get("model_id").and_then(|v| v.as_str()).unwrap_or("?"); + let model_b = json_b.get("model_id").and_then(|v| v.as_str()).unwrap_or("?"); + let date_a = json_a.get("date").and_then(|v| v.as_str()).unwrap_or("?"); + let date_b = json_b.get("date").and_then(|v| v.as_str()).unwrap_or("?"); + + println!("{:=^90}", " Benchmark Comparison "); + println!(" A: {} (model: {}, date: {})", file_a, model_a, date_a); + println!(" B: {} (model: {}, date: {})", file_b, model_b, date_b); + println!(); + + // Print comparison table + println!( + "{:<35} {:>12} {:>12} {:>10} {:>8}", + "Metric", "A", "B", "Delta", "Change" + ); + println!("{:-<35} {:->12} {:->12} {:->10} {:->8}", "", "", "", "", ""); + + for metric in METRICS { + let val_a = get_f64(&json_a, metric.key); + let val_b = get_f64(&json_b, metric.key); + + match (val_a, val_b) { + (Some(a), Some(b)) => print_diff_row(metric, a, b), + _ => { + // One or both values missing — skip + } + } + } + + // Steady-state section — both sides must have the block; otherwise render N/A. + let ss_a = json_a.get("steady_state"); + let ss_b = json_b.get("steady_state"); + let both_present = + matches!(ss_a, Some(v) if !v.is_null()) && matches!(ss_b, Some(v) if !v.is_null()); + + println!(); + println!("{:=^70}", " Steady-State Comparison "); + if !both_present { + println!("N/A — one or both runs have no steady-state window"); + } else { + let ss_a = ss_a.unwrap(); + let ss_b = ss_b.unwrap(); + for m in STEADY_STATE_METRICS { + let a = ss_a.get(m.key).and_then(|v| v.as_f64()); + let b = ss_b.get(m.key).and_then(|v| v.as_f64()); + match (a, b) { + (Some(a), Some(b)) => print_diff_row(m, a, b), + _ => println!("{:<35} N/A", m.label), + } + } + } + + println!("{:=<90}", ""); + println!(); + println!("Legend: + = improvement, - = regression (relative to A → B)"); + + Ok(()) +} + +fn print_diff_row(metric: &MetricDef, a: f64, b: f64) { + let delta = b - a; + let pct = if a.abs() > 1e-10 { + (delta / a) * 100.0 + } else if b.abs() > 1e-10 { + f64::INFINITY + } else { + 0.0 + }; + + // Determine if change is good/bad/neutral + let marker = if delta.abs() < 1e-10 { + " " + } else if metric.lower_is_better { + if delta < 0.0 { "+" } else { "-" } + } else if delta > 0.0 { + "+" + } else { + "-" + }; + + let delta_str = format_delta(delta); + let pct_str = if pct.is_infinite() { + "inf%".to_string() + } else { + format!("{:+.1}%", pct) + }; + + println!( + "{:<35} {:>12} {:>12} {:>10} {:>7}{}", + metric.label, + format_value(a), + format_value(b), + delta_str, + pct_str, + marker, + ); +} + +fn load_result_json(path: &str) -> Result { + let content = std::fs::read_to_string(path) + .map_err(|e| BenchError::Config(format!("Cannot read result file '{path}': {e}")))?; + + // Support JSONL: take the last line (most recent run) + let json_str = content.lines().rfind(|l| !l.trim().is_empty()).unwrap_or(&content); + + serde_json::from_str(json_str) + .map_err(|e| BenchError::Config(format!("Cannot parse JSON from '{path}': {e}"))) +} + +fn get_f64(json: &serde_json::Value, key: &str) -> Option { + json.get(key).and_then(|v| v.as_f64()) +} + +fn format_value(v: f64) -> String { + if v == v.floor() && v.abs() < 1e12 { + format!("{}", v as i64) + } else { + format!("{:.2}", v) + } +} + +fn format_delta(d: f64) -> String { + if d == d.floor() && d.abs() < 1e12 { + format!("{:+}", d as i64) + } else { + format!("{:+.2}", d) + } +} diff --git a/rust/src/bench/src/config.rs b/rust/src/bench/src/config.rs new file mode 100644 index 000000000000..9696612c30fd --- /dev/null +++ b/rust/src/bench/src/config.rs @@ -0,0 +1,1116 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig}; +use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt}; +use crate::error::{BenchError, Result}; + +/// Parsed goodput SLO configuration. +#[derive(Debug, Clone, Default)] +pub struct GoodputConfig { + /// TTFT SLO in milliseconds (None = not checked). + pub ttft_ms: Option, + /// TPOT SLO in milliseconds (None = not checked). + pub tpot_ms: Option, + /// E2EL SLO in milliseconds (None = not checked). + pub e2el_ms: Option, +} + +impl GoodputConfig { + pub fn is_empty(&self) -> bool { + self.ttft_ms.is_none() && self.tpot_ms.is_none() && self.e2el_ms.is_none() + } +} + +/// Ramp-up configuration. +#[derive(Debug, Clone)] +pub struct RampUpConfig { + pub strategy: RampUpStrategy, + pub start_rps: f64, + pub end_rps: f64, +} + +/// Range ratio for sampling input/output lengths, matching Python +/// `vllm bench serve`: lengths are drawn uniformly from [len*(1-r), len*(1+r)]. +/// A single float applies to both; the JSON form '{"input": r1, "output": r2}' +/// controls them independently. Each ratio must be in [0, 1). +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct RangeRatio { + pub input: f64, + pub output: f64, +} + +impl RangeRatio { + /// Parse `--random-range-ratio`: a bare float or a JSON object with + /// "input" and "output" keys. + pub fn parse(raw: &str) -> Result { + let trimmed = raw.trim(); + let (input, output) = if let Ok(v) = trimmed.parse::() { + (v, v) + } else { + let v: serde_json::Value = serde_json::from_str(trimmed).map_err(|_| { + BenchError::Config(format!( + "Invalid --random-range-ratio '{raw}': expected a float or \ + '{{\"input\": r1, \"output\": r2}}'" + )) + })?; + let obj = v.as_object().ok_or_else(|| { + BenchError::Config( + "--random-range-ratio JSON form must be an object with \ + 'input' and 'output' keys" + .into(), + ) + })?; + let get = |key: &str| -> Result { + obj.get(key).and_then(|v| v.as_f64()).ok_or_else(|| { + BenchError::Config(format!( + "--random-range-ratio JSON form must contain a numeric '{key}' key" + )) + }) + }; + (get("input")?, get("output")?) + }; + + for (name, r) in [("input", input), ("output", output)] { + if !(0.0..1.0).contains(&r) { + let hint = if r == 1.0 { + " NOTE: semantics now match Python vllm bench serve — lengths are \ + sampled from [len*(1-r), len*(1+r)] and 0.0 means fixed length. \ + The old Rust-only default 1.0 ([len*r, len]) is no longer valid." + } else { + "" + }; + return Err(BenchError::Config(format!( + "--random-range-ratio {name} ratio must be in [0, 1), got {r}.{hint}" + ))); + } + } + Ok(Self { input, output }) + } + + /// Sampling interval for input lengths: [floor(len*(1-r)), ceil(len*(1+r))]. + pub fn input_bounds(&self, len: usize) -> (usize, usize) { + let low = ((len as f64) * (1.0 - self.input)).floor() as usize; + let high = ((len as f64) * (1.0 + self.input)).ceil() as usize; + (low, high) + } + + /// Sampling interval for output lengths, clamped to at least 1 token. + pub fn output_bounds(&self, len: usize) -> (usize, usize) { + let low = (((len as f64) * (1.0 - self.output)).floor() as usize).max(1); + let high = (((len as f64) * (1.0 + self.output)).ceil() as usize).max(1); + (low, high) + } + + /// True when both ratios are 0 (every request uses the exact target lengths). + pub fn is_fixed(&self) -> bool { + self.input == 0.0 && self.output == 0.0 + } +} + +/// Validated benchmark configuration derived from CLI args. +#[derive(Debug, Clone)] +pub struct BenchConfig { + pub backend: BackendKind, + pub base_url: String, + pub api_url: String, + pub model: Option, + pub model_name: Option, + pub tokenizer_id: Option, + #[allow(dead_code)] + pub tokenizer_mode: String, + pub trust_remote_code: bool, + pub skip_tokenizer_init: bool, + pub dataset_name: DatasetName, + pub dataset_path: Option, + pub max_model_len: Option, + pub random_input_len: usize, + pub random_output_len: usize, + pub random_prefix_len: usize, + pub random_range_ratio: RangeRatio, + pub random_cache_hit_fraction: f64, + pub random_cache_ratio: f64, + /// Inputs per request for embeddings/pooling backends (1 = no batching). + pub random_batch_size: usize, + /// random-rerank: whether the served model is a reranker (default true). + pub is_reranker: bool, + pub custom_output_len: i64, + pub prefix_repetition_prefix_len: usize, + pub prefix_repetition_suffix_len: usize, + pub prefix_repetition_num_prefixes: usize, + pub prefix_repetition_output_len: usize, + pub sharegpt_output_len: Option, + pub sonnet_input_len: usize, + pub sonnet_output_len: usize, + pub sonnet_prefix_len: usize, + pub no_oversample: bool, + pub disable_shuffle: bool, + pub num_prompts: usize, + pub request_rate: f64, + pub burstiness: f64, + pub max_concurrency: Option, + pub steady_state_threshold: f64, + pub steady_state_min_window: Option, + pub no_steady_state: bool, + pub disable_tqdm: bool, + pub num_warmups: usize, + pub profile: bool, + pub profile_batch_threshold: Option, + pub profile_duration: f64, + pub save_result: bool, + pub save_detailed: bool, + pub append_result: bool, + pub result_dir: Option, + pub result_filename: Option, + pub seed: u64, + pub ignore_eos: bool, + pub insecure: bool, + pub selected_percentile_metrics: Vec, + pub selected_percentiles: Vec, + pub sweep_summary_percentiles: Vec, + pub label: Option, + pub logprobs: Option, + pub request_id_prefix: String, + pub ready_check_timeout_sec: u64, + pub extra_headers: Option>, + pub extra_body: Option, + pub metadata: Option>, + pub dry_run: bool, + pub goodput: GoodputConfig, + pub ramp_up: Option, + pub multi_turn: bool, + pub multi_turn_num_turns: usize, + pub multi_turn_min_turns: usize, + pub multi_turn_max_turns: usize, + pub sharegpt_multi_turn_max_turns: Option, + pub per_turn_input_len: usize, + pub multi_turn_concurrency: Option, + pub multi_turn_delay_ms: u64, + pub multi_turn_prefix_global_ratio: f64, + pub multi_turn_prefix_conversation_ratio: f64, + pub speed_bench_config: SpeedBenchConfig, + pub speed_bench_category: Option, + pub speed_bench_max_input_len: Option, + pub hf_split: Option, + pub hf_subset: Option, + pub hf_output_len: Option, + pub hf_text_column: Option, + pub reset_prefix_cache: bool, + pub prompt_token_ids: bool, + // --- Random multimodal dataset --- + pub random_mm_base_items_per_request: usize, + pub random_mm_num_mm_items_range_ratio: f64, + pub random_mm_limit: MmLimitPerPrompt, + pub random_mm_buckets: Vec<(MmBucketKey, f64)>, + /// Datasets that support it pre-build the chat `messages` array + /// (text + multimodal parts) instead of prompt + separate mm content. + pub enable_multimodal_chat: bool, + /// LoRA adapter names. None = no LoRA routing (use --model directly). + /// Stored as Arc so per-request override is a cheap clone. + pub lora_modules: Option>>, + pub lora_assignment: LoraAssignment, +} + +impl BenchConfig { + pub fn from_cli(cli: &Cli) -> Result { + if cli.burstiness <= 0.0 { + return Err(BenchError::Config("Burstiness must be positive".into())); + } + + if cli.num_prompts == 0 { + return Err(BenchError::Config( + "--num-prompts must be at least 1".into(), + )); + } + + if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() { + return Err(BenchError::Config( + "--request-rate must be positive (or inf)".into(), + )); + } + if cli.max_model_len == Some(0) { + return Err(BenchError::Config( + "--max-model-len must be at least 1".into(), + )); + } + + let base_url = cli.resolve_base_url(); + let api_url = cli.resolve_api_url(); + + let extra_headers = cli.parse_headers()?; + let mut extra_body = cli.parse_extra_body()?; + + // Merge sampling parameters into extra_body (matches Python behavior). + // Python collects non-None sampling params and merges them UNDER extra_body, + // meaning extra_body keys take precedence over sampling params. + { + let mut sampling_params = serde_json::Map::new(); + if let Some(v) = cli.top_p { + sampling_params.insert("top_p".into(), serde_json::json!(v)); + } + if let Some(v) = cli.top_k { + sampling_params.insert("top_k".into(), serde_json::json!(v)); + } + if let Some(v) = cli.min_p { + sampling_params.insert("min_p".into(), serde_json::json!(v)); + } + if let Some(v) = cli.temperature { + sampling_params.insert("temperature".into(), serde_json::json!(v)); + } + if let Some(v) = cli.frequency_penalty { + sampling_params.insert("frequency_penalty".into(), serde_json::json!(v)); + } + if let Some(v) = cli.presence_penalty { + sampling_params.insert("presence_penalty".into(), serde_json::json!(v)); + } + if let Some(v) = cli.repetition_penalty { + sampling_params.insert("repetition_penalty".into(), serde_json::json!(v)); + } + + if !sampling_params.is_empty() { + if !cli.backend.is_openai_compatible() { + return Err(BenchError::Config( + "Sampling parameters are only supported by openai-compatible backends." + .into(), + )); + } + + // Merge: sampling_params first, then extra_body on top (extra_body wins) + let merged = match extra_body.take() { + Some(serde_json::Value::Object(existing)) => { + sampling_params.extend(existing); + sampling_params + } + Some(other) => { + // extra_body was not an object — just use sampling params + eprintln!( + "Warning: --extra-body is not a JSON object, sampling params may be lost" + ); + let _ = other; + sampling_params + } + None => sampling_params, + }; + extra_body = Some(serde_json::Value::Object(merged)); + } + } + + // Parse metadata + let metadata = match &cli.metadata { + None => None, + Some(items) => { + let mut pairs = Vec::new(); + for item in items { + let (k, v) = item.split_once('=').ok_or_else(|| { + BenchError::Config("Invalid metadata format. Use KEY=VALUE".into()) + })?; + pairs.push((k.trim().to_string(), v.trim().to_string())); + } + Some(pairs) + } + }; + + // Parse goodput SLOs + let goodput = parse_goodput(&cli.goodput)?; + + // Parse ramp-up config + let ramp_up = parse_ramp_up(cli)?; + + // Default percentile metrics based on backend type + let default_percentile_metrics = if cli.backend.is_pooling() { + "e2el" + } else { + "ttft,tpot,itl,e2el" + }; + let percentile_metrics_str = + cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics); + let selected_percentile_metrics: Vec = + percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect(); + + let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?; + let sweep_summary_percentiles = cli + .sweep_summary_percentiles + .as_deref() + .map(|raw| parse_percentiles(raw, true)) + .transpose()? + .unwrap_or_default(); + let mut selected_percentiles = + merge_percentiles(&metric_percentiles, &sweep_summary_percentiles); + // Always include p90 so it appears in console output and sweep summary + if !selected_percentiles.contains(&90.0) { + selected_percentiles.push(90.0); + } + + let tokenizer_id = if cli.skip_tokenizer_init { + None + } else { + Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default()) + }; + + // Resolve input/output lengths + let random_input_len = cli.resolved_random_input_len(); + let random_output_len = cli.resolved_random_output_len(); + let per_turn_input_len = cli.resolved_per_turn_input_len(); + + // Normalized multi-turn turn counts (computed in validation block below, defaults + // to num_turns if multi-turn mode is not active) + let mut multi_turn_min_turns = cli.multi_turn_num_turns; + let mut multi_turn_max_turns = cli.multi_turn_num_turns; + + // For random datasets with openai-compatible backends, default to ignore_eos. + // Exception: multi-turn mode, where ignore_eos causes unbounded context growth + // across turns. Multi-turn uses min_tokens instead for output length control. + // Pooling backends don't generate tokens, so ignore_eos is irrelevant. + let ignore_eos = if cli.backend.is_pooling() { + false + } else { + cli.ignore_eos + || ((cli.dataset_name == DatasetName::Random + || cli.dataset_name == DatasetName::RandomMm) + && cli.backend.is_openai_compatible() + && !cli.multi_turn) + }; + + // Pooling backends don't support multi-turn + if cli.backend.is_pooling() && cli.multi_turn { + return Err(BenchError::Config( + "Pooling/embedding backends do not support --multi-turn".into(), + )); + } + + // LoRA validation. Adapter names must be non-empty after trim; pooling + // backends are out of scope (vLLM LoRA routing is for generative paths). + let lora_modules = match cli.lora_modules.as_ref() { + None => None, + Some(names) => { + if names.is_empty() { + return Err(BenchError::Config( + "--lora-modules requires at least one adapter name".into(), + )); + } + if cli.backend.is_pooling() { + return Err(BenchError::Config( + "--lora-modules is not supported for pooling/embedding backends".into(), + )); + } + let mut out = Vec::with_capacity(names.len()); + for n in names { + let trimmed = n.trim(); + if trimmed.is_empty() { + return Err(BenchError::Config( + "--lora-modules contains an empty adapter name".into(), + )); + } + out.push(Arc::::from(trimmed)); + } + Some(out) + } + }; + + // Random-MM validation and config parsing + let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm { + if cli.backend != BackendKind::OpenaiChat { + return Err(BenchError::Config( + "Multi-modal content (images) is only supported on 'openai-chat' backend." + .into(), + )); + } + let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt( + &cli.random_mm_limit_mm_per_prompt, + )?; + let buckets = + crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?; + (limit, buckets) + } else { + (MmLimitPerPrompt::default(), Vec::new()) + }; + + // Note: --dataset-path is optional for sharegpt (auto-downloads) and + // sonnet (uses built-in Shakespeare's sonnets). + + // Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1)) + let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?; + + // Batched inputs only make sense for pooling backends (the generation + // backends send one prompt per request). + if cli.random_batch_size == 0 { + return Err(BenchError::Config( + "--random-batch-size must be at least 1".into(), + )); + } + if cli.random_batch_size > 1 + && !cli.backend.is_pooling() + && cli.dataset_name != DatasetName::RandomRerank + { + return Err(BenchError::Config( + "--random-batch-size > 1 is only supported with embeddings/pooling backends".into(), + )); + } + + // random-rerank validation (mirrors Python RandomDatasetForReranking) + let is_reranker = !cli.no_reranker; + if cli.dataset_name == DatasetName::RandomRerank { + if !cli.backend.is_pooling() { + return Err(BenchError::Config( + "--dataset-name random-rerank requires an embeddings/pooling backend \ + (e.g. --backend vllm-rerank)" + .into(), + )); + } + if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) { + return Err(BenchError::Config( + "--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \ + (the query is folded into the first batch slot)" + .into(), + )); + } + } + + // Custom dataset validation + if cli.dataset_name == DatasetName::Custom { + match cli.dataset_path.as_deref() { + None => { + return Err(BenchError::Config( + "--dataset-path is required for --dataset-name custom \ + (a JSONL file with {\"prompt\": ..., \"output_tokens\": ...} lines)" + .into(), + )); + } + Some(p) if !p.ends_with(".jsonl") => { + return Err(BenchError::Config( + "Only JSONL format is supported for the custom dataset".into(), + )); + } + _ => {} + } + if !cli.skip_chat_template { + eprintln!( + "NOTE: client-side chat template rendering is not supported; custom \ + dataset prompts are sent raw (equivalent to --skip-chat-template)." + ); + } + } + + // Prefix repetition validation + if cli.dataset_name == DatasetName::PrefixRepetition { + if cli.prefix_repetition_num_prefixes == 0 { + return Err(BenchError::Config( + "--prefix-repetition-num-prefixes must be at least 1".into(), + )); + } + if cli.num_prompts < cli.prefix_repetition_num_prefixes { + return Err(BenchError::Config(format!( + "--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})", + cli.num_prompts, cli.prefix_repetition_num_prefixes + ))); + } + } + + // HF dataset validation + if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() { + return Err(BenchError::Config( + "--dataset-path is required for --dataset-name hf \ + (set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')" + .into(), + )); + } + if let Some(len) = cli.hf_output_len + && len == 0 + { + return Err(BenchError::Config( + "--hf-output-len must be at least 1".into(), + )); + } + + // Multi-turn validation + if cli.multi_turn { + if cli.backend != BackendKind::OpenaiChat { + return Err(BenchError::Config( + "--multi-turn requires --backend openai-chat".into(), + )); + } + if cli.multi_turn_num_turns == 0 { + return Err(BenchError::Config( + "--multi-turn-num-turns must be at least 1".into(), + )); + } + + // Normalize and validate min/max turns. ShareGPT only consumes max_turns + // (the loader walks all available turns up to the cap), so the + // min/num/max coupling used for synthetic generation does not apply. + if cli.dataset_name == DatasetName::ShareGpt { + if cli.multi_turn_max_turns == 1 { + return Err(BenchError::Config( + "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(), + )); + } + } else { + (multi_turn_min_turns, multi_turn_max_turns) = + match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) { + (0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns), + (m, 0) => (m, cli.multi_turn_num_turns), + (0, x) => (cli.multi_turn_num_turns, x), + (m, x) => (m, x), + }; + if multi_turn_min_turns < 1 { + return Err(BenchError::Config( + "--multi-turn-min-turns must be at least 1".into(), + )); + } + if multi_turn_min_turns > multi_turn_max_turns { + return Err(BenchError::Config( + "--multi-turn-min-turns must be <= --multi-turn-max-turns".into(), + )); + } + } + + if ignore_eos { + eprintln!( + "WARNING: --ignore-eos is set with --multi-turn. The server may not \ + respect output length limits, causing unbounded context growth." + ); + } + + // Validate prefix sharing ratios + let pg = cli.multi_turn_prefix_global_ratio; + let pc = cli.multi_turn_prefix_conversation_ratio; + if !(0.0..=1.0).contains(&pg) { + return Err(BenchError::Config( + "--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(), + )); + } + if !(0.0..=1.0).contains(&pc) { + return Err(BenchError::Config( + "--multi-turn-prefix-conversation-ratio must be in [0.0, 1.0]".into(), + )); + } + if pg + pc >= 1.0 { + return Err(BenchError::Config( + "--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(), + )); + } + if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random { + return Err(BenchError::Config( + "Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(), + )); + } + } + + if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) { + return Err(BenchError::Config(format!( + "--steady-state-threshold must be in (0.0, 1.0], got {}", + cli.steady_state_threshold + ))); + } + if let Some(mw) = cli.steady_state_min_window + && mw < 0.0 + { + return Err(BenchError::Config(format!( + "--steady-state-min-window must be >= 0, got {mw}" + ))); + } + + if cli.profile_batch_threshold.is_some() && !cli.profile { + return Err(BenchError::Config( + "--profile-batch-threshold requires --profile".into(), + )); + } + if cli.profile_duration <= 0.0 { + return Err(BenchError::Config( + "--profile-duration must be positive".into(), + )); + } + if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 { + return Err(BenchError::Config( + "--profile-duration requires --profile-batch-threshold".into(), + )); + } + + Ok(BenchConfig { + backend: cli.backend, + base_url, + api_url, + model: cli.model.clone(), + model_name: cli.served_model_name.clone(), + tokenizer_id, + tokenizer_mode: cli.tokenizer_mode.clone(), + trust_remote_code: cli.trust_remote_code, + skip_tokenizer_init: cli.skip_tokenizer_init, + dataset_name: cli.dataset_name, + dataset_path: cli.dataset_path.clone(), + max_model_len: cli.max_model_len, + random_input_len, + random_output_len, + random_prefix_len: cli.random_prefix_len, + random_range_ratio, + random_batch_size: cli.random_batch_size, + is_reranker, + custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len), + prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len, + prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len, + prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes, + prefix_repetition_output_len: cli + .output_len + .unwrap_or(cli.prefix_repetition_output_len), + random_cache_hit_fraction: cli.random_cache_hit_fraction, + random_cache_ratio: cli.random_cache_ratio, + sharegpt_output_len: cli.sharegpt_output_len, + sonnet_input_len: cli.sonnet_input_len, + sonnet_output_len: cli.sonnet_output_len, + sonnet_prefix_len: cli.sonnet_prefix_len, + no_oversample: cli.no_oversample, + disable_shuffle: cli.disable_shuffle, + num_prompts: cli.num_prompts, + request_rate: cli.request_rate, + burstiness: cli.burstiness, + max_concurrency: cli.max_concurrency, + steady_state_threshold: cli.steady_state_threshold, + steady_state_min_window: cli.steady_state_min_window, + no_steady_state: cli.no_steady_state, + disable_tqdm: cli.disable_tqdm, + num_warmups: cli.num_warmups, + profile: cli.profile, + profile_batch_threshold: cli.profile_batch_threshold, + profile_duration: cli.profile_duration, + save_result: cli.save_result, + save_detailed: cli.save_detailed, + append_result: cli.append_result, + result_dir: cli.result_dir.clone(), + result_filename: cli.result_filename.clone(), + seed: cli.seed, + ignore_eos, + insecure: cli.insecure, + selected_percentile_metrics, + selected_percentiles, + sweep_summary_percentiles, + label: cli.label.clone(), + logprobs: cli.logprobs, + request_id_prefix: cli.get_request_id_prefix(), + ready_check_timeout_sec: cli.ready_check_timeout_sec, + extra_headers, + extra_body, + metadata, + dry_run: cli.dry_run, + goodput, + ramp_up, + multi_turn: cli.multi_turn, + multi_turn_num_turns: cli.multi_turn_num_turns, + multi_turn_min_turns, + multi_turn_max_turns, + sharegpt_multi_turn_max_turns: if cli.multi_turn + && cli.dataset_name == DatasetName::ShareGpt + && cli.multi_turn_max_turns != 0 + { + Some(cli.multi_turn_max_turns) + } else { + None + }, + per_turn_input_len, + multi_turn_concurrency: cli.multi_turn_concurrency, + multi_turn_delay_ms: cli.multi_turn_delay_ms, + multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio, + multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio, + speed_bench_config: cli.speed_bench_config, + speed_bench_category: cli.speed_bench_category.clone(), + speed_bench_max_input_len: cli.speed_bench_max_input_len, + hf_split: cli.hf_split.clone(), + hf_subset: cli.hf_subset.clone(), + hf_output_len: cli.hf_output_len, + hf_text_column: cli.hf_text_column.clone(), + reset_prefix_cache: cli.reset_prefix_cache, + prompt_token_ids: cli.prompt_token_ids, + random_mm_base_items_per_request: cli.random_mm_base_items_per_request, + random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio, + random_mm_limit, + random_mm_buckets, + enable_multimodal_chat: cli.enable_multimodal_chat, + lora_modules, + lora_assignment: cli.lora_assignment, + }) + } +} + +fn parse_percentiles(raw: &str, dedupe: bool) -> Result> { + let mut percentiles = Vec::new(); + for s in raw.split(',') { + let p = s + .trim() + .parse::() + .map_err(|_| BenchError::Config(format!("Invalid percentile: {s}")))?; + if !(0.0..=100.0).contains(&p) { + return Err(BenchError::Config(format!( + "Percentile must be in [0, 100], got: {p}" + ))); + } + if !dedupe || !percentiles.contains(&p) { + percentiles.push(p); + } + } + Ok(percentiles) +} + +fn merge_percentiles(metric_percentiles: &[f64], summary_percentiles: &[f64]) -> Vec { + let mut merged = Vec::with_capacity(metric_percentiles.len() + summary_percentiles.len()); + for &p in metric_percentiles { + if !merged.contains(&p) { + merged.push(p); + } + } + for &p in summary_percentiles { + if !merged.contains(&p) { + merged.push(p); + } + } + merged +} + +fn parse_goodput(goodput_args: &Option>) -> Result { + let items = match goodput_args { + None => return Ok(GoodputConfig::default()), + Some(items) => items, + }; + + let valid_names = ["ttft", "tpot", "e2el"]; + let mut config = GoodputConfig::default(); + + for item in items { + let (name, val_str) = item.split_once(':').ok_or_else(|| { + BenchError::Config(format!( + "Invalid goodput format: '{item}'. Use KEY:VALUE (e.g. ttft:100)" + )) + })?; + + let val: f64 = val_str + .trim() + .parse() + .map_err(|_| BenchError::Config(format!("Invalid goodput value: '{val_str}'")))?; + + if val < 0.0 { + return Err(BenchError::Config(format!( + "Goodput SLO value must be non-negative, got: {name}:{val}" + ))); + } + + if !valid_names.contains(&name) { + return Err(BenchError::Config(format!( + "Invalid goodput metric: '{name}'. Valid: {valid_names:?}" + ))); + } + + match name { + "ttft" => config.ttft_ms = Some(val), + "tpot" => config.tpot_ms = Some(val), + "e2el" => config.e2el_ms = Some(val), + _ => unreachable!(), + } + } + + Ok(config) +} + +fn parse_ramp_up(cli: &Cli) -> Result> { + let strategy = match cli.ramp_up_strategy { + None => return Ok(None), + Some(s) => s, + }; + + let start_rps = cli.ramp_up_start_rps.ok_or_else(|| { + BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into()) + })?; + + let end_rps = cli.ramp_up_end_rps.ok_or_else(|| { + BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into()) + })?; + + if start_rps <= 0.0 || end_rps <= 0.0 { + return Err(BenchError::Config( + "Ramp-up RPS values must be positive".into(), + )); + } + + Ok(Some(RampUpConfig { + strategy, + start_rps, + end_rps, + })) +} + +#[cfg(test)] +mod tests { + use clap::Parser; + + use super::*; + use crate::cli::Cli; + + fn base_multi_turn_args() -> Vec<&'static str> { + vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + ] + } + + #[test] + fn test_prefix_sharing_defaults_to_zero() { + let args = base_multi_turn_args(); + let cli = Cli::parse_from(args); + let config = BenchConfig::from_cli(&cli).unwrap(); + assert_eq!(config.multi_turn_prefix_global_ratio, 0.0); + assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0); + } + + #[test] + fn test_prefix_sharing_valid() { + let mut args = base_multi_turn_args(); + args.extend([ + "--multi-turn-prefix-global-ratio", + "0.1", + "--multi-turn-prefix-conversation-ratio", + "0.8", + ]); + let cli = Cli::parse_from(args); + let config = BenchConfig::from_cli(&cli).unwrap(); + assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10); + assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10); + } + + #[test] + fn test_prefix_sharing_sum_exceeds_one_fails() { + let mut args = base_multi_turn_args(); + args.extend([ + "--multi-turn-prefix-global-ratio", + "0.6", + "--multi-turn-prefix-conversation-ratio", + "0.6", + ]); + let cli = Cli::parse_from(args); + assert!(BenchConfig::from_cli(&cli).is_err()); + } + + #[test] + fn test_prefix_sharing_sum_equals_one_fails() { + let mut args = base_multi_turn_args(); + args.extend([ + "--multi-turn-prefix-global-ratio", + "0.5", + "--multi-turn-prefix-conversation-ratio", + "0.5", + ]); + let cli = Cli::parse_from(args); + assert!(BenchConfig::from_cli(&cli).is_err()); + } + + #[test] + fn test_prefix_sharing_out_of_range_fails() { + let mut args = base_multi_turn_args(); + args.extend(["--multi-turn-prefix-global-ratio", "1.5"]); + let cli = Cli::parse_from(args); + assert!(BenchConfig::from_cli(&cli).is_err()); + } + + #[test] + fn test_prefix_sharing_requires_random_dataset() { + let args = vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + "--dataset-name", + "sharegpt", + "--multi-turn-prefix-global-ratio", + "0.1", + ]; + let cli = Cli::parse_from(args); + assert!(BenchConfig::from_cli(&cli).is_err()); + } + + #[test] + fn test_sharegpt_multi_turn_max_turns_default_uncapped() { + let args = vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + "--dataset-name", + "sharegpt", + ]; + let cli = Cli::parse_from(args); + let config = BenchConfig::from_cli(&cli).unwrap(); + + assert_eq!(config.multi_turn_max_turns, 3); + assert_eq!(config.sharegpt_multi_turn_max_turns, None); + } + + #[test] + fn test_sharegpt_multi_turn_max_turns_2_succeeds() { + // Regression: previously the (0, x) normalization arm produced + // (min=multi_turn_num_turns=3, max=2), tripping the min>max check + // with a misleading error about a flag the user never set. + let args = vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + "--dataset-name", + "sharegpt", + "--multi-turn-max-turns", + "2", + ]; + let cli = Cli::parse_from(args); + let config = BenchConfig::from_cli(&cli).unwrap(); + assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2)); + } + + #[test] + fn test_sharegpt_multi_turn_max_turns_1_rejected() { + let args = vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + "--dataset-name", + "sharegpt", + "--multi-turn-max-turns", + "1", + ]; + let cli = Cli::parse_from(args); + let err = BenchConfig::from_cli(&cli).unwrap_err().to_string(); + assert!( + err.contains("at least 2 for ShareGPT"), + "expected ShareGPT-specific error, got: {err}" + ); + } + + #[test] + fn test_sharegpt_multi_turn_max_turns_explicit_cap() { + let args = vec![ + "vllm-bench", + "--backend", + "openai-chat", + "--multi-turn", + "--model", + "test-model", + "--dataset-name", + "sharegpt", + "--multi-turn-max-turns", + "20", + ]; + let cli = Cli::parse_from(args); + let config = BenchConfig::from_cli(&cli).unwrap(); + + assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20)); + } + + #[test] + fn test_sweep_summary_percentiles_default_empty() { + let args = base_multi_turn_args(); + let cli = Cli::parse_from(args); + let config = BenchConfig::from_cli(&cli).unwrap(); + + assert!(config.sweep_summary_percentiles.is_empty()); + assert_eq!(config.selected_percentiles, vec![99.0, 90.0]); + } + + #[test] + fn test_sweep_summary_percentiles_are_deduped_and_merged() { + let mut args = base_multi_turn_args(); + args.extend([ + "--metric-percentiles", + "99,95", + "--sweep-summary-percentiles", + "90,95,90", + ]); + let cli = Cli::parse_from(args); + let config = BenchConfig::from_cli(&cli).unwrap(); + + assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]); + assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]); + } + + #[test] + fn test_invalid_sweep_summary_percentile_fails() { + let mut args = base_multi_turn_args(); + args.extend(["--sweep-summary-percentiles", "101"]); + let cli = Cli::parse_from(args); + assert!(BenchConfig::from_cli(&cli).is_err()); + } + + #[test] + fn test_max_model_len_is_configured() { + let args = vec![ + "vllm-bench", + "--model", + "test-model", + "--max-model-len", + "4096", + ]; + let cli = Cli::parse_from(args); + let config = BenchConfig::from_cli(&cli).unwrap(); + + assert_eq!(config.max_model_len, Some(4096)); + } + + #[test] + fn test_zero_max_model_len_fails() { + let args = vec![ + "vllm-bench", + "--model", + "test-model", + "--max-model-len", + "0", + ]; + let cli = Cli::parse_from(args); + + assert!(BenchConfig::from_cli(&cli).is_err()); + } + #[test] + fn test_range_ratio_parse_float() { + let rr = RangeRatio::parse("0.2").unwrap(); + assert_eq!(rr.input, 0.2); + assert_eq!(rr.output, 0.2); + assert!(!rr.is_fixed()); + assert!(RangeRatio::parse("0.0").unwrap().is_fixed()); + } + + #[test] + fn test_range_ratio_parse_dict() { + let rr = RangeRatio::parse(r#"{"input": 0.1, "output": 0.5}"#).unwrap(); + assert_eq!(rr.input, 0.1); + assert_eq!(rr.output, 0.5); + } + + #[test] + fn test_range_ratio_rejects_old_default_with_hint() { + let err = RangeRatio::parse("1.0").unwrap_err().to_string(); + assert!(err.contains("semantics now match Python"), "got: {err}"); + assert!(RangeRatio::parse("-0.1").is_err()); + assert!(RangeRatio::parse("{\"input\": 0.1}").is_err()); + assert!(RangeRatio::parse("abc").is_err()); + } + + #[test] + fn test_range_ratio_bounds_python_semantics() { + // Python: [floor(len*(1-r)), ceil(len*(1+r))], output low clamped to 1 + let rr = RangeRatio::parse("0.2").unwrap(); + assert_eq!(rr.input_bounds(100), (80, 120)); + assert_eq!(rr.output_bounds(100), (80, 120)); + let rr = RangeRatio::parse("0.99").unwrap(); + assert_eq!(rr.output_bounds(1).0, 1); // clamped + let fixed = RangeRatio::parse("0.0").unwrap(); + assert_eq!(fixed.input_bounds(8192), (8192, 8192)); + } +} diff --git a/rust/src/bench/src/datasets/custom.rs b/rust/src/bench/src/datasets/custom.rs new file mode 100644 index 000000000000..04bbfcfdbac2 --- /dev/null +++ b/rust/src/bench/src/datasets/custom.rs @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Custom dataset: JSONL file with one request per line. +//! +//! ```jsonl +//! {"prompt": "What is the capital of India?", "output_tokens": 10} +//! {"prompt": "What is the capital of Iran?", "output_tokens": 1520} +//! ``` +//! +//! Mirrors Python's `CustomDataset`. `output_tokens` is optional unless +//! `--custom-output-len -1` is passed. Unlike Python, prompts are always sent +//! raw (no client-side chat template; see `--skip-chat-template`). + +use std::sync::Arc; + +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use serde::Deserialize; + +use super::SampleRequest; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +#[derive(Deserialize)] +struct CustomLine { + prompt: String, + output_tokens: Option, +} + +/// Load the custom JSONL dataset. +/// +/// `output_len < 0` means "use the per-line output_tokens field" (Python's +/// `--custom-output-len -1`); otherwise `output_len` applies to every request. +pub fn load_custom_dataset( + tokenizer: &TokenizerKind, + path: &str, + num_requests: usize, + output_len: i64, + seed: u64, + request_id_prefix: &str, + no_oversample: bool, + disable_shuffle: bool, +) -> Result> { + let content = std::fs::read_to_string(path) + .map_err(|e| BenchError::Config(format!("Failed to read custom dataset '{path}': {e}")))?; + + let mut lines: Vec = Vec::new(); + for (lineno, line) in content.lines().enumerate() { + let line = line.trim(); + if line.is_empty() { + continue; + } + let parsed: CustomLine = serde_json::from_str(line).map_err(|e| { + BenchError::Config(format!( + "Invalid JSONL at {path}:{}: {e} (each line must be an object \ + with a 'prompt' field)", + lineno + 1 + )) + })?; + lines.push(parsed); + } + if lines.is_empty() { + return Err(BenchError::Config(format!( + "Custom dataset '{path}' contains no entries" + ))); + } + + // Python shuffles the loaded data (seeded) before taking num_requests. + if !disable_shuffle { + let mut rng = StdRng::seed_from_u64(seed); + lines.shuffle(&mut rng); + } + + let mut requests: Vec = Vec::with_capacity(num_requests.min(lines.len())); + for (i, item) in lines.iter().enumerate() { + if requests.len() >= num_requests { + break; + } + + let expected_output_len = if output_len < 0 { + let raw = item.output_tokens.as_ref().ok_or_else(|| { + BenchError::Config( + "custom dataset: --custom-output-len -1 requires an \ + 'output_tokens' field on every line" + .into(), + ) + })?; + raw.as_i64().filter(|v| *v > 0).ok_or_else(|| { + BenchError::Config(format!( + "custom dataset: invalid 'output_tokens' value {raw}: \ + must be a positive integer" + )) + })? as usize + } else { + output_len as usize + }; + + let prompt_len = tokenizer.encode(&item.prompt, true)?.len(); + requests.push(SampleRequest { + prompt: Arc::from(item.prompt.as_str()), + prompt_len, + expected_output_len, + request_id: Some(format!("{request_id_prefix}{i}")), + ..Default::default() + }); + } + + super::oversample_requests( + &mut requests, + num_requests, + request_id_prefix, + no_oversample, + ); + Ok(requests) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn write_temp_jsonl(name: &str, content: &str) -> String { + let path = std::env::temp_dir().join(format!("vllm-bench-custom-{name}.jsonl")); + std::fs::write(&path, content).unwrap(); + path.to_string_lossy().into_owned() + } + + /// gpt2 via built-in tiktoken encoding — loads without network access. + fn test_tokenizer() -> TokenizerKind { + crate::tokenizer::load_tokenizer("gpt2", false, None) + .expect("gpt2 built-in tiktoken should always load without network") + } + + #[test] + fn test_load_custom_dataset_basic() { + let path = write_temp_jsonl( + "basic", + r#"{"prompt": "hello world", "output_tokens": 10} +{"prompt": "foo bar baz", "output_tokens": 20} +"#, + ); + let reqs = load_custom_dataset(&test_tokenizer(), &path, 2, 256, 0, "t-", true, true) + .expect("load should succeed"); + assert_eq!(reqs.len(), 2); + // Fixed output_len (256) wins over per-line output_tokens by default + assert!(reqs.iter().all(|r| r.expected_output_len == 256)); + assert_eq!(&*reqs[0].prompt, "hello world"); + assert!(reqs[0].prompt_len > 0); + } + + #[test] + fn test_load_custom_dataset_per_line_output_tokens() { + let path = write_temp_jsonl( + "perline", + r#"{"prompt": "hello", "output_tokens": 10} +{"prompt": "world", "output_tokens": 20} +"#, + ); + let reqs = load_custom_dataset(&test_tokenizer(), &path, 2, -1, 0, "t-", true, true) + .expect("load should succeed"); + assert_eq!(reqs[0].expected_output_len, 10); + assert_eq!(reqs[1].expected_output_len, 20); + } + + #[test] + fn test_load_custom_dataset_missing_output_tokens_errors() { + let path = write_temp_jsonl("missing", r#"{"prompt": "hello"}"#); + let err = load_custom_dataset(&test_tokenizer(), &path, 1, -1, 0, "t-", true, true) + .expect_err("should fail without output_tokens"); + assert!(err.to_string().contains("output_tokens")); + } + + #[test] + fn test_load_custom_dataset_missing_prompt_errors() { + let path = write_temp_jsonl("noprompt", r#"{"text": "hello"}"#); + assert!( + load_custom_dataset(&test_tokenizer(), &path, 1, 256, 0, "t-", true, true).is_err() + ); + } +} diff --git a/rust/src/bench/src/datasets/hf_dataset.rs b/rust/src/bench/src/datasets/hf_dataset.rs new file mode 100644 index 000000000000..35932a39833f --- /dev/null +++ b/rust/src/bench/src/datasets/hf_dataset.rs @@ -0,0 +1,1239 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; + +use super::SampleRequest; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +/// Cache directory for downloaded HF datasets. +fn cache_dir() -> std::path::PathBuf { + dirs::cache_dir() + .unwrap_or_else(|| std::path::PathBuf::from("/tmp")) + .join("vllm-bench") + .join("datasets") +} + +/// Sanitize a dataset name for use in filenames (replace `/` and other unsafe chars). +fn sanitize_name(name: &str) -> String { + name.chars() + .map(|c| { + if c.is_alphanumeric() || c == '-' || c == '_' || c == '.' { + c + } else { + '_' + } + }) + .collect() +} + +/// Detected column format for extracting prompts from HF dataset rows. +enum ColumnFormat { + /// Column contains chat messages (array of {role, content} or {from, value}). + Chat(String), + /// Single text column for prompt, optional output column. + Text { + prompt_col: String, + output_col: Option, + }, + /// Multiple columns combined (e.g., context + input for LongBench). + Combined { + cols: Vec, + output_col: Option, + }, +} + +/// Make a GET request with retry logic (3 retries with exponential backoff). +/// Returns the parsed JSON response. +fn get_with_retry( + client: &reqwest::blocking::Client, + url: &str, + label: &str, +) -> Result { + let max_retries = 3; + for attempt in 0..=max_retries { + let resp = match client.get(url).send() { + Ok(r) => r, + Err(e) => { + if attempt < max_retries { + std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + continue; + } + return Err(BenchError::Config(format!( + "{label} failed after {max_retries} retries: {e}" + ))); + } + }; + + // Detect gated/private dataset errors + let status = resp.status(); + if status.as_u16() == 401 || status.as_u16() == 403 { + return Err(BenchError::Config(format!( + "{label} returned HTTP {status}. This dataset may be gated or private. \ + Try setting the HF_TOKEN environment variable with a valid HuggingFace token." + ))); + } + + if status.is_server_error() && attempt < max_retries { + std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + continue; + } + + if !status.is_success() { + return Err(BenchError::Config(format!( + "{label} returned HTTP {status}" + ))); + } + + let data: serde_json::Value = resp + .json() + .map_err(|e| BenchError::Config(format!("Failed to parse {label} response: {e}")))?; + return Ok(data); + } + unreachable!() +} + +/// Download an arbitrary HuggingFace dataset via the datasets-server REST API. +/// +/// Returns `(cache_path, resolved_config, resolved_split)`. +/// +/// If both `subset` and `split` are provided, the `/info` call is skipped as an optimization. +/// Paginated download fetches rows in pages of 100 until `num_rows_needed` are collected +/// or the dataset is exhausted. +pub fn download_hf_dataset( + dataset: &str, + subset: Option<&str>, + split: Option<&str>, + num_rows_needed: usize, +) -> Result<(String, String, String)> { + let encoded_dataset: String = + url::form_urlencoded::byte_serialize(dataset.as_bytes()).collect(); + + let mut client_builder = + reqwest::blocking::Client::builder().timeout(std::time::Duration::from_secs(120)); + + // Add HF_TOKEN auth header if available + if let Ok(token) = std::env::var("HF_TOKEN") { + let mut headers = reqwest::header::HeaderMap::new(); + let header_value = reqwest::header::HeaderValue::from_str(&format!("Bearer {token}")) + .map_err(|e| BenchError::Config(format!("Invalid HF_TOKEN: {e}")))?; + headers.insert(reqwest::header::AUTHORIZATION, header_value); + client_builder = client_builder.default_headers(headers); + } + + let client = client_builder + .build() + .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?; + + // Resolve config and split + let (resolved_config, resolved_split) = if let (Some(cfg), Some(spl)) = (subset, split) { + // Both provided — skip /info call + (cfg.to_string(), spl.to_string()) + } else { + // Call /info to discover available configs and splits + let info_url = + format!("https://datasets-server.huggingface.co/info?dataset={encoded_dataset}"); + let info = get_with_retry(&client, &info_url, "HF dataset /info")?; + + let dataset_info = + info.get("dataset_info").and_then(|d| d.as_object()).ok_or_else(|| { + BenchError::Config(format!( + "No 'dataset_info' in /info response for '{dataset}'. \ + The dataset may not exist or may not be accessible." + )) + })?; + + // Resolve config + let resolved_config = if let Some(user_cfg) = subset { + if !dataset_info.contains_key(user_cfg) { + let available: Vec<&String> = dataset_info.keys().collect(); + return Err(BenchError::Config(format!( + "Config '{user_cfg}' not found in dataset '{dataset}'. Available: {available:?}" + ))); + } + user_cfg.to_string() + } else if dataset_info.contains_key("default") { + "default".to_string() + } else { + dataset_info.keys().next().cloned().ok_or_else(|| { + BenchError::Config(format!("No configs found in dataset '{dataset}'")) + })? + }; + + // Resolve split from the config's splits + let splits_info = dataset_info + .get(&resolved_config) + .and_then(|c| c.get("splits")) + .and_then(|s| s.as_object()); + + let resolved_split = if let Some(user_spl) = split { + if let Some(si) = splits_info + && !si.contains_key(user_spl) + { + let available: Vec<&String> = si.keys().collect(); + return Err(BenchError::Config(format!( + "Split '{user_spl}' not found in config '{resolved_config}' \ + of dataset '{dataset}'. Available: {available:?}" + ))); + } + user_spl.to_string() + } else if let Some(si) = splits_info { + // Priority: train > test > validation > first + let priority = ["train", "test", "validation"]; + let mut found = None; + for p in &priority { + if si.contains_key(*p) { + found = Some(p.to_string()); + break; + } + } + found + .unwrap_or_else(|| si.keys().next().cloned().unwrap_or_else(|| "train".to_string())) + } else { + "train".to_string() + }; + + (resolved_config, resolved_split) + }; + + println!("HF dataset: {dataset} (config={resolved_config}, split={resolved_split})"); + + // Check cache + let dir = cache_dir(); + std::fs::create_dir_all(&dir)?; + let cache_path = dir.join(format!( + "hf-{}-{}-{}.json", + sanitize_name(dataset), + sanitize_name(&resolved_config), + sanitize_name(&resolved_split) + )); + + if cache_path.exists() { + let path_str = cache_path.to_string_lossy().to_string(); + println!("HF dataset cached: {path_str}"); + return Ok((path_str, resolved_config, resolved_split)); + } + + println!("Downloading HF dataset '{dataset}' from datasets-server..."); + + let encoded_config: String = + url::form_urlencoded::byte_serialize(resolved_config.as_bytes()).collect(); + let encoded_split: String = + url::form_urlencoded::byte_serialize(resolved_split.as_bytes()).collect(); + + let mut all_rows: Vec = Vec::new(); + let mut offset = 0usize; + let page_size = 100usize; + + loop { + let url = format!( + "https://datasets-server.huggingface.co/rows\ + ?dataset={encoded_dataset}\ + &config={encoded_config}\ + &split={encoded_split}\ + &offset={offset}\ + &length={page_size}" + ); + + let data = get_with_retry(&client, &url, "HF dataset /rows")?; + + let rows = data["rows"] + .as_array() + .ok_or_else(|| BenchError::Config("No 'rows' in API response".into()))?; + + if rows.is_empty() { + break; + } + + for row in rows { + if let Some(row_data) = row.get("row") { + all_rows.push(row_data.clone()); + } + } + + let fetched = rows.len(); + offset += fetched; + + let total = data["num_rows_total"].as_u64().unwrap_or(0); + eprint!("\r Fetched {offset}/{total} rows..."); + + // Stop if we have enough rows or reached end of dataset + if all_rows.len() >= num_rows_needed || fetched < page_size { + break; + } + } + eprintln!(); // newline after progress + + if all_rows.is_empty() { + return Err(BenchError::Config(format!( + "HF dataset '{dataset}' download returned no rows" + ))); + } + + // Save to cache + let json_str = serde_json::to_string(&all_rows)?; + std::fs::write(&cache_path, &json_str)?; + + let path_str = cache_path.to_string_lossy().to_string(); + println!("HF dataset: {} rows saved to {path_str}", all_rows.len()); + Ok((path_str, resolved_config, resolved_split)) +} + +/// Check if a JSON value looks like a chat message (has role+content or from+value). +fn is_chat_message(val: &serde_json::Value) -> bool { + if let Some(obj) = val.as_object() { + (obj.contains_key("role") && obj.contains_key("content")) + || (obj.contains_key("from") && obj.contains_key("value")) + } else { + false + } +} + +/// Check if a JSON value is an array of chat messages. +fn is_chat_array(val: &serde_json::Value) -> bool { + val.as_array() + .map(|arr| !arr.is_empty() && arr.iter().all(is_chat_message)) + .unwrap_or(false) +} + +/// Known column names for chat-format data. +const CHAT_COLUMNS: &[&str] = &["conversation", "conversations", "messages"]; + +/// Known column names for single text prompts (in priority order). +const TEXT_COLUMNS: &[&str] = &[ + "prompt", + "question", + "problem", + "input", + "text", + "content", + "instruction", +]; + +/// Known column names for output/completion data. +const OUTPUT_COLUMNS: &[&str] = &[ + "completion", + "response", + "answer", + "output", + "solution", + "answers", +]; + +/// Detect the column format from the first row of the dataset. +fn detect_column_format( + row: &serde_json::Value, + text_column_override: Option<&str>, +) -> Result { + let obj = row + .as_object() + .ok_or_else(|| BenchError::Config("HF dataset row is not a JSON object".into()))?; + + // Helper: find first matching output column + let find_output_col = || -> Option { + for col in OUTPUT_COLUMNS { + if obj.contains_key(*col) { + return Some(col.to_string()); + } + } + None + }; + + // 1. User override via --hf-text-column + if let Some(col_name) = text_column_override { + if !obj.contains_key(col_name) { + let available: Vec<&String> = obj.keys().collect(); + return Err(BenchError::Config(format!( + "Column '{col_name}' not found in dataset. Available columns: {available:?}" + ))); + } + let val = &obj[col_name]; + if is_chat_array(val) { + return Ok(ColumnFormat::Chat(col_name.to_string())); + } + return Ok(ColumnFormat::Text { + prompt_col: col_name.to_string(), + output_col: find_output_col(), + }); + } + + // 2. Chat columns + for col in CHAT_COLUMNS { + if let Some(val) = obj.get(*col) + && is_chat_array(val) + { + return Ok(ColumnFormat::Chat(col.to_string())); + } + } + + // 3. "turns" column — array of strings + if let Some(val) = obj.get("turns") + && let Some(arr) = val.as_array() + && !arr.is_empty() + && arr[0].is_string() + { + return Ok(ColumnFormat::Text { + prompt_col: "turns".to_string(), + output_col: find_output_col(), + }); + } + + // 4. Combined: context + input + if obj.contains_key("context") && obj.contains_key("input") { + return Ok(ColumnFormat::Combined { + cols: vec!["context".to_string(), "input".to_string()], + output_col: find_output_col(), + }); + } + + // 5. Single text columns + for col in TEXT_COLUMNS { + if obj.contains_key(*col) { + return Ok(ColumnFormat::Text { + prompt_col: col.to_string(), + output_col: find_output_col(), + }); + } + } + + // Fallback: list available columns + let available: Vec<&String> = obj.keys().collect(); + Err(BenchError::Config(format!( + "Could not auto-detect prompt column in HF dataset. \ + Available columns: {available:?}. \ + Use --hf-text-column to specify the column containing prompts." + ))) +} + +/// Extract the first user message from a chat message array. +/// Supports both {role, content} and {from, value} formats. +fn extract_chat_prompt(messages: &[serde_json::Value]) -> Option { + for msg in messages { + let role = msg + .get("role") + .and_then(|r| r.as_str()) + .or_else(|| msg.get("from").and_then(|f| f.as_str())); + let content = msg + .get("content") + .and_then(|c| c.as_str()) + .or_else(|| msg.get("value").and_then(|v| v.as_str())); + + if let (Some(role), Some(content)) = (role, content) + && (role == "user" || role == "human") + { + return Some(content.to_string()); + } + } + None +} + +/// Extract the first assistant message from a chat message array. +fn extract_chat_completion(messages: &[serde_json::Value]) -> Option { + for msg in messages { + let role = msg + .get("role") + .and_then(|r| r.as_str()) + .or_else(|| msg.get("from").and_then(|f| f.as_str())); + let content = msg + .get("content") + .and_then(|c| c.as_str()) + .or_else(|| msg.get("value").and_then(|v| v.as_str())); + + if let (Some(role), Some(content)) = (role, content) + && (role == "assistant" || role == "gpt") + { + return Some(content.to_string()); + } + } + None +} + +/// Load an HF dataset from the cached JSON file and convert to SampleRequests. +pub fn load_hf_dataset( + tokenizer: &TokenizerKind, + dataset_path: &str, + num_requests: usize, + hf_output_len: Option, + seed: u64, + request_id_prefix: &str, + text_column_override: Option<&str>, + no_oversample: bool, + disable_shuffle: bool, +) -> Result> { + let content = std::fs::read_to_string(dataset_path).map_err(|e| { + BenchError::Config(format!( + "Failed to read HF dataset file '{dataset_path}': {e}" + )) + })?; + + let entries: Vec = serde_json::from_str(&content) + .map_err(|e| BenchError::Config(format!("Invalid JSON in HF dataset file: {e}")))?; + + if entries.is_empty() { + return Err(BenchError::Config( + "HF dataset file contains no rows".into(), + )); + } + + // Detect column format from first row + let format = detect_column_format(&entries[0], text_column_override)?; + + // Print detected format + match &format { + ColumnFormat::Chat(col) => println!("HF dataset: detected chat column '{col}'"), + ColumnFormat::Text { + prompt_col, + output_col, + } => { + let out_msg = output_col.as_deref().unwrap_or("none"); + println!("HF dataset: detected text column '{prompt_col}', output column: {out_msg}"); + } + ColumnFormat::Combined { cols, output_col } => { + let out_msg = output_col.as_deref().unwrap_or("none"); + println!( + "HF dataset: detected combined columns {:?}, output column: {out_msg}", + cols + ); + } + } + + // Build shuffled indices + let mut indices: Vec = (0..entries.len()).collect(); + let mut rng = StdRng::seed_from_u64(seed); + if !disable_shuffle { + indices.shuffle(&mut rng); + } + + let mut samples = Vec::new(); + let mut idx = 0; + let mut warned_no_output = false; + + for &entry_idx in &indices { + if samples.len() >= num_requests { + break; + } + + let row = &entries[entry_idx]; + + // Extract prompt and optional completion based on format + let (prompt, completion) = match &format { + ColumnFormat::Chat(col) => { + let messages = + row.get(col.as_str()).and_then(|v| v.as_array()).cloned().unwrap_or_default(); + let prompt = match extract_chat_prompt(&messages) { + Some(p) => p, + None => continue, + }; + let completion = extract_chat_completion(&messages); + (prompt, completion) + } + ColumnFormat::Text { + prompt_col, + output_col, + } => { + let prompt_val = match row.get(prompt_col.as_str()) { + Some(v) => v, + None => continue, + }; + + // Handle "turns" column (array of strings — take first element) + let prompt = if prompt_col == "turns" { + match prompt_val.as_array().and_then(|arr| arr.first()) { + Some(v) => v.as_str().unwrap_or("").to_string(), + None => continue, + } + } else { + prompt_val.as_str().unwrap_or("").to_string() + }; + + let completion = output_col.as_ref().and_then(|col| { + row.get(col.as_str()).and_then(|v| { + // Handle "answers" which may be an array + if let Some(arr) = v.as_array() { + arr.first().and_then(|a| a.as_str()).map(|s| s.to_string()) + } else { + v.as_str().map(|s| s.to_string()) + } + }) + }); + + (prompt, completion) + } + ColumnFormat::Combined { cols, output_col } => { + let parts: Vec = cols + .iter() + .filter_map(|col| { + row.get(col.as_str()).and_then(|v| v.as_str()).map(|s| s.to_string()) + }) + .collect(); + + if parts.is_empty() { + continue; + } + + let prompt = parts.join("\n\n"); + + let completion = output_col.as_ref().and_then(|col| { + row.get(col.as_str()).and_then(|v| v.as_str()).map(|s| s.to_string()) + }); + + (prompt, completion) + } + }; + + if prompt.is_empty() { + continue; + } + + // Tokenize prompt + let prompt_ids = tokenizer.encode(&prompt, false)?; + let prompt_len = prompt_ids.len(); + + if prompt_len < 4 { + continue; + } + + // Determine output length + let output_len = if let Some(fixed_len) = hf_output_len { + fixed_len + } else if let Some(ref comp) = completion { + if !comp.is_empty() { + let comp_ids = tokenizer.encode(comp, false)?; + let len = comp_ids.len(); + if len == 0 { 128 } else { len } + } else { + if !warned_no_output { + eprintln!( + "WARNING: No output column detected and --hf-output-len not set. \ + Using default output length of 128 tokens." + ); + warned_no_output = true; + } + 128 + } + } else { + if !warned_no_output { + eprintln!( + "WARNING: No output column detected and --hf-output-len not set. \ + Using default output length of 128 tokens." + ); + warned_no_output = true; + } + 128 + }; + + samples.push(SampleRequest { + prompt: Arc::from(prompt.as_str()), + prompt_len, + expected_output_len: output_len, + request_id: Some(format!("{request_id_prefix}{idx}")), + ..Default::default() + }); + idx += 1; + } + + // Oversample if needed + if samples.len() < num_requests { + if no_oversample { + println!( + "Skipping oversampling. Total samples: {} (requested: {num_requests})", + samples.len() + ); + } else if !samples.is_empty() { + let original_len = samples.len(); + let needed = num_requests - original_len; + for i in 0..needed { + let mut req = samples[rng.random_range(0..original_len)].clone(); + req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); + samples.push(req); + } + println!( + "Oversampled HF dataset from {original_len} to {} total samples.", + samples.len() + ); + } + } + + if samples.is_empty() { + return Err(BenchError::Config( + "No valid samples after processing HF dataset. \ + Try a different --hf-text-column or check the dataset format." + .into(), + )); + } + + Ok(samples) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_detect_chat_conversation() { + let row = serde_json::json!({ + "conversation_id": "abc", + "conversation": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"} + ] + }); + let format = detect_column_format(&row, None).unwrap(); + assert!(matches!(format, ColumnFormat::Chat(ref col) if col == "conversation")); + } + + #[test] + fn test_detect_chat_messages() { + let row = serde_json::json!({ + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi"} + ] + }); + let format = detect_column_format(&row, None).unwrap(); + assert!(matches!(format, ColumnFormat::Chat(ref col) if col == "messages")); + } + + #[test] + fn test_detect_chat_sharegpt_format() { + let row = serde_json::json!({ + "conversations": [ + {"from": "human", "value": "Hello"}, + {"from": "gpt", "value": "Hi!"} + ] + }); + let format = detect_column_format(&row, None).unwrap(); + assert!(matches!(format, ColumnFormat::Chat(ref col) if col == "conversations")); + } + + #[test] + fn test_detect_plain_prompt() { + let row = serde_json::json!({ + "prompt": "What is 2+2?", + "completion": "4" + }); + let format = detect_column_format(&row, None).unwrap(); + match format { + ColumnFormat::Text { + ref prompt_col, + ref output_col, + } => { + assert_eq!(prompt_col, "prompt"); + assert_eq!(output_col.as_deref(), Some("completion")); + } + _ => panic!("Expected Text format"), + } + } + + #[test] + fn test_detect_question_answer() { + let row = serde_json::json!({ + "question": "What is AI?", + "answer": "Artificial intelligence" + }); + let format = detect_column_format(&row, None).unwrap(); + match format { + ColumnFormat::Text { + ref prompt_col, + ref output_col, + } => { + assert_eq!(prompt_col, "question"); + assert_eq!(output_col.as_deref(), Some("answer")); + } + _ => panic!("Expected Text format"), + } + } + + #[test] + fn test_detect_combined_context_input() { + let row = serde_json::json!({ + "context": "The quick brown fox...", + "input": "What animal was mentioned?", + "answers": ["fox"] + }); + let format = detect_column_format(&row, None).unwrap(); + match format { + ColumnFormat::Combined { + ref cols, + ref output_col, + } => { + assert_eq!(cols, &["context", "input"]); + assert_eq!(output_col.as_deref(), Some("answers")); + } + _ => panic!("Expected Combined format"), + } + } + + #[test] + fn test_detect_turns_array() { + let row = serde_json::json!({ + "turns": ["First turn prompt", "Second turn"], + "answer": "The answer" + }); + let format = detect_column_format(&row, None).unwrap(); + match format { + ColumnFormat::Text { + ref prompt_col, + ref output_col, + } => { + assert_eq!(prompt_col, "turns"); + assert_eq!(output_col.as_deref(), Some("answer")); + } + _ => panic!("Expected Text format"), + } + } + + #[test] + fn test_detect_user_override() { + let row = serde_json::json!({ + "my_custom_col": "Hello world", + "answer": "response" + }); + let format = detect_column_format(&row, Some("my_custom_col")).unwrap(); + match format { + ColumnFormat::Text { + ref prompt_col, + ref output_col, + } => { + assert_eq!(prompt_col, "my_custom_col"); + assert_eq!(output_col.as_deref(), Some("answer")); + } + _ => panic!("Expected Text format"), + } + } + + #[test] + fn test_detect_user_override_chat_column() { + let row = serde_json::json!({ + "my_chat": [ + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello"} + ] + }); + let format = detect_column_format(&row, Some("my_chat")).unwrap(); + assert!(matches!(format, ColumnFormat::Chat(ref col) if col == "my_chat")); + } + + #[test] + fn test_detect_user_override_missing_column() { + let row = serde_json::json!({"text": "hello"}); + let result = detect_column_format(&row, Some("nonexistent")); + assert!(result.is_err()); + } + + #[test] + fn test_detect_no_known_columns() { + let row = serde_json::json!({"id": 1, "label": "positive"}); + let result = detect_column_format(&row, None); + assert!(result.is_err()); + } + + #[test] + fn test_extract_chat_prompt_role_content() { + let messages = vec![ + serde_json::json!({"role": "user", "content": "What is AI?"}), + serde_json::json!({"role": "assistant", "content": "AI is..."}), + ]; + assert_eq!( + extract_chat_prompt(&messages), + Some("What is AI?".to_string()) + ); + } + + #[test] + fn test_extract_chat_prompt_from_value() { + let messages = vec![ + serde_json::json!({"from": "human", "value": "Hello"}), + serde_json::json!({"from": "gpt", "value": "Hi!"}), + ]; + assert_eq!(extract_chat_prompt(&messages), Some("Hello".to_string())); + } + + #[test] + fn test_extract_chat_prompt_system_first() { + let messages = vec![ + serde_json::json!({"role": "system", "content": "You are helpful"}), + serde_json::json!({"role": "user", "content": "Tell me a joke"}), + serde_json::json!({"role": "assistant", "content": "Why did..."}), + ]; + assert_eq!( + extract_chat_prompt(&messages), + Some("Tell me a joke".to_string()) + ); + } + + #[test] + fn test_extract_chat_prompt_no_user() { + let messages = vec![serde_json::json!({"role": "system", "content": "You are helpful"})]; + assert_eq!(extract_chat_prompt(&messages), None); + } + + #[test] + fn test_extract_chat_completion_role_content() { + let messages = vec![ + serde_json::json!({"role": "user", "content": "Hi"}), + serde_json::json!({"role": "assistant", "content": "Hello!"}), + ]; + assert_eq!( + extract_chat_completion(&messages), + Some("Hello!".to_string()) + ); + } + + #[test] + fn test_extract_chat_completion_gpt() { + let messages = vec![ + serde_json::json!({"from": "human", "value": "Hi"}), + serde_json::json!({"from": "gpt", "value": "Hello!"}), + ]; + assert_eq!( + extract_chat_completion(&messages), + Some("Hello!".to_string()) + ); + } + + #[test] + fn test_is_chat_message_valid() { + assert!(is_chat_message( + &serde_json::json!({"role": "user", "content": "hi"}) + )); + assert!(is_chat_message( + &serde_json::json!({"from": "human", "value": "hi"}) + )); + } + + #[test] + fn test_is_chat_message_invalid() { + assert!(!is_chat_message(&serde_json::json!({"text": "hi"}))); + assert!(!is_chat_message(&serde_json::json!("just a string"))); + assert!(!is_chat_message(&serde_json::json!(42))); + } + + #[test] + fn test_is_chat_array_valid() { + let val = serde_json::json!([ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"} + ]); + assert!(is_chat_array(&val)); + } + + #[test] + fn test_is_chat_array_invalid() { + assert!(!is_chat_array(&serde_json::json!([]))); + assert!(!is_chat_array(&serde_json::json!(["a", "b"]))); + assert!(!is_chat_array(&serde_json::json!("not an array"))); + } + + #[test] + fn test_sanitize_name() { + assert_eq!( + sanitize_name("allenai/WildChat-4.8M"), + "allenai_WildChat-4.8M" + ); + assert_eq!(sanitize_name("simple-name"), "simple-name"); + assert_eq!(sanitize_name("org/sub/deep"), "org_sub_deep"); + } + + // --- extract_chat_prompt edge cases --- + + #[test] + fn test_extract_chat_prompt_empty_messages() { + let messages: Vec = vec![]; + assert_eq!(extract_chat_prompt(&messages), None); + } + + #[test] + fn test_extract_chat_prompt_model_role_skipped() { + // "model" role is not "user" or "human" — should be skipped + let messages = vec![ + serde_json::json!({"role": "model", "content": "I am the model"}), + serde_json::json!({"role": "assistant", "content": "Hello"}), + ]; + assert_eq!(extract_chat_prompt(&messages), None); + } + + #[test] + fn test_extract_chat_prompt_model_role_before_user() { + // "model" role before user: only "user" should be matched + let messages = vec![ + serde_json::json!({"role": "model", "content": "intro"}), + serde_json::json!({"role": "user", "content": "actual user message"}), + ]; + assert_eq!( + extract_chat_prompt(&messages), + Some("actual user message".to_string()) + ); + } + + // --- extract_chat_completion edge cases --- + + #[test] + fn test_extract_chat_completion_no_assistant() { + let messages = vec![ + serde_json::json!({"role": "user", "content": "What is 2+2?"}), + serde_json::json!({"role": "system", "content": "You are helpful"}), + ]; + assert_eq!(extract_chat_completion(&messages), None); + } + + #[test] + fn test_extract_chat_completion_empty_messages() { + let messages: Vec = vec![]; + assert_eq!(extract_chat_completion(&messages), None); + } + + #[test] + fn test_extract_chat_completion_model_role_not_matched() { + // "model" role is not "assistant" or "gpt" — should return None + let messages = vec![ + serde_json::json!({"role": "user", "content": "Hello"}), + serde_json::json!({"role": "model", "content": "I should not be returned"}), + ]; + assert_eq!(extract_chat_completion(&messages), None); + } + + // --- load_hf_dataset integration tests (require built-in tiktoken, no network) --- + + /// Build a gpt2 tokenizer using built-in tiktoken encoding (no network required). + fn builtin_tokenizer() -> crate::tokenizer::TokenizerKind { + crate::tokenizer::load_tokenizer("gpt2", false, None) + .expect("gpt2 built-in tiktoken should always load without network") + } + + /// Write JSON data to a unique temp file and return the path string. + fn write_temp_json(name: &str, data: &serde_json::Value) -> String { + let path = std::env::temp_dir().join(format!("vllm-bench-test-{name}.json")); + std::fs::write(&path, serde_json::to_string(data).unwrap()).unwrap(); + path.to_string_lossy().to_string() + } + + #[test] + fn test_load_hf_dataset_chat_format() { + let tok = builtin_tokenizer(); + let data = serde_json::json!([ + { + "conversation": [ + {"role": "user", "content": "What is the meaning of life, the universe, and everything?"}, + {"role": "assistant", "content": "The answer is 42, according to Douglas Adams."} + ] + }, + { + "conversation": [ + {"role": "user", "content": "Tell me about quantum computing and its applications in modern science."}, + {"role": "assistant", "content": "Quantum computing uses quantum bits to perform calculations."} + ] + } + ]); + let path = write_temp_json("chat-format", &data); + + let result = + load_hf_dataset(&tok, &path, 2, Some(50), 42, "test-", None, false, false).unwrap(); + + assert_eq!(result.len(), 2); + assert!( + result.iter().all(|r| r.expected_output_len == 50), + "all samples should have fixed output len 50" + ); + } + + #[test] + fn test_load_hf_dataset_plain_text_format() { + let tok = builtin_tokenizer(); + let data = serde_json::json!([ + {"prompt": "What is the capital of France and why is it important in European history?", "completion": "Paris is the capital of France."}, + {"prompt": "Explain the difference between machine learning and deep learning in simple terms.", "completion": "Machine learning is a subset of AI."}, + {"prompt": "How does photosynthesis work in plants and what role does chlorophyll play?", "completion": "Photosynthesis converts light to energy."} + ]); + let path = write_temp_json("plain-text-format", &data); + + let result = load_hf_dataset(&tok, &path, 3, None, 0, "req-", None, false, false).unwrap(); + + assert_eq!(result.len(), 3); + // Without hf_output_len override, output len is derived from the completion tokens + assert!(result.iter().all(|r| r.expected_output_len > 0)); + } + + #[test] + fn test_load_hf_dataset_combined_format() { + let tok = builtin_tokenizer(); + let data = serde_json::json!([ + { + "context": "The quick brown fox jumped over the lazy dog near the river bank.", + "input": "What animal jumped over the dog in this sentence?", + "answers": ["The quick brown fox"] + }, + { + "context": "Albert Einstein developed the theory of relativity in the early twentieth century.", + "input": "Who developed the theory of relativity and when approximately?", + "answers": ["Albert Einstein"] + } + ]); + let path = write_temp_json("combined-format", &data); + + let result = + load_hf_dataset(&tok, &path, 2, Some(64), 1, "comb-", None, false, false).unwrap(); + + assert_eq!(result.len(), 2); + assert!(result.iter().all(|r| r.expected_output_len == 64)); + // Combined format joins context + input with "\n\n" + assert!(result.iter().all(|r| r.prompt.contains('\n'))); + } + + #[test] + fn test_load_hf_dataset_empty_returns_error() { + let tok = builtin_tokenizer(); + let data = serde_json::json!([]); + let path = write_temp_json("empty-dataset", &data); + + let result = load_hf_dataset(&tok, &path, 5, None, 0, "test-", None, false, false); + + assert!(result.is_err(), "empty dataset should return an error"); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("no rows") || err.contains("No rows") || err.contains("empty"), + "error should mention empty/no rows: {err}" + ); + } + + #[test] + fn test_load_hf_dataset_all_rows_filtered_too_short() { + let tok = builtin_tokenizer(); + // Very short prompts that will tokenize to fewer than 4 tokens and be filtered out + let data = serde_json::json!([ + {"prompt": "Hi", "completion": "Ok"}, + {"prompt": "Yes", "completion": "No"}, + {"prompt": "Ok", "completion": "Fine"} + ]); + let path = write_temp_json("all-filtered", &data); + + let result = load_hf_dataset(&tok, &path, 3, None, 0, "test-", None, false, false); + + assert!( + result.is_err(), + "all rows being too short should return an error" + ); + } + + #[test] + fn test_load_hf_dataset_hf_output_len_override() { + let tok = builtin_tokenizer(); + let data = serde_json::json!([ + {"prompt": "Describe the history of the Roman Empire and its eventual decline over centuries.", "completion": "The Roman Empire fell for many complex reasons."}, + {"prompt": "What are the key differences between supervised and unsupervised machine learning?", "completion": "Supervised learning uses labeled data while unsupervised does not."}, + {"prompt": "Explain how neural networks are inspired by the human brain structure and function.", "completion": "Neural networks mimic brain neurons with layers of nodes."} + ]); + let path = write_temp_json("output-len-override", &data); + + let fixed_len = 77usize; + let result = load_hf_dataset( + &tok, + &path, + 3, + Some(fixed_len), + 42, + "t-", + None, + false, + false, + ) + .unwrap(); + + assert!(!result.is_empty()); + assert!( + result.iter().all(|r| r.expected_output_len == fixed_len), + "all samples must use the fixed output length {fixed_len}" + ); + } + + #[test] + fn test_load_hf_dataset_no_oversample() { + let tok = builtin_tokenizer(); + // 2 valid rows, but request 10 with no_oversample=true + let data = serde_json::json!([ + {"prompt": "Explain how photosynthesis converts sunlight to energy in plant cells.", "completion": "Plants use chlorophyll to absorb sunlight."}, + {"prompt": "What is the difference between a virus and a bacterium in terms of biology?", "completion": "Viruses need host cells while bacteria are self-sufficient."} + ]); + let path = write_temp_json("no-oversample", &data); + + let result = load_hf_dataset( + &tok, + &path, + 10, + Some(32), + 0, + "t-", + None, + true, // no_oversample + false, + ) + .unwrap(); + + // Should have at most 2 samples (the actual dataset size), not 10 + assert!( + result.len() <= 2, + "no_oversample should cap result at dataset size, got {}", + result.len() + ); + } + + #[test] + fn test_load_hf_dataset_disable_shuffle_preserves_order() { + let tok = builtin_tokenizer(); + let data = serde_json::json!([ + {"prompt": "Alpha prompt about the first topic in our alphabetically ordered sequence.", "completion": "Alpha answer"}, + {"prompt": "Beta prompt about the second topic continuing from our ordered sequence here.", "completion": "Beta answer"}, + {"prompt": "Gamma prompt about the third item in our clearly ordered alphabetical sequence.", "completion": "Gamma answer"} + ]); + let path = write_temp_json("disable-shuffle", &data); + + let result = load_hf_dataset( + &tok, + &path, + 3, + Some(20), + 99, + "t-", + None, + false, + true, // disable_shuffle + ) + .unwrap(); + + assert_eq!(result.len(), 3); + // With disable_shuffle, rows come out in original order: + // Alpha < Beta < Gamma (alphabetical), so first prompt contains "Alpha" + assert!( + result[0].prompt.to_lowercase().contains("alpha"), + "first result should be Alpha prompt with shuffle disabled, got: {}", + result[0].prompt + ); + assert!( + result[1].prompt.to_lowercase().contains("beta"), + "second result should be Beta prompt, got: {}", + result[1].prompt + ); + } + + #[test] + fn test_load_hf_dataset_request_id_prefix() { + let tok = builtin_tokenizer(); + let data = serde_json::json!([ + {"prompt": "What is the boiling point of water at sea level under standard atmospheric pressure?", "completion": "Water boils at 100 degrees Celsius."}, + {"prompt": "Describe the structure of DNA and how genetic information is encoded within it.", "completion": "DNA is a double helix with base pairs."} + ]); + let path = write_temp_json("request-id-prefix", &data); + + let result = + load_hf_dataset(&tok, &path, 2, Some(10), 0, "myprefix-", None, false, false).unwrap(); + + assert_eq!(result.len(), 2); + assert_eq!(result[0].request_id.as_deref(), Some("myprefix-0")); + assert_eq!(result[1].request_id.as_deref(), Some("myprefix-1")); + } +} diff --git a/rust/src/bench/src/datasets/mod.rs b/rust/src/bench/src/datasets/mod.rs new file mode 100644 index 000000000000..3917b4d833c1 --- /dev/null +++ b/rust/src/bench/src/datasets/mod.rs @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +pub mod custom; +pub mod hf_dataset; +pub mod multi_turn; +pub mod prefix_repetition; +pub mod random; +pub mod random_mm; +pub mod random_rerank; +pub mod sharegpt; +pub mod sonnet; +pub mod speed_bench; + +use std::sync::Arc; + +/// Represents a single inference request for benchmarking. +/// Matches Python's SampleRequest dataclass from datasets.py:71-82. +/// +/// `prompt` uses `Arc` to avoid expensive String clones when distributing +/// requests across tokio tasks. At 100k prompts with 8k tokens each, this saves +/// ~3GB of peak memory vs cloning String per task. +#[derive(Debug, Clone)] +pub struct SampleRequest { + pub prompt: Arc, + pub prompt_len: usize, + pub expected_output_len: usize, + pub request_id: Option, + /// Pre-computed token IDs for this prompt. + /// When set, the completions backend sends these directly via `prompt_token_ids` + /// instead of the text `prompt`, avoiding server-side re-tokenization. + pub prompt_token_ids: Option>, + /// Multimodal content items as pre-serialized JSON fragments. + /// Each `Arc` is a complete JSON object string, e.g. + /// `{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,..."}}` + /// + /// Pre-serialized to avoid: + /// 1. `serde_json::Value` tree overhead (3 Maps + keys per image) + /// 2. Deep-cloning ~200KB+ base64 data when building request payloads + /// + /// Double-`Arc` for zero-cost sharing: outer Arc for the slice, inner Arc for each fragment. + pub multi_modal_content: Option]>>, + /// Pre-serialized OpenAI chat `messages` array as a complete JSON string, + /// e.g. `[{"role":"user","content":[{"type":"text","text":"..."},{"type":"image_url",...}]}]`. + /// + /// Set by datasets when `--enable-multimodal-chat` is on (mirrors Python's + /// `apply_multimodal_chat_transformation`: the dataset builds the chat messages + /// and the backend sends them verbatim). When set, `multi_modal_content` is None + /// and the mm items are embedded here instead. `prompt` still holds the text part + /// for token accounting and /tokenize verification. + pub chat_messages_json: Option>, + /// Multiple text inputs for one request (pooling backends only). + /// Embeddings send it as `"input": [t1, t2, ...]` (--random-batch-size); + /// rerank sends `[0]` as the query and `[1..]` as documents (random-rerank). + /// Mirrors Python's list-valued `SampleRequest.prompt`. + pub prompt_list: Option]>>, +} + +impl Default for SampleRequest { + /// Empty request; struct-update base so dataset builders only spell out the + /// fields they set (new optional fields then don't touch every call site). + fn default() -> Self { + Self { + prompt: Arc::from(""), + prompt_len: 0, + expected_output_len: 0, + request_id: None, + prompt_token_ids: None, + multi_modal_content: None, + chat_messages_json: None, + prompt_list: None, + } + } +} + +/// Oversample `requests` up to `num_requests` by cloning random entries +/// (seeded by list length for determinism), renumbering their request ids. +/// No-op when enough samples exist, `no_oversample` is set, or the list is empty. +/// Mirrors Python `BenchmarkDataset.maybe_oversample_requests`. +pub fn oversample_requests( + requests: &mut Vec, + num_requests: usize, + request_id_prefix: &str, + no_oversample: bool, +) { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + if requests.len() >= num_requests || requests.is_empty() { + return; + } + if no_oversample { + println!( + "Skipping oversampling. Total samples: {} (requested: {num_requests})", + requests.len() + ); + return; + } + let original_len = requests.len(); + let mut rng = StdRng::seed_from_u64(original_len as u64); + for i in 0..(num_requests - original_len) { + let mut req = requests[rng.random_range(0..original_len)].clone(); + req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); + requests.push(req); + } + println!( + "Oversampled requests from {original_len} to {} total samples.", + requests.len() + ); +} + +/// Group already-generated single-input requests into batched requests of +/// `batch_size` inputs each (embeddings/pooling only). Mirrors Python +/// `RandomDataset.sample` batching: prompt becomes a list, prompt_len is the +/// sum over the batch, request ids are renumbered per batch. +/// `batch_size <= 1` returns the input unchanged. +pub fn batch_requests( + requests: Vec, + batch_size: usize, + request_id_prefix: &str, +) -> Vec { + if batch_size <= 1 { + return requests; + } + requests + .chunks(batch_size) + .enumerate() + .map(|(batch_idx, batch)| SampleRequest { + prompt_list: Some(batch.iter().map(|r| r.prompt.clone()).collect()), + prompt_len: batch.iter().map(|r| r.prompt_len).sum(), + expected_output_len: 0, + request_id: Some(format!("{request_id_prefix}{batch_idx}")), + ..Default::default() + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn req(prompt: &str, len: usize) -> SampleRequest { + SampleRequest { + prompt: Arc::from(prompt), + prompt_len: len, + expected_output_len: 128, + ..Default::default() + } + } + + #[test] + fn test_batch_requests_groups_and_sums() { + let reqs = vec![ + req("a", 10), + req("b", 20), + req("c", 30), + req("d", 40), + req("e", 50), + ]; + let batched = batch_requests(reqs, 2, "t-"); + assert_eq!(batched.len(), 3); // 2 + 2 + 1 + let first = batched[0].prompt_list.as_ref().unwrap(); + assert_eq!(first.len(), 2); + assert_eq!(&*first[0], "a"); + assert_eq!(batched[0].prompt_len, 30); + assert_eq!(batched[0].expected_output_len, 0); + assert_eq!(batched[0].request_id.as_deref(), Some("t-0")); + assert_eq!(batched[2].prompt_list.as_ref().unwrap().len(), 1); + assert_eq!(batched[2].prompt_len, 50); + } + + #[test] + fn test_batch_requests_size_one_is_identity() { + let reqs = vec![req("a", 10), req("b", 20)]; + let out = batch_requests(reqs, 1, "t-"); + assert_eq!(out.len(), 2); + assert!(out[0].prompt_list.is_none()); + assert_eq!(&*out[0].prompt, "a"); + } + + #[test] + fn test_oversample_requests() { + let mut reqs = vec![req("a", 10), req("b", 20)]; + oversample_requests(&mut reqs, 5, "t-", false); + assert_eq!(reqs.len(), 5); + assert_eq!(reqs[4].request_id.as_deref(), Some("t-4")); + + let mut reqs = vec![req("a", 10)]; + oversample_requests(&mut reqs, 5, "t-", true); // no_oversample + assert_eq!(reqs.len(), 1); + } +} + +/// A single turn in a multi-turn conversation. +#[derive(Debug, Clone)] +pub struct ConversationTurn { + pub user_message: Arc, + pub user_message_len: usize, + pub expected_output_len: usize, +} + +/// A complete multi-turn conversation with all turns pre-generated. +#[derive(Debug, Clone)] +pub struct MultiTurnConversation { + pub conversation_id: String, + pub turns: Vec, +} diff --git a/rust/src/bench/src/datasets/multi_turn.rs b/rust/src/bench/src/datasets/multi_turn.rs new file mode 100644 index 000000000000..43098b9725e3 --- /dev/null +++ b/rust/src/bench/src/datasets/multi_turn.rs @@ -0,0 +1,777 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; + +use super::{ConversationTurn, MultiTurnConversation}; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +/// Configuration for generating random multi-turn conversations. +#[derive(Debug, Clone)] +pub struct MultiTurnRandomConfig { + pub num_conversations: usize, + pub min_turns: usize, + pub max_turns: usize, + /// Shared prefix length prepended to the conversation. + /// + /// In normal accumulated-history mode this is added to turn 0, so all + /// later turns inherit it through history. In no-history prefix-sharing + /// mode it is added to every independent turn. + pub prefix_len: usize, + /// Input length for turn 0. + pub input_len: usize, + /// Input length for turns 1+. 0 = fallback to input_len. + pub per_turn_input_len: usize, + pub output_len: usize, + pub seed: u64, + pub request_id_prefix: String, + pub prefix_sharing_config: Option, +} + +/// Configuration for 3-tier prefix sharing in multi-turn user messages. +#[derive(Debug, Clone)] +pub struct PrefixSharingConfig { + /// Fraction of per-turn input tokens shared across ALL conversations. + pub global_ratio: f64, + /// Fraction of per-turn input tokens shared within each conversation. + pub conversation_ratio: f64, +} + +/// Generate a deterministic token sequence from allowed tokens using offset+modulo. +fn make_token_seq(allowed_tokens: &[u32], offset: usize, len: usize) -> Vec { + let at_len = allowed_tokens.len(); + (0..len).map(|i| allowed_tokens[(offset + i) % at_len]).collect() +} + +/// Generate synthetic multi-turn conversations with random user messages. +/// +/// Each conversation has `num_turns` turns, each with a random user prompt +/// of `input_len` tokens and `output_len` expected output tokens. +pub fn generate_multi_turn_random( + tokenizer: &TokenizerKind, + cfg: &MultiTurnRandomConfig, +) -> Result> { + let num_conversations = cfg.num_conversations; + let min_turns = cfg.min_turns; + let max_turns = cfg.max_turns; + let prefix_len = cfg.prefix_len; + let input_len = cfg.input_len; + let output_len = cfg.output_len; + let seed = cfg.seed; + let request_id_prefix = &cfg.request_id_prefix; + let allowed_tokens = tokenizer.get_allowed_tokens(); + if allowed_tokens.is_empty() { + return Err(BenchError::Tokenizer("No allowed tokens found".into())); + } + + let vocab_size = tokenizer.vocab_size() as usize; + let num_special = tokenizer.num_special_tokens_to_add(); + let real_input_len = input_len.saturating_sub(num_special); + let real_per_turn_len = if cfg.per_turn_input_len > 0 { + cfg.per_turn_input_len.saturating_sub(num_special) + } else { + real_input_len + }; + + if real_input_len < 1 { + return Err(BenchError::Config(format!( + "--random-input-len too small: with {num_special} special tokens, \ + effective input length is {real_input_len}" + ))); + } + if real_per_turn_len < 1 { + return Err(BenchError::Config(format!( + "--per-turn-input-len too small: with {num_special} special tokens, \ + effective per-turn input length is {real_per_turn_len}" + ))); + } + + // Prefix sharing mode: generate 3-tier prefixed messages + let mut rng = StdRng::seed_from_u64(seed); + if let Some(ref ps_cfg) = cfg.prefix_sharing_config { + return generate_prefix_sharing_conversations( + tokenizer, + cfg, + ps_cfg, + &allowed_tokens, + &mut rng, + ); + } + let shared_prefix_text = + generate_shared_prefix_text(tokenizer, &allowed_tokens, prefix_len, seed)?; + + // Pre-generate per-conversation turn counts and per-turn offsets deterministically. + // Turn counts are drawn first so the RNG sequence is stable regardless of vocab_size. + let conv_turn_counts: Vec = (0..num_conversations) + .map(|_| { + if min_turns == max_turns { + min_turns + } else { + rng.random_range(min_turns..=max_turns) + } + }) + .collect(); + + let offsets: Vec> = conv_turn_counts + .iter() + .map(|&n| (0..n).map(|_| rng.random_range(0..vocab_size)).collect()) + .collect(); + + // Parallel generation across conversations + offsets + .par_iter() + .enumerate() + .map(|(conv_idx, conv_offsets)| { + let mut turns = Vec::with_capacity(conv_offsets.len()); + for (turn_idx, &offset) in conv_offsets.iter().enumerate() { + let target_len = if turn_idx == 0 { + real_input_len + } else { + real_per_turn_len + }; + // Use max_turns stride to keep offsets unique across variable-length convs + let inner_seq = make_token_seq( + &allowed_tokens, + offset + conv_idx * max_turns + turn_idx, + target_len, + ); + + let (prompt, adjusted) = + gen_prompt_to_target_len(tokenizer, &inner_seq, target_len)?; + let (prompt, token_len) = if turn_idx == 0 && !shared_prefix_text.is_empty() { + let combined = format!("{}{}", &*shared_prefix_text, prompt); + let token_len = tokenizer.encode(&combined, false)?.len(); + (combined, token_len) + } else { + (prompt, adjusted.len()) + }; + + turns.push(ConversationTurn { + user_message: Arc::from(prompt), + user_message_len: token_len, + expected_output_len: output_len, + }); + } + + Ok(MultiTurnConversation { + conversation_id: format!("{request_id_prefix}conv-{conv_idx}"), + turns, + }) + }) + .collect() +} + +/// Generate conversations with 3-tier prefix sharing. +/// +/// Each turn's user message = [global_prefix][conversation_prefix][unique_suffix]. +/// No history accumulation — each turn sends only its own fixed-length message. +fn generate_prefix_sharing_conversations( + tokenizer: &TokenizerKind, + cfg: &MultiTurnRandomConfig, + ps_cfg: &PrefixSharingConfig, + allowed_tokens: &[u32], + rng: &mut StdRng, +) -> Result> { + let num_conversations = cfg.num_conversations; + let min_turns = cfg.min_turns; + let max_turns = cfg.max_turns; + let prefix_len = cfg.prefix_len; + let output_len = cfg.output_len; + let request_id_prefix = &cfg.request_id_prefix; + + let num_special = tokenizer.num_special_tokens_to_add(); + let real_input_len = cfg.input_len.saturating_sub(num_special); + let real_per_turn_len = if cfg.per_turn_input_len > 0 { + cfg.per_turn_input_len.saturating_sub(num_special) + } else { + real_input_len + }; + + // Compute segment lengths from turn-0 (real_input_len) so the shared prefix + // bytes stay byte-identical across all turns regardless of per_turn_input_len. + let global_len = (real_input_len as f64 * ps_cfg.global_ratio).floor() as usize; + let conv_len = (real_input_len as f64 * ps_cfg.conversation_ratio).floor() as usize; + let unique_len = real_input_len.saturating_sub(global_len + conv_len); + + // Validate that turns 1+ still have room for a non-empty unique suffix + if real_per_turn_len <= global_len + conv_len { + return Err(BenchError::Config(format!( + "--per-turn-input-len ({real_per_turn_len} after special tokens) is too small: \ + global_len={global_len} + conv_len={conv_len} already fills the budget. \ + Increase --per-turn-input-len or reduce prefix ratios." + ))); + } + + let at_len = allowed_tokens.len(); + let shared_prefix_text = + generate_shared_prefix_text(tokenizer, allowed_tokens, prefix_len, cfg.seed)?; + + // Generate global prefix text once + let global_text: Arc = if global_len > 0 { + let offset: usize = rng.random_range(0..at_len); + let seq = make_token_seq(allowed_tokens, offset, global_len); + let (text, _) = gen_prompt_to_target_len(tokenizer, &seq, global_len)?; + Arc::from(text) + } else { + Arc::from("") + }; + + // Generate per-conversation prefix texts + let conv_texts: Vec> = if conv_len > 0 { + let mut texts = Vec::with_capacity(num_conversations); + for conv_idx in 0..num_conversations { + let offset: usize = rng.random_range(0..at_len); + let seq = make_token_seq(allowed_tokens, offset + conv_idx, conv_len); + let (text, _) = gen_prompt_to_target_len(tokenizer, &seq, conv_len)?; + texts.push(Arc::from(text)); + } + texts + } else { + vec![Arc::from(""); num_conversations] + }; + + // Pre-generate per-conversation turn counts and unique offsets deterministically. + let vocab_size = tokenizer.vocab_size() as usize; + let conv_turn_counts: Vec = (0..num_conversations) + .map(|_| { + if min_turns == max_turns { + min_turns + } else { + rng.random_range(min_turns..=max_turns) + } + }) + .collect(); + + let unique_offsets: Vec> = conv_turn_counts + .iter() + .map(|&n| (0..n).map(|_| rng.random_range(0..vocab_size)).collect()) + .collect(); + + // Parallel generation across conversations + unique_offsets + .par_iter() + .enumerate() + .map(|(conv_idx, conv_offsets)| { + let mut turns = Vec::with_capacity(conv_offsets.len()); + for (turn_idx, &offset) in conv_offsets.iter().enumerate() { + // Turn 0 uses unique_len derived from real_input_len; + // turns 1+ use per-turn unique_len (prefix bytes stay identical). + let turn_unique_len = if turn_idx == 0 { + unique_len + } else { + real_per_turn_len.saturating_sub(global_len + conv_len) + }; + + // Generate unique suffix + let unique_text = if turn_unique_len > 0 { + let seq = make_token_seq( + allowed_tokens, + offset + conv_idx * max_turns + turn_idx, + turn_unique_len, + ); + let (text, _) = gen_prompt_to_target_len(tokenizer, &seq, turn_unique_len)?; + text + } else { + String::new() + }; + + // Concatenate: optional random prefix + global + conversation + unique. + // Prefix-sharing mode sends each turn independently, so the random + // prefix must be included on every turn to be present in every request. + let combined = format!( + "{}{}{}{}", + &*shared_prefix_text, &*global_text, &*conv_texts[conv_idx], unique_text + ); + // Re-encode to get actual token count (BPE boundary effects) + let token_len = tokenizer.encode(&combined, false)?.len(); + + turns.push(ConversationTurn { + user_message: Arc::from(combined), + user_message_len: token_len, + expected_output_len: output_len, + }); + } + + Ok(MultiTurnConversation { + conversation_id: format!("{request_id_prefix}conv-{conv_idx}"), + turns, + }) + }) + .collect() +} + +fn generate_shared_prefix_text( + tokenizer: &TokenizerKind, + allowed_tokens: &[u32], + prefix_len: usize, + seed: u64, +) -> Result> { + if prefix_len == 0 { + return Ok(Arc::from("")); + } + + let mut rng = StdRng::seed_from_u64(seed.wrapping_add(0xDEAD)); + let tokens: Vec = (0..prefix_len) + .map(|_| allowed_tokens[rng.random_range(0..allowed_tokens.len())]) + .collect(); + let (text, _) = gen_prompt_to_target_len(tokenizer, &tokens, prefix_len)?; + Ok(Arc::from(text)) +} + +/// Load multi-turn conversations from a ShareGPT dataset. +/// +/// Walks ALL turns in each entry (not just first 2). Filters entries +/// with at least 4 messages (2 user + 2 assistant = 2 real turns). +pub fn load_sharegpt_multi_turn( + tokenizer: &TokenizerKind, + dataset_path: &str, + num_conversations: usize, + output_len_override: Option, + max_turns: Option, + seed: u64, + request_id_prefix: &str, +) -> Result> { + let content = std::fs::read_to_string(dataset_path).map_err(|e| { + BenchError::Config(format!( + "Failed to read ShareGPT file '{dataset_path}': {e}" + )) + })?; + + let data: serde_json::Value = serde_json::from_str(&content) + .map_err(|e| BenchError::Config(format!("Invalid JSON in ShareGPT file: {e}")))?; + + let entries = data + .as_array() + .ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?; + + // Filter entries with at least 4 messages (2 turns: user+assistant+user+assistant) + let mut filtered: Vec<&serde_json::Value> = entries + .iter() + .filter(|entry| { + entry + .get("conversations") + .and_then(|c| c.as_array()) + .map(|a| a.len() >= 4) + .unwrap_or(false) + }) + .collect(); + + if filtered.is_empty() { + return Err(BenchError::Config( + "No valid multi-turn entries in ShareGPT file (need at least 4 messages per entry)" + .into(), + )); + } + + // Shuffle + let mut rng = StdRng::seed_from_u64(seed); + filtered.shuffle(&mut rng); + + let mut conversations = Vec::new(); + + for entry in &filtered { + if conversations.len() >= num_conversations { + break; + } + + let msgs = entry["conversations"].as_array().unwrap(); + let mut turns = Vec::new(); + + // Walk alternating human/gpt pairs, stopping early once max_turns reached + // to avoid tokenizing turns that would be discarded by truncate(). + let mut i = 0; + while i + 1 < msgs.len() { + if let Some(m) = max_turns + && turns.len() >= m + { + break; + } + let from = msgs[i].get("from").and_then(|f| f.as_str()).unwrap_or(""); + let user_text = msgs[i].get("value").and_then(|v| v.as_str()).unwrap_or(""); + let assistant_text = msgs[i + 1].get("value").and_then(|v| v.as_str()).unwrap_or(""); + + // Expect human then gpt + if from != "human" || user_text.is_empty() { + i += 1; + continue; + } + + let user_ids = tokenizer.encode(user_text, false)?; + let user_len = user_ids.len(); + + let expected_output_len = if let Some(override_len) = output_len_override { + override_len + } else { + let assistant_ids = tokenizer.encode(assistant_text, false)?; + assistant_ids.len().max(1) + }; + + turns.push(ConversationTurn { + user_message: Arc::from(user_text), + user_message_len: user_len, + expected_output_len, + }); + + i += 2; + } + + if turns.len() >= 2 { + let conv_idx = conversations.len(); + conversations.push(MultiTurnConversation { + conversation_id: format!("{request_id_prefix}conv-{conv_idx}"), + turns, + }); + } + } + + if conversations.is_empty() { + return Err(BenchError::Config( + "No valid multi-turn conversations after filtering ShareGPT dataset.".into(), + )); + } + + // Oversample if needed + if conversations.len() < num_conversations { + let original_len = conversations.len(); + let needed = num_conversations - original_len; + for i in 0..needed { + let mut conv = conversations[rng.random_range(0..original_len)].clone(); + conv.conversation_id = format!("{request_id_prefix}conv-{}", original_len + i); + conversations.push(conv); + } + println!( + "Oversampled multi-turn conversations from {original_len} to {} total.", + conversations.len() + ); + } + + Ok(conversations) +} + +/// Ensure decoded-then-encoded prompt length matches the target. +fn gen_prompt_to_target_len( + tokenizer: &TokenizerKind, + token_sequence: &[u32], + target_len: usize, +) -> Result<(String, Vec)> { + let max_retry = 20; + let mut tokens = token_sequence.to_vec(); + + for retry in 0..=max_retry { + let prompt = tokenizer.decode(&tokens, true)?; + tokens = tokenizer.encode(&prompt, false)?; + + if retry >= max_retry { + // BPE tokenizers can oscillate by ±1 on certain boundaries. + // For benchmark random content, accept close-enough and truncate/pad. + if tokens.len() > target_len { + tokens.truncate(target_len); + } + // If still short by 1-2 tokens, accept as-is — negligible for benchmarks. + // Re-decode after truncation to ensure prompt string matches token vector. + let prompt = tokenizer.decode(&tokens, true)?; + return Ok((prompt, tokens)); + } + + if tokens.len() == target_len { + return Ok((prompt, tokens)); + } else if tokens.len() < target_len { + let allowed = tokenizer.get_allowed_tokens(); + let needed = target_len - tokens.len(); + if allowed.is_empty() { + let vocab_size = tokenizer.vocab_size() as usize; + for j in 0..needed { + tokens.push(((tokens.len() + j) % vocab_size) as u32); + } + } else { + for j in 0..needed { + tokens.push(allowed[(tokens.len() + j) % allowed.len()]); + } + } + } else { + tokens.truncate(target_len); + } + } + + unreachable!() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn common_prefix_bytes(strings: &[&str]) -> usize { + if strings.is_empty() { + return 0; + } + let first = strings[0].as_bytes(); + let mut len = first.len(); + for s in &strings[1..] { + let b = s.as_bytes(); + len = len.min(b.len()); + for i in 0..len { + if first[i] != b[i] { + len = i; + break; + } + } + } + len + } + + #[test] + #[ignore] + fn test_prefix_sharing_structure() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + + let cfg = MultiTurnRandomConfig { + num_conversations: 5, + min_turns: 3, + max_turns: 3, + prefix_len: 0, + input_len: 1000, + per_turn_input_len: 0, + output_len: 100, + seed: 42, + request_id_prefix: "test-".to_string(), + prefix_sharing_config: Some(PrefixSharingConfig { + global_ratio: 0.1, + conversation_ratio: 0.8, + }), + }; + + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap(); + assert_eq!(conversations.len(), 5); + + let messages: Vec> = conversations + .iter() + .map(|c| c.turns.iter().map(|t| &*t.user_message).collect()) + .collect(); + + // 1. Global prefix: all messages share a common prefix + let all_msgs: Vec<&str> = messages.iter().flat_map(|v| v.iter().copied()).collect(); + let global_prefix = common_prefix_bytes(&all_msgs); + println!("Global prefix bytes: {global_prefix}"); + assert!(global_prefix > 0, "Global prefix must be non-empty"); + + // 2. Conversation prefix: turns within same conversation share more + for (i, conv_msgs) in messages.iter().enumerate() { + let conv_prefix = common_prefix_bytes(conv_msgs); + println!("Conv {i} prefix bytes: {conv_prefix} (global: {global_prefix})"); + assert!( + conv_prefix > global_prefix, + "Conv prefix ({conv_prefix}) must exceed global prefix ({global_prefix})" + ); + } + + // 3. Different conversations diverge after global prefix + let cross = common_prefix_bytes(&[messages[0][0], messages[1][0]]); + let within = common_prefix_bytes(&messages[0]); + println!("Cross-conv prefix: {cross}, within-conv prefix: {within}"); + assert!( + cross < within, + "Cross-conv ({cross}) must be < within-conv ({within})" + ); + + // 4. Turns within same conversation are not identical (unique suffix) + for (i, conv_msgs) in messages.iter().enumerate() { + for a in 0..conv_msgs.len() { + for b in (a + 1)..conv_msgs.len() { + assert_ne!( + conv_msgs[a], conv_msgs[b], + "Conv {i} turn {a} and {b} must differ" + ); + } + } + } + + // 5. Token lengths approximately match target + for (i, conv) in conversations.iter().enumerate() { + for (j, turn) in conv.turns.iter().enumerate() { + let diff = (turn.user_message_len as i64 - 1000).abs(); + println!( + "Conv {i} turn {j}: {} tokens (diff {diff})", + turn.user_message_len + ); + assert!( + diff <= 10, + "Token len {} too far from 1000", + turn.user_message_len + ); + } + } + + println!("All prefix sharing checks passed!"); + } + + #[test] + #[ignore] + fn test_per_turn_input_len_default_mode() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + + let cfg = MultiTurnRandomConfig { + num_conversations: 4, + min_turns: 3, + max_turns: 3, + prefix_len: 0, + input_len: 512, + per_turn_input_len: 128, + output_len: 64, + seed: 1, + request_id_prefix: "test-".to_string(), + prefix_sharing_config: None, + }; + + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap(); + assert_eq!(conversations.len(), 4); + + for (i, conv) in conversations.iter().enumerate() { + assert_eq!(conv.turns.len(), 3); + for (j, turn) in conv.turns.iter().enumerate() { + let expected = if j == 0 { 512usize } else { 128usize }; + let diff = (turn.user_message_len as i64 - expected as i64).abs(); + println!( + "Conv {i} turn {j}: {} tokens (expected ~{expected}, diff {diff})", + turn.user_message_len + ); + assert!( + diff <= 5, + "Conv {i} turn {j}: token len {} too far from {expected}", + turn.user_message_len + ); + } + } + println!("per_turn_input_len default-mode checks passed!"); + } + + #[test] + #[ignore] + fn test_variable_turns_range() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + + let cfg = MultiTurnRandomConfig { + num_conversations: 50, + min_turns: 2, + max_turns: 5, + prefix_len: 0, + input_len: 256, + per_turn_input_len: 0, + output_len: 32, + seed: 7, + request_id_prefix: "test-".to_string(), + prefix_sharing_config: None, + }; + + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap(); + assert_eq!(conversations.len(), 50); + + let mut distinct_counts = std::collections::HashSet::new(); + for conv in &conversations { + let n = conv.turns.len(); + assert!((2..=5).contains(&n), "turn count {n} out of [2,5]"); + distinct_counts.insert(n); + } + assert!( + distinct_counts.len() >= 2, + "expected at least 2 distinct turn counts, got {distinct_counts:?}" + ); + println!("variable_turns_range checks passed! counts: {distinct_counts:?}"); + } + + #[test] + #[ignore] + fn test_variable_turns_fixed() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + + let cfg = MultiTurnRandomConfig { + num_conversations: 10, + min_turns: 4, + max_turns: 4, + prefix_len: 0, + input_len: 256, + per_turn_input_len: 0, + output_len: 32, + seed: 42, + request_id_prefix: "test-".to_string(), + prefix_sharing_config: None, + }; + + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap(); + for conv in &conversations { + assert_eq!(conv.turns.len(), 4, "expected exactly 4 turns"); + } + println!("variable_turns_fixed checks passed!"); + } + + #[test] + #[ignore] + fn test_per_turn_input_len_prefix_sharing() { + let tok = crate::tokenizer::load_tokenizer("nvidia/Kimi-K2.5-NVFP4", false, None).unwrap(); + + // Turn 0 input_len=1000, turns 1+ per_turn_input_len=600 + // global_len ≈ 100 (10%), conv_len ≈ 800 (80%), unique ≈ 100 + // per-turn unique ≈ 600 - 900 = negative → would error; use smaller ratios + // global=0.05 (50), conv=0.5 (500), unique_t0=450, unique_t1=600-550=50 + let cfg = MultiTurnRandomConfig { + num_conversations: 4, + min_turns: 3, + max_turns: 3, + prefix_len: 0, + input_len: 1000, + per_turn_input_len: 600, + output_len: 64, + seed: 3, + request_id_prefix: "test-".to_string(), + prefix_sharing_config: Some(PrefixSharingConfig { + global_ratio: 0.05, + conversation_ratio: 0.50, + }), + }; + + let conversations = generate_multi_turn_random(&tok, &cfg).unwrap(); + assert_eq!(conversations.len(), 4); + + let messages: Vec> = conversations + .iter() + .map(|c| c.turns.iter().map(|t| &*t.user_message).collect()) + .collect(); + + // Global prefix bytes shared across all turns of all conversations + let all_msgs: Vec<&str> = messages.iter().flat_map(|v| v.iter().copied()).collect(); + let global_prefix = common_prefix_bytes(&all_msgs); + assert!(global_prefix > 0, "Global prefix must be non-empty"); + + // Within each conversation, prefix grows (conv prefix longer than global) + for (i, conv_msgs) in messages.iter().enumerate() { + let conv_prefix = common_prefix_bytes(conv_msgs); + assert!( + conv_prefix > global_prefix, + "Conv {i}: conv_prefix ({conv_prefix}) must exceed global ({global_prefix})" + ); + } + + // Turn 0 length ≈ 1000, turns 1+ ≈ 600 + for (i, conv) in conversations.iter().enumerate() { + for (j, turn) in conv.turns.iter().enumerate() { + let expected = if j == 0 { 1000usize } else { 600usize }; + let diff = (turn.user_message_len as i64 - expected as i64).abs(); + println!( + "Conv {i} turn {j}: {} tokens (expected ~{expected})", + turn.user_message_len + ); + assert!( + diff <= 10, + "Conv {i} turn {j}: token len {} too far from {expected}", + turn.user_message_len + ); + } + } + println!("per_turn_input_len prefix-sharing checks passed!"); + } +} diff --git a/rust/src/bench/src/datasets/prefix_repetition.rs b/rust/src/bench/src/datasets/prefix_repetition.rs new file mode 100644 index 000000000000..7b06ef6077d1 --- /dev/null +++ b/rust/src/bench/src/datasets/prefix_repetition.rs @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Prefix repetition dataset: N distinct shared prefixes, each reused by +//! `num_prompts / num_prefixes` requests with a fresh random suffix. +//! The standard prefix-cache stress workload; mirrors Python's +//! `PrefixRepetitionRandomDataset`. + +use std::sync::Arc; + +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; + +use super::SampleRequest; +use super::random::gen_prompt_decode_to_target_len; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +/// Generate the prefix repetition dataset. +/// +/// Like Python, `num_requests % num_prefixes` remainder requests are dropped: +/// the total is `(num_requests / num_prefixes) * num_prefixes`. +pub fn generate_prefix_repetition_dataset( + tokenizer: &TokenizerKind, + num_requests: usize, + prefix_len: usize, + suffix_len: usize, + num_prefixes: usize, + output_len: usize, + seed: u64, + request_id_prefix: &str, + disable_shuffle: bool, +) -> Result> { + let prompts_per_prefix = num_requests / num_prefixes; + if prompts_per_prefix == 0 { + return Err(BenchError::Config(format!( + "num_prompts ({num_requests}) must be >= num_prefixes ({num_prefixes})" + ))); + } + let total = prompts_per_prefix * num_prefixes; + if total != num_requests { + println!( + "prefix_repetition: generating {total} requests \ + ({num_prefixes} prefixes x {prompts_per_prefix} prompts each; \ + {} dropped to divide evenly)", + num_requests - total + ); + } + + let allowed_tokens = tokenizer.get_allowed_tokens(); + if allowed_tokens.is_empty() { + return Err(BenchError::Tokenizer("No allowed tokens found".into())); + } + let allowed_ref = &allowed_tokens; + + // Exact-length random token block: decode -> re-encode -> converge to target. + let gen_block = |target_len: usize, item_seed: u64| -> Result> { + let mut rng = StdRng::seed_from_u64(item_seed); + let tokens: Vec = (0..target_len) + .map(|_| allowed_ref[rng.random_range(0..allowed_ref.len())]) + .collect(); + let (_, adjusted) = + gen_prompt_decode_to_target_len(tokenizer, &tokens, target_len, false, allowed_ref)?; + Ok(adjusted) + }; + + // Generate the shared prefixes (one per group), then suffixes in parallel. + let prefixes: Vec> = (0..num_prefixes) + .map(|p| gen_block(prefix_len, seed.wrapping_add(0xF1F0).wrapping_add(p as u64))) + .collect::>>()?; + + let rid_prefix = request_id_prefix.to_string(); + let mut requests: Vec = (0..total) + .into_par_iter() + .map(|i| { + let prefix_tokens = &prefixes[i / prompts_per_prefix]; + let suffix_tokens = gen_block(suffix_len, seed.wrapping_add(0xBEEF + i as u64))?; + + let mut combined = Vec::with_capacity(prefix_tokens.len() + suffix_tokens.len()); + combined.extend_from_slice(prefix_tokens); + combined.extend_from_slice(&suffix_tokens); + let prompt = tokenizer.decode(&combined, true)?; + + Ok(SampleRequest { + prompt: Arc::from(prompt), + prompt_len: combined.len(), + expected_output_len: output_len, + request_id: Some(format!("{rid_prefix}{i}")), + ..Default::default() + }) + }) + .collect::>>()?; + + // Interleave prefixes (Python shuffles too) so one prefix group isn't sent + // as a contiguous burst. + if !disable_shuffle { + let mut rng = StdRng::seed_from_u64(seed); + requests.shuffle(&mut rng); + } + + Ok(requests) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// gpt2 via built-in tiktoken encoding — loads without network access. + fn test_tokenizer() -> TokenizerKind { + crate::tokenizer::load_tokenizer("gpt2", false, None) + .expect("gpt2 built-in tiktoken should always load without network") + } + + #[test] + fn test_prefix_repetition_structure() { + let tok = test_tokenizer(); + // 7 requests / 3 prefixes -> 2 per prefix, 6 total (remainder dropped like Python) + let reqs = generate_prefix_repetition_dataset(&tok, 7, 32, 16, 3, 64, 0, "t-", true) + .expect("generation should succeed"); + assert_eq!(reqs.len(), 6); + assert!(reqs.iter().all(|r| r.expected_output_len == 64)); + // Exact-length blocks: prompt_len == prefix + suffix + assert!( + reqs.iter().all(|r| r.prompt_len == 32 + 16), + "lens: {:?}", + reqs.iter().map(|r| r.prompt_len).collect::>() + ); + // Consecutive pairs (shuffle disabled) share a common prefix; requests + // from different groups don't. + let common = |a: &str, b: &str| -> usize { + a.bytes().zip(b.bytes()).take_while(|(x, y)| x == y).count() + }; + let same_group = common(&reqs[0].prompt, &reqs[1].prompt); + let diff_group = common(&reqs[0].prompt, &reqs[2].prompt); + assert!( + same_group > diff_group, + "same-group shared prefix ({same_group}) should exceed cross-group ({diff_group})" + ); + } + + #[test] + fn test_prefix_repetition_too_few_requests_errors() { + let tok = test_tokenizer(); + assert!(generate_prefix_repetition_dataset(&tok, 2, 32, 16, 3, 64, 0, "t-", true).is_err()); + } +} diff --git a/rust/src/bench/src/datasets/random.rs b/rust/src/bench/src/datasets/random.rs new file mode 100644 index 000000000000..7cf5311278f5 --- /dev/null +++ b/rust/src/bench/src/datasets/random.rs @@ -0,0 +1,491 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use rayon::prelude::*; + +use super::SampleRequest; +use crate::config::RangeRatio; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +/// Generate random dataset with rayon parallelism. +/// +/// This is the key performance win — Python does sequential tokenizer calls +/// while Rust parallelizes across CPU cores with native tokenizer speed. +/// +/// Mirrors Python's RandomDataset.sample() from datasets.py:470-560. +pub fn generate_random_dataset( + tokenizer: &TokenizerKind, + num_requests: usize, + input_len: usize, + output_len: usize, + prefix_len: usize, + range_ratio: RangeRatio, + cache_hit_fraction: f64, + cache_ratio: f64, + seed: u64, + request_id_prefix: &str, + use_token_ids: bool, + batch_size: usize, +) -> Result> { + let vocab_size = tokenizer.vocab_size(); + let allowed_tokens = tokenizer.get_allowed_tokens(); + if allowed_tokens.is_empty() { + return Err(BenchError::Tokenizer("No allowed tokens found".into())); + } + + if batch_size > 1 && use_token_ids { + return Err(BenchError::Config( + "--random-batch-size > 1 is not supported with --prompt-token-ids".into(), + )); + } + + let num_special = tokenizer.num_special_tokens_to_add(); + let real_input_len = input_len.saturating_sub(num_special); + + // Python semantics: sample uniformly from [len*(1-r), len*(1+r)]. + let (input_low, input_high) = range_ratio.input_bounds(real_input_len); + let (output_low, output_high) = range_ratio.output_bounds(output_len); + if !range_ratio.is_fixed() { + println!( + "Sampling input_len from [{input_low}, {input_high}] and \ + output_len from [{output_low}, {output_high}]" + ); + } + + // Bimodal prefix-cache mode: a fraction of prompts (warm) reuse a shared cached + // prefix covering `cache_ratio` of their length; the rest (cold) are fully unique. + // Models e.g. "80% of prompts have 95% of input cached" with + // --random-cache-hit-fraction 0.8 --random-cache-ratio 0.95. In this mode + // --random-input-len is the TOTAL prompt length L (the cached prefix is part of L), + // and --random-prefix-len is ignored. + let bimodal = cache_hit_fraction > 0.0 && cache_ratio > 0.0; + if bimodal { + if !use_token_ids { + return Err(BenchError::Config( + "bimodal prefix-cache (--random-cache-hit-fraction) requires --prompt-token-ids \ + so warm prompts send identical token IDs and actually hit the prefix cache" + .into(), + )); + } + if cache_hit_fraction > 1.0 || cache_ratio > 1.0 { + return Err(BenchError::Config( + "--random-cache-hit-fraction and --random-cache-ratio must be in [0, 1]".into(), + )); + } + } + + // Length of the shared cached base prefix. + let base_len = if bimodal { + ((input_high as f64) * cache_ratio).ceil() as usize + } else { + prefix_len + }; + + // Validate (non-bimodal keeps the original check) + if !bimodal { + let min_total = prefix_len + input_low; + if min_total < 1 { + return Err(BenchError::Config(format!( + "--random-input-len too small: with {num_special} special tokens and \ + range_ratio={:?}, minimum total input is {min_total}", + range_ratio + ))); + } + } + + // Generate the shared base prefix once (sequential, only happens once). + let prefix_token_ids = if base_len > 0 { + generate_prefix(tokenizer, &allowed_tokens, base_len, seed)? + } else { + Vec::new() + }; + + // Pre-generate per-request sampling params using deterministic RNG + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + let mut rng = StdRng::seed_from_u64(seed); + + struct RequestParams { + cached_len: usize, // tokens taken from the shared base (cache-hittable) + suffix_len: usize, // unique tokens appended after the cached prefix + output_len: usize, + offset: usize, + } + + let params: Vec = (0..num_requests) + .map(|_| { + let ol = if output_low == output_high { + output_low + } else { + rng.random_range(output_low..=output_high) + }; + let off = rng.random_range(0..vocab_size as usize); + if bimodal { + // Total length L from the input distribution; prefix is part of L. + let l = if input_low == input_high { + input_low + } else { + rng.random_range(input_low..=input_high) + }; + let warm = rng.random::() < cache_hit_fraction; + let cached = if warm { + (((l as f64) * cache_ratio).round() as usize).min(base_len).min(l) + } else { + 0 + }; + RequestParams { + cached_len: cached, + suffix_len: l - cached, + output_len: ol, + offset: off, + } + } else { + // Original behavior: full shared prefix + variable unique input. + let il = if input_low == input_high { + input_low + } else { + rng.random_range(input_low..=input_high) + }; + RequestParams { + cached_len: prefix_len, + suffix_len: il, + output_len: ol, + offset: off, + } + } + }) + .collect(); + + // Phase 1: Generate all token sequences (parallel, fast — just array ops) + let prefix_ref = &prefix_token_ids; + let allowed_ref = &allowed_tokens; + let rid_prefix = request_id_prefix.to_string(); + + let token_sequences: Vec> = params + .par_iter() + .enumerate() + .map(|(i, p)| { + let at_len = allowed_ref.len(); + let mut seq = Vec::with_capacity(p.cached_len + p.suffix_len); + seq.extend_from_slice(&prefix_ref[..p.cached_len]); + for j in 0..p.suffix_len { + seq.push(allowed_ref[(p.offset + i + j) % at_len]); + } + seq + }) + .collect(); + + let target_lens: Vec = params.iter().map(|p| p.cached_len + p.suffix_len).collect(); + + if use_token_ids { + // Fast path: store token IDs directly. The completions backend sends + // them as `"prompt": [id1, id2, ...]`, bypassing both client-side decode + // and server-side tokenization. Token counts are exact by construction. + let result: Vec = token_sequences + .into_par_iter() + .enumerate() + .map(|(i, tokens)| SampleRequest { + prompt: Arc::from(""), + prompt_len: target_lens[i], + expected_output_len: params[i].output_len, + request_id: Some(format!("{rid_prefix}{i}")), + prompt_token_ids: Some(Arc::from(tokens)), + ..Default::default() + }) + .collect(); + Ok(result) + } else { + // Default path: decode tokens to text, re-encode, + // truncate to target length, decode again. Sends text prompts for maximum + let result: Vec = token_sequences + .into_par_iter() + .enumerate() + .map(|(i, tokens)| { + let target = target_lens[i]; + // decode → encode → truncate → decode + let prompt_text = tokenizer.decode(&tokens, true)?; + let mut re_encoded = tokenizer.encode(&prompt_text, false)?; + re_encoded.truncate(target); + let prompt = tokenizer.decode(&re_encoded, true)?; + let prompt_len = re_encoded.len(); + Ok(SampleRequest { + prompt: Arc::from(prompt), + prompt_len, + expected_output_len: params[i].output_len, + request_id: Some(format!("{rid_prefix}{i}")), + ..Default::default() + }) + }) + .collect::>>()?; + Ok(super::batch_requests(result, batch_size, request_id_prefix)) + } +} + +fn generate_prefix( + tokenizer: &TokenizerKind, + allowed_tokens: &[u32], + prefix_len: usize, + seed: u64, +) -> Result> { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + let mut rng = StdRng::seed_from_u64(seed.wrapping_add(0xDEAD)); + let tokens: Vec = (0..prefix_len) + .map(|_| allowed_tokens[rng.random_range(0..allowed_tokens.len())]) + .collect(); + + let (_, adjusted) = + gen_prompt_decode_to_target_len(tokenizer, &tokens, prefix_len, false, allowed_tokens)?; + Ok(adjusted) +} + +/// Ensure decoded-then-encoded prompt length matches the target. +/// +/// Mirrors Python's `gen_prompt_decode_to_target_len` from datasets.py:381-435. +pub(crate) fn gen_prompt_decode_to_target_len( + tokenizer: &TokenizerKind, + token_sequence: &[u32], + target_len: usize, + add_special_tokens: bool, + allowed_tokens: &[u32], +) -> Result<(String, Vec)> { + let max_retry = 20; + let mut tokens = token_sequence.to_vec(); + + for retry in 0..=max_retry { + let prompt = tokenizer.decode(&tokens, true)?; + tokens = tokenizer.encode(&prompt, add_special_tokens)?; + + if retry >= max_retry { + if tokens.len() != target_len { + return Err(BenchError::Tokenizer(format!( + "Token length mismatch after {max_retry} retries: \ + target={target_len}, actual={}. \ + encode/decode roundtrip cannot converge.", + tokens.len() + ))); + } + return Ok((prompt, tokens)); + } + + if tokens.len() == target_len { + return Ok((prompt, tokens)); + } else if tokens.len() < target_len { + // Pad with tokens from the allowed set (UTF-8-safe for tiktoken) + let needed = target_len - tokens.len(); + if allowed_tokens.is_empty() { + let vocab_size = tokenizer.vocab_size() as usize; + for j in 0..needed { + tokens.push(((tokens.len() + j) % vocab_size) as u32); + } + } else { + for j in 0..needed { + tokens.push(allowed_tokens[(tokens.len() + j) % allowed_tokens.len()]); + } + } + } else { + // Truncate + tokens.truncate(target_len); + } + } + + unreachable!() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tokenizer; + + // Integration test requires a tokenizer, so only run with --ignored + #[test] + #[ignore] + fn test_generate_random_dataset_token_ids() { + let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let requests = generate_random_dataset( + &tokenizer, + 10, // num_requests + 128, // input_len + 32, // output_len + 0, // prefix_len + RangeRatio { + input: 0.0, + output: 0.0, + }, // range_ratio (0.0 = fixed length) + 0.0, // cache_hit_fraction (0 = bimodal off) + 0.0, // cache_ratio + 42, // seed + "test-", + true, // use_token_ids + 1, // batch_size + ) + .unwrap(); + + assert_eq!(requests.len(), 10); + for req in &requests { + assert!(req.prompt_token_ids.is_some()); + assert_eq!(req.prompt_token_ids.as_ref().unwrap().len(), req.prompt_len); + assert!(req.prompt_len > 0); + assert_eq!(req.expected_output_len, 32); + } + } + + #[test] + #[ignore] + fn test_generate_random_dataset_text() { + let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let requests = generate_random_dataset( + &tokenizer, + 10, // num_requests + 128, // input_len + 32, // output_len + 0, // prefix_len + RangeRatio { + input: 0.0, + output: 0.0, + }, // range_ratio (0.0 = fixed length) + 0.0, // cache_hit_fraction (0 = bimodal off) + 0.0, // cache_ratio + 42, // seed + "test-", + false, // use_token_ids = false → text prompts + 1, // batch_size + ) + .unwrap(); + + assert_eq!(requests.len(), 10); + for req in &requests { + assert!(req.prompt_token_ids.is_none()); + assert!(!req.prompt.is_empty()); + assert!(req.prompt_len > 0); + assert!(req.prompt_len <= 128); + assert_eq!(req.expected_output_len, 32); + } + } + + /// Test that generated prompts have EXACT target token length (token ID mode). + #[test] + #[ignore] + fn test_token_length_exact_local() { + let tokenizer = tokenizer::load_tokenizer("gpt2", false, None).unwrap(); + let target_len = 512; + let requests = generate_random_dataset( + &tokenizer, + 50, + target_len, + 64, + 0, + RangeRatio { + input: 0.0, + output: 0.0, + }, + 0.0, + 0.0, + 123, + "len-test-", + true, + 1, + ) + .unwrap(); + + for (i, req) in requests.iter().enumerate() { + let token_ids = req.prompt_token_ids.as_ref().expect("should have token IDs"); + assert_eq!( + token_ids.len(), + target_len, + "Request {i}: expected {target_len} token IDs, got {}", + token_ids.len() + ); + assert_eq!(req.prompt_len, target_len); + } + } + + /// Test that tiktoken tokenizer produces exact target token lengths (token ID mode). + #[test] + #[ignore] + fn test_token_length_exact_tiktoken() { + // Use Qwen2.5 which has a tiktoken-format tokenizer + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None); + let tokenizer = match tokenizer { + Ok(t) => t, + Err(e) => { + eprintln!("Skipping tiktoken test (tokenizer unavailable): {e}"); + return; + } + }; + + // Verify it's actually a tiktoken tokenizer or local — either way test convergence + let target_len = 256; + let requests = generate_random_dataset( + &tokenizer, + 20, + target_len, + 32, + 0, + RangeRatio { + input: 0.0, + output: 0.0, + }, + 0.0, + 0.0, + 42, + "tiktoken-test-", + true, + 1, + ) + .unwrap(); + + for (i, req) in requests.iter().enumerate() { + let token_ids = req.prompt_token_ids.as_ref().expect("should have token IDs"); + assert_eq!( + token_ids.len(), + target_len, + "Request {i}: expected {target_len} token IDs, got {}", + token_ids.len() + ); + assert_eq!(req.prompt_len, target_len); + } + } + + /// Test encode/decode roundtrip stability for tiktoken. + /// After one decode→encode cycle with UTF-8-safe tokens, length must not drift. + #[test] + #[ignore] + fn test_tiktoken_roundtrip_stability() { + let tokenizer = tokenizer::load_tokenizer("Qwen/Qwen2.5-0.5B", false, None); + let tokenizer = match tokenizer { + Ok(t) => t, + Err(e) => { + eprintln!("Skipping roundtrip test (tokenizer unavailable): {e}"); + return; + } + }; + + let allowed = tokenizer.get_allowed_tokens(); + assert!(!allowed.is_empty(), "allowed tokens should not be empty"); + + // Build a sequence from allowed tokens only + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + let mut rng = StdRng::seed_from_u64(99); + let seq: Vec = (0..512).map(|_| allowed[rng.random_range(0..allowed.len())]).collect(); + + let decoded = tokenizer.decode(&seq, true).unwrap(); + let re_encoded = tokenizer.encode(&decoded, false).unwrap(); + let re_decoded = tokenizer.decode(&re_encoded, true).unwrap(); + let re_re_encoded = tokenizer.encode(&re_decoded, false).unwrap(); + + // After first cycle, length should stabilize + assert_eq!( + re_encoded.len(), + re_re_encoded.len(), + "Roundtrip should stabilize: first re-encode={}, second re-encode={}", + re_encoded.len(), + re_re_encoded.len() + ); + } +} diff --git a/rust/src/bench/src/datasets/random_mm.rs b/rust/src/bench/src/datasets/random_mm.rs new file mode 100644 index 000000000000..22bb4334f716 --- /dev/null +++ b/rust/src/bench/src/datasets/random_mm.rs @@ -0,0 +1,657 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::io::Cursor; +use std::sync::Arc; + +use base64::Engine as _; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; + +use super::SampleRequest; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +/// A bucket key: (height, width, num_frames). num_frames=1 means image, >1 means video. +#[derive(Debug, Clone)] +pub struct MmBucketKey { + pub height: u32, + pub width: u32, + pub num_frames: u32, +} + +/// Per-modality hard caps. +#[derive(Debug, Clone)] +pub struct MmLimitPerPrompt { + pub image: usize, + pub video: usize, +} + +impl Default for MmLimitPerPrompt { + fn default() -> Self { + Self { + image: 255, + video: 1, + } + } +} + +/// Parse the limit-mm-per-prompt JSON string, e.g. `{"image": 3, "video": 0}`. +pub fn parse_limit_mm_per_prompt(s: &str) -> Result { + let v: serde_json::Value = serde_json::from_str(s) + .map_err(|e| BenchError::Config(format!("Invalid --random-mm-limit-mm-per-prompt: {e}")))?; + let obj = v.as_object().ok_or_else(|| { + BenchError::Config("--random-mm-limit-mm-per-prompt must be a JSON object".into()) + })?; + let image = obj.get("image").and_then(|v| v.as_u64()).unwrap_or(255) as usize; + let video = obj.get("video").and_then(|v| v.as_u64()).unwrap_or(1) as usize; + Ok(MmLimitPerPrompt { image, video }) +} + +/// Parse the bucket config string in Python-style syntax. +/// +/// Accepts: `{(256,256,1): 0.5, (720,1280,1): 0.5}` +/// Each key is `(height, width, num_frames)` and value is the probability weight. +pub fn parse_bucket_config(s: &str) -> Result> { + let trimmed = s.trim(); + let inner = trimmed + .strip_prefix('{') + .and_then(|s| s.strip_suffix('}')) + .ok_or_else(|| BenchError::Config("Bucket config must be wrapped in {}".into()))?; + + let mut buckets = Vec::new(); + let chars: Vec = inner.chars().collect(); + let len = chars.len(); + let mut i = 0; + + while i < len { + // Skip whitespace and commas + while i < len && (chars[i].is_whitespace() || chars[i] == ',') { + i += 1; + } + if i >= len { + break; + } + + // Expect '(' + if chars[i] != '(' { + return Err(BenchError::Config(format!( + "Expected '(' in bucket config at position {i}" + ))); + } + i += 1; + + // Read until ')' + let tuple_start = i; + while i < len && chars[i] != ')' { + i += 1; + } + if i >= len { + return Err(BenchError::Config("Unclosed '(' in bucket config".into())); + } + let tuple_str: String = chars[tuple_start..i].iter().collect(); + i += 1; // skip ')' + + // Skip whitespace, then expect ':' + while i < len && chars[i].is_whitespace() { + i += 1; + } + if i >= len || chars[i] != ':' { + return Err(BenchError::Config( + "Expected ':' after tuple in bucket config".into(), + )); + } + i += 1; + + // Skip whitespace + while i < len && chars[i].is_whitespace() { + i += 1; + } + + // Read the probability value until ',' or end + let val_start = i; + while i < len && chars[i] != ',' { + i += 1; + } + let val_str: String = chars[val_start..i].iter().collect(); + + // Parse tuple + let parts: Vec<&str> = tuple_str.split(',').collect(); + if parts.len() != 3 { + return Err(BenchError::Config(format!( + "Bucket key must have 3 values (height,width,num_frames), got: ({tuple_str})" + ))); + } + + let height: u32 = parts[0].trim().parse().map_err(|_| { + BenchError::Config(format!( + "Invalid height in bucket config: '{}'", + parts[0].trim() + )) + })?; + let width: u32 = parts[1].trim().parse().map_err(|_| { + BenchError::Config(format!( + "Invalid width in bucket config: '{}'", + parts[1].trim() + )) + })?; + let num_frames: u32 = parts[2].trim().parse().map_err(|_| { + BenchError::Config(format!( + "Invalid num_frames in bucket config: '{}'", + parts[2].trim() + )) + })?; + let prob: f64 = val_str.trim().parse().map_err(|_| { + BenchError::Config(format!( + "Invalid probability in bucket config: '{}'", + val_str.trim() + )) + })?; + + if prob < 0.0 { + return Err(BenchError::Config(format!( + "Bucket probability must be non-negative, got: {prob}" + ))); + } + + buckets.push(( + MmBucketKey { + height, + width, + num_frames, + }, + prob, + )); + } + + if buckets.is_empty() { + return Err(BenchError::Config( + "Bucket config must have at least one entry".into(), + )); + } + + Ok(buckets) +} + +/// JSON fragment prefix/suffix for image content blocks. +const IMG_JSON_PREFIX: &str = r#"{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"#; +const IMG_JSON_SUFFIX: &str = r#""}}"#; + +/// Generate a synthetic random JPEG image and return it as a pre-serialized JSON fragment. +/// +/// Builds the complete JSON string in a single allocation: +/// `{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"}}` +/// +/// The base64 data is written directly into the final string — no intermediate +/// String or format!() copy. +fn generate_random_image(width: u32, height: u32, rng: &mut StdRng) -> Result> { + let pixel_count = (width as usize) * (height as usize) * 3; + let mut pixels = vec![0u8; pixel_count]; + rng.fill(pixels.as_mut_slice()); + + let img = image::RgbImage::from_raw(width, height, pixels) + .ok_or_else(|| BenchError::Config("Failed to create image from random pixels".into()))?; + + // Pre-allocate JPEG buffer (random pixels compress poorly, estimate ~60% of raw) + let estimated_jpeg = pixel_count * 3 / 5; + let mut buf = Cursor::new(Vec::with_capacity(estimated_jpeg)); + img.write_to(&mut buf, image::ImageFormat::Jpeg) + .map_err(|e| BenchError::Config(format!("Failed to encode JPEG: {e}")))?; + + let jpeg_bytes = buf.into_inner(); + + // Pre-compute exact output size: prefix + base64_len + suffix + let b64_len = jpeg_bytes.len().div_ceil(3) * 4; + let total_len = IMG_JSON_PREFIX.len() + b64_len + IMG_JSON_SUFFIX.len(); + + // Single allocation: write base64 directly into the JSON fragment string + let mut json_fragment = String::with_capacity(total_len); + json_fragment.push_str(IMG_JSON_PREFIX); + base64::engine::general_purpose::STANDARD.encode_string(&jpeg_bytes, &mut json_fragment); + json_fragment.push_str(IMG_JSON_SUFFIX); + + Ok(Arc::from(json_fragment)) +} + +/// Sample multimodal items for a single request. +/// +/// Returns a list of (height, width, num_frames) tuples. +fn sample_mm_items( + rng: &mut StdRng, + min_items: usize, + max_items: usize, + buckets: &[(MmBucketKey, f64)], + limit: &MmLimitPerPrompt, +) -> Vec { + let num_items = if min_items == max_items { + min_items + } else { + rng.random_range(min_items..=max_items) + }; + + // Filter to non-zero probability buckets + let active_buckets: Vec<&(MmBucketKey, f64)> = + buckets.iter().filter(|(_, p)| *p > 0.0).collect(); + if active_buckets.is_empty() || num_items == 0 { + return Vec::new(); + } + + let total_weight: f64 = active_buckets.iter().map(|(_, p)| p).sum(); + if total_weight <= 0.0 { + return Vec::new(); + } + + let mut result = Vec::with_capacity(num_items); + let mut image_count = 0usize; + let mut video_count = 0usize; + + for _ in 0..num_items { + // Build normalized weights considering remaining capacity + let mut weights: Vec = Vec::with_capacity(active_buckets.len()); + for (key, prob) in &active_buckets { + let is_video = key.num_frames > 1; + let at_limit = if is_video { + video_count >= limit.video + } else { + image_count >= limit.image + }; + weights.push(if at_limit { 0.0 } else { *prob }); + } + + let w_total: f64 = weights.iter().sum(); + if w_total <= 0.0 { + break; // All modalities at limit + } + + // Weighted random selection (strict `<` to avoid selecting zero-weight buckets) + let r = rng.random::() * w_total; + let mut cumulative = 0.0; + // Default to last non-zero-weight bucket (floating-point accumulation fallback) + let mut selected_idx = weights.iter().rposition(|w| *w > 0.0).unwrap_or(0); + for (i, w) in weights.iter().enumerate() { + cumulative += w; + if r < cumulative { + selected_idx = i; + break; + } + } + + let (key, _) = &active_buckets[selected_idx]; + if key.num_frames > 1 { + video_count += 1; + } else { + image_count += 1; + } + result.push(key.clone()); + } + + result +} + +/// Generate random multimodal dataset. +/// +/// Mirrors Python's RandomMultiModalDataset.sample() from datasets.py. +/// Generates text prompts with exact token lengths and random images/videos. +pub fn generate_random_mm_dataset( + tokenizer: &TokenizerKind, + num_requests: usize, + input_len: usize, + output_len: usize, + prefix_len: usize, + range_ratio: crate::config::RangeRatio, + seed: u64, + request_id_prefix: &str, + base_items_per_request: usize, + num_mm_items_range_ratio: f64, + limit: &MmLimitPerPrompt, + buckets: &[(MmBucketKey, f64)], + enable_multimodal_chat: bool, +) -> Result> { + if !(0.0..=1.0).contains(&num_mm_items_range_ratio) { + return Err(BenchError::Config( + "num_mm_items_range_ratio must be in [0, 1]".into(), + )); + } + + // Check for video buckets with non-zero probability + for (key, prob) in buckets { + if key.num_frames > 1 && *prob > 0.0 { + return Err(BenchError::Config( + "Video generation (num_frames > 1) is not yet supported in Rust. \ + Set video bucket probabilities to 0.0." + .into(), + )); + } + } + + // Compute item count bounds + let n = base_items_per_request as f64; + let r = num_mm_items_range_ratio; + let min_items = (n * (1.0 - r)).floor().max(0.0) as usize; + let max_items = (n * (1.0 + r)).ceil() as usize; + // Clamp to total modality limit + let total_limit = limit.image + limit.video; + let max_items = max_items.min(total_limit); + let min_items = min_items.min(max_items); + + let vocab_size = tokenizer.vocab_size(); + let allowed_tokens = tokenizer.get_allowed_tokens(); + if allowed_tokens.is_empty() { + return Err(BenchError::Tokenizer("No allowed tokens found".into())); + } + + let num_special = tokenizer.num_special_tokens_to_add(); + let real_input_len = input_len.saturating_sub(num_special); + + // Python semantics: sample uniformly from [len*(1-r), len*(1+r)]. + let (input_low, input_high) = range_ratio.input_bounds(real_input_len); + let (output_low, output_high) = range_ratio.output_bounds(output_len); + + // Pre-generate per-request params + let mut rng = StdRng::seed_from_u64(seed); + + struct RequestParams { + input_len: usize, + output_len: usize, + offset: usize, + } + + let params: Vec = (0..num_requests) + .map(|_| { + let il = if input_low == input_high { + input_low + } else { + rng.random_range(input_low..=input_high) + }; + let ol = if output_low == output_high { + output_low + } else { + rng.random_range(output_low..=output_high) + }; + let off = rng.random_range(0..vocab_size as usize); + RequestParams { + input_len: il, + output_len: ol, + offset: off, + } + }) + .collect(); + + // Pre-generate multimodal item configs per request + let mm_configs: Vec> = (0..num_requests) + .map(|_| sample_mm_items(&mut rng, min_items, max_items, buckets, limit)) + .collect(); + + // Generate text prompts (need text for chat backend, not just token IDs) + let prefix_token_ids = if prefix_len > 0 { + generate_prefix(tokenizer, &allowed_tokens, prefix_len, seed)? + } else { + Vec::new() + }; + + // Generate token sequences + let prefix_ref = &prefix_token_ids; + let allowed_ref = &allowed_tokens; + + let token_sequences: Vec> = params + .par_iter() + .enumerate() + .map(|(i, p)| { + let at_len = allowed_ref.len(); + let mut seq = Vec::with_capacity(prefix_ref.len() + p.input_len); + seq.extend_from_slice(prefix_ref); + for j in 0..p.input_len { + seq.push(allowed_ref[(p.offset + i + j) % at_len]); + } + seq + }) + .collect(); + + let target_lens: Vec = params.iter().map(|p| prefix_len + p.input_len).collect(); + + // Decode tokens to text (chat backend needs text prompts for multimodal) + let prompts: Result> = token_sequences + .into_par_iter() + .enumerate() + .map(|(i, tokens)| { + let (text, _adjusted) = super::random::gen_prompt_decode_to_target_len( + tokenizer, + &tokens, + target_lens[i], + false, + allowed_ref, + )?; + Ok(text) + }) + .collect(); + let prompts = prompts?; + + // Generate images for each request (parallel per request) + // Each request gets its own RNG seeded deterministically. + let rid_prefix = request_id_prefix.to_string(); + let result: Vec = prompts + .into_par_iter() + .enumerate() + .map(|(i, prompt)| { + let mut item_rng = + StdRng::seed_from_u64(seed.wrapping_add(i as u64).wrapping_add(0xBEEF)); + let mm_items: Vec> = mm_configs[i] + .iter() + .map(|key| { + generate_random_image(key.width, key.height, &mut item_rng) + .expect("Image generation should not fail") + }) + .collect(); + + let mm_content: Option]>> = if mm_items.is_empty() { + None + } else { + Some(Arc::from(mm_items)) + }; + + // --enable-multimodal-chat: pre-build the full chat `messages` array + // (text part + mm items) at dataset time, mirroring Python's + // apply_multimodal_chat_transformation. mm content moves inside the + // messages string; the backend splices it verbatim. + let (mm_content, chat_messages_json) = if enable_multimodal_chat { + let msgs = build_chat_messages_json(&prompt, mm_content.as_deref()); + (None, Some(Arc::from(msgs.as_str()))) + } else { + (mm_content, None) + }; + + SampleRequest { + prompt: Arc::from(prompt.as_str()), + prompt_len: target_lens[i], + expected_output_len: params[i].output_len, + request_id: Some(format!("{rid_prefix}{i}")), + multi_modal_content: mm_content, + chat_messages_json, + ..Default::default() + } + }) + .collect(); + + Ok(result) +} + +/// Pre-serialize the OpenAI chat `messages` array for --enable-multimodal-chat. +/// +/// Produces `[{"role":"user","content":[{"type":"text","text":"..."},,...]}]` +/// by concatenating the JSON-escaped prompt with the pre-serialized mm fragments, +/// so the ~200KB+ base64 image data is never parsed or re-serialized. +pub(crate) fn build_chat_messages_json(prompt: &str, mm_items: Option<&[Arc]>) -> String { + let mm_total: usize = + mm_items.map(|items| items.iter().map(|f| f.len() + 1).sum()).unwrap_or(0); + let mut msgs = String::with_capacity(64 + prompt.len() * 2 + mm_total); + msgs.push_str(r#"[{"role":"user","content":[{"type":"text","text":"#); + // serde_json::to_string on &str produces a JSON-escaped quoted string + msgs.push_str(&serde_json::to_string(prompt).unwrap()); + msgs.push('}'); + for fragment in mm_items.unwrap_or(&[]) { + msgs.push(','); + msgs.push_str(fragment); + } + msgs.push_str("]}]"); + msgs +} + +fn generate_prefix( + tokenizer: &TokenizerKind, + allowed_tokens: &[u32], + prefix_len: usize, + seed: u64, +) -> Result> { + let mut rng = StdRng::seed_from_u64(seed.wrapping_add(0xDEAD)); + let tokens: Vec = (0..prefix_len) + .map(|_| allowed_tokens[rng.random_range(0..allowed_tokens.len())]) + .collect(); + + let (_, adjusted) = super::random::gen_prompt_decode_to_target_len( + tokenizer, + &tokens, + prefix_len, + false, + allowed_tokens, + )?; + Ok(adjusted) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_bucket_config_basic() { + let input = "{(256,256,1): 0.5, (720,1280,1): 0.5}"; + let buckets = parse_bucket_config(input).unwrap(); + assert_eq!(buckets.len(), 2); + assert_eq!(buckets[0].0.height, 256); + assert_eq!(buckets[0].0.width, 256); + assert_eq!(buckets[0].0.num_frames, 1); + assert!((buckets[0].1 - 0.5).abs() < 1e-10); + assert_eq!(buckets[1].0.height, 720); + assert_eq!(buckets[1].0.width, 1280); + } + + #[test] + fn test_parse_bucket_config_single() { + let input = "{(1024, 800, 1): 1.0}"; + let buckets = parse_bucket_config(input).unwrap(); + assert_eq!(buckets.len(), 1); + assert_eq!(buckets[0].0.height, 1024); + assert_eq!(buckets[0].0.width, 800); + assert_eq!(buckets[0].0.num_frames, 1); + assert!((buckets[0].1 - 1.0).abs() < 1e-10); + } + + #[test] + fn test_parse_bucket_config_with_video() { + let input = "{(256,256,1): 0.4, (720,1280,1): 0.4, (720,1280,16): 0.2}"; + let buckets = parse_bucket_config(input).unwrap(); + assert_eq!(buckets.len(), 3); + assert_eq!(buckets[2].0.num_frames, 16); + } + + #[test] + fn test_build_chat_messages_json_valid_and_ordered() { + let frag: Arc = + Arc::from(r#"{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,AAAA"}}"#); + let msgs = build_chat_messages_json("hi \"there\"\nline2", Some(&[frag])); + let v: serde_json::Value = serde_json::from_str(&msgs).expect("must be valid JSON"); + assert_eq!(v.as_array().unwrap().len(), 1); + assert_eq!(v[0]["role"], "user"); + let content = v[0]["content"].as_array().unwrap(); + assert_eq!(content.len(), 2); + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[0]["text"], "hi \"there\"\nline2"); + assert_eq!(content[1]["type"], "image_url"); + } + + #[test] + fn test_build_chat_messages_json_text_only() { + let msgs = build_chat_messages_json("plain", None); + let v: serde_json::Value = serde_json::from_str(&msgs).unwrap(); + let content = v[0]["content"].as_array().unwrap(); + assert_eq!(content.len(), 1); + assert_eq!(content[0]["text"], "plain"); + } + + #[test] + fn test_parse_limit_mm_per_prompt() { + let input = r#"{"image": 3, "video": 0}"#; + let limit = parse_limit_mm_per_prompt(input).unwrap(); + assert_eq!(limit.image, 3); + assert_eq!(limit.video, 0); + } + + #[test] + fn test_parse_limit_mm_per_prompt_defaults() { + let input = r#"{}"#; + let limit = parse_limit_mm_per_prompt(input).unwrap(); + assert_eq!(limit.image, 255); + assert_eq!(limit.video, 1); + } + + #[test] + fn test_sample_mm_items_basic() { + let mut rng = StdRng::seed_from_u64(42); + let buckets = vec![ + ( + MmBucketKey { + height: 256, + width: 256, + num_frames: 1, + }, + 0.5, + ), + ( + MmBucketKey { + height: 720, + width: 1280, + num_frames: 1, + }, + 0.5, + ), + ]; + let limit = MmLimitPerPrompt { image: 5, video: 0 }; + let items = sample_mm_items(&mut rng, 2, 3, &buckets, &limit); + assert!(items.len() >= 2 && items.len() <= 3); + for item in &items { + assert_eq!(item.num_frames, 1); + } + } + + #[test] + fn test_sample_mm_items_respects_limit() { + let mut rng = StdRng::seed_from_u64(42); + let buckets = vec![( + MmBucketKey { + height: 256, + width: 256, + num_frames: 1, + }, + 1.0, + )]; + let limit = MmLimitPerPrompt { image: 2, video: 0 }; + let items = sample_mm_items(&mut rng, 5, 5, &buckets, &limit); + // Should be capped at 2 due to image limit + assert_eq!(items.len(), 2); + } + + #[test] + fn test_generate_random_image() { + let mut rng = StdRng::seed_from_u64(42); + let result = generate_random_image(64, 64, &mut rng).unwrap(); + // Result is a pre-serialized JSON fragment + assert!( + result + .starts_with(r#"{"type":"image_url","image_url":{"url":"data:image/jpeg;base64,"#) + ); + assert!(result.ends_with(r#""}}"#)); + // Verify it's valid JSON + let parsed: serde_json::Value = serde_json::from_str(&result).unwrap(); + assert_eq!(parsed["type"], "image_url"); + } +} diff --git a/rust/src/bench/src/datasets/random_rerank.rs b/rust/src/bench/src/datasets/random_rerank.rs new file mode 100644 index 000000000000..9c335e7e44f0 --- /dev/null +++ b/rust/src/bench/src/datasets/random_rerank.rs @@ -0,0 +1,185 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Random dataset specialized for scoring/rerank benchmarks: each request is +//! one query plus a batch of documents. Mirrors Python's +//! `RandomDatasetForReranking`. +//! +//! With `is_reranker` (default): the query and each document share the +//! request's token budget (`query + sep + doc ~= input_len`), and every +//! batched request counts the query once per document pair. +//! With `--no-reranker` (embedding-based scoring): the query is just another +//! embedding input occupying the first batch slot. + +use std::sync::Arc; + +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use rayon::prelude::*; + +use super::SampleRequest; +use crate::config::RangeRatio; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +pub fn generate_random_rerank_dataset( + tokenizer: &TokenizerKind, + num_requests: usize, + input_len: usize, + range_ratio: RangeRatio, + seed: u64, + request_id_prefix: &str, + batch_size: usize, + is_reranker: bool, +) -> Result> { + let allowed_tokens = tokenizer.get_allowed_tokens(); + if allowed_tokens.is_empty() { + return Err(BenchError::Tokenizer("No allowed tokens found".into())); + } + let allowed_ref = &allowed_tokens; + + let num_special = tokenizer.num_special_tokens_to_add(); + let real_input_len = input_len.saturating_sub(num_special); + + let n_sep_tokens = usize::from(is_reranker); + let query_len_param = if is_reranker { + (real_input_len / 2).saturating_sub(n_sep_tokens) + } else { + real_input_len + }; + + let mut rng = StdRng::seed_from_u64(seed); + let sample = |rng: &mut StdRng, (low, high): (usize, usize)| -> usize { + if low == high { + low + } else { + rng.random_range(low..=high) + } + }; + + // One query length for the whole run, like Python. + let query_len = sample(&mut rng, range_ratio.input_bounds(query_len_param)); + + // --no-reranker folds the query into the first batch slot. + let (num_docs, docs_per_batch, doc_len_param) = if is_reranker { + let doc_len = real_input_len.saturating_sub(query_len).saturating_sub(n_sep_tokens); + (num_requests, batch_size, doc_len) + } else { + (num_requests - 1, batch_size - 1, real_input_len) + }; + if doc_len_param == 0 { + return Err(BenchError::Config(format!( + "random-rerank: --random-input-len {input_len} leaves no budget for documents \ + (query_len={query_len})" + ))); + } + + // Pre-sample per-document lengths and offsets deterministically. + let doc_bounds = range_ratio.input_bounds(doc_len_param); + let doc_params: Vec<(usize, usize)> = (0..num_docs) + .map(|_| { + ( + sample(&mut rng, doc_bounds), + rng.random_range(0..allowed_ref.len()), + ) + }) + .collect(); + let query_offset = rng.random_range(0..allowed_ref.len()); + + // Exact-length text: token sequence -> decode -> re-encode -> truncate -> decode. + let gen_text = |target: usize, offset: usize, index: usize| -> Result<(Arc, usize)> { + let at_len = allowed_ref.len(); + let tokens: Vec = + (0..target).map(|j| allowed_ref[(offset + index + j) % at_len]).collect(); + let text = tokenizer.decode(&tokens, true)?; + let mut re_encoded = tokenizer.encode(&text, false)?; + re_encoded.truncate(target); + let final_text = tokenizer.decode(&re_encoded, true)?; + Ok((Arc::from(final_text), re_encoded.len())) + }; + + let (query_prompt, query_input_len) = gen_text(query_len, query_offset, 0)?; + + let docs: Vec<(Arc, usize)> = doc_params + .par_iter() + .enumerate() + .map(|(i, (len, offset))| gen_text(*len, *offset, i + 1)) + .collect::>>()?; + + // Batch documents; every request is [query, doc1, doc2, ...]. + let rid_prefix = request_id_prefix.to_string(); + let requests = docs + .chunks(docs_per_batch) + .enumerate() + .map(|(batch_idx, batch)| { + let query_contrib = if is_reranker { + (query_input_len + n_sep_tokens) * batch.len() + } else { + query_input_len + }; + let mut prompt_list: Vec> = Vec::with_capacity(batch.len() + 1); + prompt_list.push(query_prompt.clone()); + prompt_list.extend(batch.iter().map(|(text, _)| text.clone())); + SampleRequest { + prompt_list: Some(Arc::from(prompt_list)), + prompt_len: query_contrib + batch.iter().map(|(_, len)| len).sum::(), + expected_output_len: 0, + request_id: Some(format!("{rid_prefix}{batch_idx}")), + ..Default::default() + } + }) + .collect(); + + Ok(requests) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// gpt2 via built-in tiktoken encoding — loads without network access. + fn test_tokenizer() -> TokenizerKind { + crate::tokenizer::load_tokenizer("gpt2", false, None) + .expect("gpt2 built-in tiktoken should always load without network") + } + + fn fixed_ratio() -> RangeRatio { + RangeRatio::parse("0.0").unwrap() + } + + #[test] + fn test_random_rerank_reranker_mode() { + let tok = test_tokenizer(); + // 6 docs in batches of 3 -> 2 requests of [query, d1, d2, d3] + let reqs = generate_random_rerank_dataset(&tok, 6, 128, fixed_ratio(), 0, "t-", 3, true) + .expect("generation should succeed"); + assert_eq!(reqs.len(), 2); + for r in &reqs { + let list = r.prompt_list.as_ref().expect("prompt_list must be set"); + assert_eq!(list.len(), 4); + assert_eq!(r.expected_output_len, 0); + assert!(r.prompt_len > 0); + } + // Same query shared across requests + assert_eq!( + reqs[0].prompt_list.as_ref().unwrap()[0], + reqs[1].prompt_list.as_ref().unwrap()[0] + ); + // Reranker budget: query+sep+doc pairs stay near input_len per pair + // (query ~63, doc ~64 for input_len=128, gpt2 has no special tokens) + let list = reqs[0].prompt_list.as_ref().unwrap(); + let query_tokens = tok.encode(&list[0], false).unwrap().len(); + assert!(query_tokens <= 64, "query too long: {query_tokens}"); + } + + #[test] + fn test_random_rerank_no_reranker_mode() { + let tok = test_tokenizer(); + // no-reranker: query occupies first slot; 5 non-query docs in batches of 2 + let reqs = generate_random_rerank_dataset(&tok, 6, 64, fixed_ratio(), 0, "t-", 3, false) + .expect("generation should succeed"); + // 6-1=5 docs, batches of 3-1=2 -> 3 requests + assert_eq!(reqs.len(), 3); + assert_eq!(reqs[0].prompt_list.as_ref().unwrap().len(), 3); + } +} diff --git a/rust/src/bench/src/datasets/sharegpt.rs b/rust/src/bench/src/datasets/sharegpt.rs new file mode 100644 index 000000000000..2d3901820df6 --- /dev/null +++ b/rust/src/bench/src/datasets/sharegpt.rs @@ -0,0 +1,209 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; + +use super::SampleRequest; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +/// Default validation bounds matching Python's is_valid_sequence() defaults. +const MIN_LEN: usize = 4; +const MAX_PROMPT_LEN: usize = 1024; +const MAX_TOTAL_LEN: usize = 2048; + +/// Default HuggingFace dataset repo and filename for ShareGPT. +const DEFAULT_SHAREGPT_REPO: &str = "anon8231489123/ShareGPT_Vicuna_unfiltered"; +const DEFAULT_SHAREGPT_FILE: &str = "ShareGPT_V3_unfiltered_cleaned_split.json"; + +/// Download the default ShareGPT dataset from HuggingFace Hub. +/// Uses hf-hub's built-in cache — subsequent calls return the cached path instantly. +pub fn download_sharegpt_dataset() -> Result { + println!( + "Downloading ShareGPT dataset from {DEFAULT_SHAREGPT_REPO}/{DEFAULT_SHAREGPT_FILE} ..." + ); + let repo = crate::hub::HubRepo::dataset(DEFAULT_SHAREGPT_REPO.to_string()); + let path = repo.get(DEFAULT_SHAREGPT_FILE).map_err(|e| { + BenchError::Config(format!( + "Failed to download ShareGPT dataset from '{DEFAULT_SHAREGPT_REPO}': {e}" + )) + })?; + let path_str = path.to_string_lossy().to_string(); + println!("ShareGPT dataset ready: {path_str}"); + Ok(path_str) +} + +/// Load and sample from a ShareGPT-format JSON dataset. +/// +/// Mirrors Python's ShareGPTDataset from datasets.py:1230-1313. +pub fn load_sharegpt_dataset( + tokenizer: &TokenizerKind, + dataset_path: &str, + num_requests: usize, + output_len_override: Option, + seed: u64, + request_id_prefix: &str, + no_oversample: bool, + disable_shuffle: bool, +) -> Result> { + // Load JSON file + let content = std::fs::read_to_string(dataset_path).map_err(|e| { + BenchError::Config(format!( + "Failed to read ShareGPT file '{dataset_path}': {e}" + )) + })?; + + let data: serde_json::Value = serde_json::from_str(&content) + .map_err(|e| BenchError::Config(format!("Invalid JSON in ShareGPT file: {e}")))?; + + let entries = data + .as_array() + .ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?; + + // Filter entries with at least 2 conversation turns + let mut filtered: Vec<&serde_json::Value> = entries + .iter() + .filter(|entry| { + entry + .get("conversations") + .and_then(|c| c.as_array()) + .map(|a| a.len() >= 2) + .unwrap_or(false) + }) + .collect(); + + if filtered.is_empty() { + return Err(BenchError::Config( + "No valid entries in ShareGPT file (need at least 2 conversation turns)".into(), + )); + } + + // Shuffle (unless disabled) + let mut rng = StdRng::seed_from_u64(seed); + if !disable_shuffle { + filtered.shuffle(&mut rng); + } + + // Sample requests + let mut samples = Vec::new(); + let mut ind = 0; + + for entry in &filtered { + if samples.len() >= num_requests { + break; + } + + let conversations = entry["conversations"].as_array().unwrap(); + let prompt = conversations[0]["value"].as_str().unwrap_or(""); + let completion = conversations[1]["value"].as_str().unwrap_or(""); + + if prompt.is_empty() { + continue; + } + + // Tokenize prompt and completion + let prompt_ids = tokenizer.encode(prompt, false)?; + let prompt_len = prompt_ids.len(); + + let new_output_len = if let Some(override_len) = output_len_override { + override_len + } else { + let completion_ids = tokenizer.encode(completion, false)?; + completion_ids.len() + }; + + // Validate sequence lengths (matching Python's is_valid_sequence) + let skip_min_output = output_len_override.is_some(); + if !is_valid_sequence(prompt_len, new_output_len, skip_min_output) { + continue; + } + + samples.push(SampleRequest { + prompt: Arc::from(prompt), + prompt_len, + expected_output_len: new_output_len, + request_id: Some(format!("{request_id_prefix}{ind}")), + ..Default::default() + }); + ind += 1; + } + + // Oversample if dataset is smaller than requested + if samples.len() < num_requests { + if no_oversample { + println!( + "Skipping oversampling. Total samples: {} (requested: {num_requests})", + samples.len() + ); + } else if !samples.is_empty() { + let needed = num_requests - samples.len(); + let original_len = samples.len(); + for i in 0..needed { + let mut req = samples[rng.random_range(0..original_len)].clone(); + req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); + samples.push(req); + } + println!( + "Oversampled requests from {original_len} to {} total samples.", + samples.len() + ); + } + } + + if samples.is_empty() { + return Err(BenchError::Config( + "No valid samples after filtering ShareGPT dataset. \ + Try relaxing constraints or using a larger dataset." + .into(), + )); + } + + Ok(samples) +} + +/// Validate a sequence based on prompt and output lengths. +/// Mirrors Python's is_valid_sequence() from datasets.py:260-284. +fn is_valid_sequence( + prompt_len: usize, + output_len: usize, + skip_min_output_len_check: bool, +) -> bool { + if prompt_len < MIN_LEN { + return false; + } + if !skip_min_output_len_check && output_len < MIN_LEN { + return false; + } + if prompt_len > MAX_PROMPT_LEN { + return false; + } + if prompt_len + output_len > MAX_TOTAL_LEN { + return false; + } + true +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_is_valid_sequence() { + // Valid + assert!(is_valid_sequence(100, 50, false)); + // Prompt too short + assert!(!is_valid_sequence(3, 50, false)); + // Output too short + assert!(!is_valid_sequence(100, 3, false)); + // Output too short but skip check + assert!(is_valid_sequence(100, 1, true)); + // Prompt too long + assert!(!is_valid_sequence(1025, 50, false)); + // Combined too long + assert!(!is_valid_sequence(1024, 1025, false)); + } +} diff --git a/rust/src/bench/src/datasets/sonnet.rs b/rust/src/bench/src/datasets/sonnet.rs new file mode 100644 index 000000000000..6fc7a75c0b5b --- /dev/null +++ b/rust/src/bench/src/datasets/sonnet.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand::seq::IndexedRandom; + +use super::SampleRequest; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +/// Default values mirror Python's SonnetDataset defaults (datasets.py). +pub const DEFAULT_PREFIX_LEN: usize = 200; +pub const DEFAULT_INPUT_LEN: usize = 550; +pub const DEFAULT_OUTPUT_LEN: usize = 150; + +const BASE_PROMPT: &str = "Pick as many lines as you can from these poem lines:\n"; + +/// Shakespeare's sonnets, public domain. Bundled so `--dataset-name sonnet` works +/// out of the box without `--dataset-path`. Source: +/// https://raw.githubusercontent.com/vllm-project/vllm/main/benchmarks/sonnet.txt +const BUILTIN_SONNET: &str = include_str!("sonnet.txt"); + +/// Load the sonnet dataset and generate `num_requests` prompts targeting `input_len` +/// total prompt tokens. +/// +/// Mirrors Python's `SonnetDataset.sample()` from `vllm/benchmarks/datasets/datasets.py`. +/// The Rust port skips `apply_chat_template` (no Jinja runtime here): `base_offset` +/// and `prompt_len` are computed from the raw text. The resulting prompts are slightly +/// shorter than Python's chat-template-formatted version (off by the chat scaffolding +/// tokens, typically <20). +pub fn load_sonnet_dataset( + tokenizer: &TokenizerKind, + dataset_path: Option<&str>, + num_requests: usize, + input_len: usize, + output_len: usize, + prefix_len: usize, + seed: u64, + request_id_prefix: &str, +) -> Result> { + let content = match dataset_path { + Some(path) => std::fs::read_to_string(path) + .map_err(|e| BenchError::Config(format!("Failed to read sonnet file '{path}': {e}")))?, + None => BUILTIN_SONNET.to_string(), + }; + + // Match Python's f.readlines(): keep trailing newlines so that joining lines + // reconstructs the original text without inserting extra separators. + let lines: Vec = content.split_inclusive('\n').map(|s| s.to_string()).collect(); + + if lines.is_empty() { + let src = dataset_path.unwrap_or(""); + return Err(BenchError::Config(format!("Sonnet file '{src}' is empty"))); + } + + // Average tokens per line (used to estimate how many lines to draw). + let mut total_tokens: usize = 0; + for line in &lines { + let ids = tokenizer.encode(line, false)?; + total_tokens += ids.len(); + } + let avg_len = total_tokens as f64 / lines.len() as f64; + if avg_len <= 0.0 { + return Err(BenchError::Config( + "Sonnet lines tokenized to zero tokens on average".into(), + )); + } + + let base_ids = tokenizer.encode(BASE_PROMPT, false)?; + let base_offset = base_ids.len(); + if input_len <= base_offset { + return Err(BenchError::Config(format!( + "--sonnet-input-len ({input_len}) must be larger than the base prompt length ({base_offset})" + ))); + } + + let num_input_lines = ((input_len - base_offset) as f64 / avg_len).round() as i64; + let num_prefix_lines = + (((prefix_len as i64 - base_offset as i64) as f64) / avg_len).round().max(0.0) as i64; + let num_input_lines = num_input_lines.max(0) as usize; + let num_prefix_lines = (num_prefix_lines as usize).min(lines.len()); + let num_input_lines = num_input_lines.max(num_prefix_lines); + + let prefix_lines: &[String] = &lines[..num_prefix_lines]; + let extras_per_request = num_input_lines - num_prefix_lines; + + let mut rng = StdRng::seed_from_u64(seed); + let mut samples = Vec::with_capacity(num_requests); + let mut ind = 0usize; + let mut attempts = 0usize; + let max_attempts = num_requests.saturating_mul(20).max(1000); + + while samples.len() < num_requests { + if attempts >= max_attempts { + return Err(BenchError::Config(format!( + "Could not assemble {num_requests} sonnet prompts under input_len={input_len} \ + after {attempts} attempts. Try increasing --sonnet-input-len." + ))); + } + attempts += 1; + + let mut prompt = String::with_capacity(BASE_PROMPT.len() + 256 * num_input_lines); + prompt.push_str(BASE_PROMPT); + for line in prefix_lines { + prompt.push_str(line); + } + for _ in 0..extras_per_request { + // random.choices with replacement — duplicates are allowed. + let line = lines.choose(&mut rng).unwrap(); + prompt.push_str(line); + } + + let prompt_ids = tokenizer.encode(&prompt, false)?; + let prompt_len = prompt_ids.len(); + if prompt_len <= input_len { + samples.push(SampleRequest { + prompt: Arc::from(prompt), + prompt_len, + expected_output_len: output_len, + request_id: Some(format!("{request_id_prefix}{ind}")), + ..Default::default() + }); + ind += 1; + } + } + + Ok(samples) +} diff --git a/rust/src/bench/src/datasets/sonnet.txt b/rust/src/bench/src/datasets/sonnet.txt new file mode 100644 index 000000000000..34c444e8ce8e --- /dev/null +++ b/rust/src/bench/src/datasets/sonnet.txt @@ -0,0 +1,518 @@ +FROM fairest creatures we desire increase, +That thereby beauty's rose might never die, +But as the riper should by time decease, +His tender heir might bear his memory: +But thou, contracted to thine own bright eyes, +Feed'st thy light'st flame with self-substantial fuel, +Making a famine where abundance lies, +Thyself thy foe, to thy sweet self too cruel. +Thou that art now the world's fresh ornament +And only herald to the gaudy spring, +Within thine own bud buriest thy content +And, tender churl, makest waste in niggarding. +Pity the world, or else this glutton be, +To eat the world's due, by the grave and thee. +When forty winters shall beseige thy brow, +And dig deep trenches in thy beauty's field, +Thy youth's proud livery, so gazed on now, +Will be a tatter'd weed, of small worth held: +Then being ask'd where all thy beauty lies, +Where all the treasure of thy lusty days, +To say, within thine own deep-sunken eyes, +Were an all-eating shame and thriftless praise. +How much more praise deserved thy beauty's use, +If thou couldst answer 'This fair child of mine +Shall sum my count and make my old excuse,' +Proving his beauty by succession thine! +This were to be new made when thou art old, +And see thy blood warm when thou feel'st it cold. +Look in thy glass, and tell the face thou viewest +Now is the time that face should form another; +Whose fresh repair if now thou not renewest, +Thou dost beguile the world, unbless some mother. +For where is she so fair whose unear'd womb +Disdains the tillage of thy husbandry? +Or who is he so fond will be the tomb +Of his self-love, to stop posterity? +Thou art thy mother's glass, and she in thee +Calls back the lovely April of her prime: +So thou through windows of thine age shall see +Despite of wrinkles this thy golden time. +But if thou live, remember'd not to be, +Die single, and thine image dies with thee. +Unthrifty loveliness, why dost thou spend +Upon thyself thy beauty's legacy? +Nature's bequest gives nothing but doth lend, +And being frank she lends to those are free. +Then, beauteous niggard, why dost thou abuse +The bounteous largess given thee to give? +Profitless usurer, why dost thou use +So great a sum of sums, yet canst not live? +For having traffic with thyself alone, +Thou of thyself thy sweet self dost deceive. +Then how, when nature calls thee to be gone, +What acceptable audit canst thou leave? +Thy unused beauty must be tomb'd with thee, +Which, used, lives th' executor to be. +Those hours, that with gentle work did frame +The lovely gaze where every eye doth dwell, +Will play the tyrants to the very same +And that unfair which fairly doth excel: +For never-resting time leads summer on +To hideous winter and confounds him there; +Sap cheque'd with frost and lusty leaves quite gone, +Beauty o'ersnow'd and bareness every where: +Then, were not summer's distillation left, +A liquid prisoner pent in walls of glass, +Beauty's effect with beauty were bereft, +Nor it nor no remembrance what it was: +But flowers distill'd though they with winter meet, +Leese but their show; their substance still lives sweet. +Then let not winter's ragged hand deface +In thee thy summer, ere thou be distill'd: +Make sweet some vial; treasure thou some place +With beauty's treasure, ere it be self-kill'd. +That use is not forbidden usury, +Which happies those that pay the willing loan; +That's for thyself to breed another thee, +Or ten times happier, be it ten for one; +Ten times thyself were happier than thou art, +If ten of thine ten times refigured thee: +Then what could death do, if thou shouldst depart, +Leaving thee living in posterity? +Be not self-will'd, for thou art much too fair +To be death's conquest and make worms thine heir. +Lo! in the orient when the gracious light +Lifts up his burning head, each under eye +Doth homage to his new-appearing sight, +Serving with looks his sacred majesty; +And having climb'd the steep-up heavenly hill, +Resembling strong youth in his middle age, +yet mortal looks adore his beauty still, +Attending on his golden pilgrimage; +But when from highmost pitch, with weary car, +Like feeble age, he reeleth from the day, +The eyes, 'fore duteous, now converted are +From his low tract and look another way: +So thou, thyself out-going in thy noon, +Unlook'd on diest, unless thou get a son. +Music to hear, why hear'st thou music sadly? +Sweets with sweets war not, joy delights in joy. +Why lovest thou that which thou receivest not gladly, +Or else receivest with pleasure thine annoy? +If the true concord of well-tuned sounds, +By unions married, do offend thine ear, +They do but sweetly chide thee, who confounds +In singleness the parts that thou shouldst bear. +Mark how one string, sweet husband to another, +Strikes each in each by mutual ordering, +Resembling sire and child and happy mother +Who all in one, one pleasing note do sing: +Whose speechless song, being many, seeming one, +Sings this to thee: 'thou single wilt prove none.' +Is it for fear to wet a widow's eye +That thou consumest thyself in single life? +Ah! if thou issueless shalt hap to die. +The world will wail thee, like a makeless wife; +The world will be thy widow and still weep +That thou no form of thee hast left behind, +When every private widow well may keep +By children's eyes her husband's shape in mind. +Look, what an unthrift in the world doth spend +Shifts but his place, for still the world enjoys it; +But beauty's waste hath in the world an end, +And kept unused, the user so destroys it. +No love toward others in that bosom sits +That on himself such murderous shame commits. +For shame! deny that thou bear'st love to any, +Who for thyself art so unprovident. +Grant, if thou wilt, thou art beloved of many, +But that thou none lovest is most evident; +For thou art so possess'd with murderous hate +That 'gainst thyself thou stick'st not to conspire. +Seeking that beauteous roof to ruinate +Which to repair should be thy chief desire. +O, change thy thought, that I may change my mind! +Shall hate be fairer lodged than gentle love? +Be, as thy presence is, gracious and kind, +Or to thyself at least kind-hearted prove: +Make thee another self, for love of me, +That beauty still may live in thine or thee. +As fast as thou shalt wane, so fast thou growest +In one of thine, from that which thou departest; +And that fresh blood which youngly thou bestowest +Thou mayst call thine when thou from youth convertest. +Herein lives wisdom, beauty and increase: +Without this, folly, age and cold decay: +If all were minded so, the times should cease +And threescore year would make the world away. +Let those whom Nature hath not made for store, +Harsh featureless and rude, barrenly perish: +Look, whom she best endow'd she gave the more; +Which bounteous gift thou shouldst in bounty cherish: +She carved thee for her seal, and meant thereby +Thou shouldst print more, not let that copy die. +When I do count the clock that tells the time, +And see the brave day sunk in hideous night; +When I behold the violet past prime, +And sable curls all silver'd o'er with white; +When lofty trees I see barren of leaves +Which erst from heat did canopy the herd, +And summer's green all girded up in sheaves +Borne on the bier with white and bristly beard, +Then of thy beauty do I question make, +That thou among the wastes of time must go, +Since sweets and beauties do themselves forsake +And die as fast as they see others grow; +And nothing 'gainst Time's scythe can make defence +Save breed, to brave him when he takes thee hence. +O, that you were yourself! but, love, you are +No longer yours than you yourself here live: +Against this coming end you should prepare, +And your sweet semblance to some other give. +So should that beauty which you hold in lease +Find no determination: then you were +Yourself again after yourself's decease, +When your sweet issue your sweet form should bear. +Who lets so fair a house fall to decay, +Which husbandry in honour might uphold +Against the stormy gusts of winter's day +And barren rage of death's eternal cold? +O, none but unthrifts! Dear my love, you know +You had a father: let your son say so. +Not from the stars do I my judgment pluck; +And yet methinks I have astronomy, +But not to tell of good or evil luck, +Of plagues, of dearths, or seasons' quality; +Nor can I fortune to brief minutes tell, +Pointing to each his thunder, rain and wind, +Or say with princes if it shall go well, +By oft predict that I in heaven find: +But from thine eyes my knowledge I derive, +And, constant stars, in them I read such art +As truth and beauty shall together thrive, +If from thyself to store thou wouldst convert; +Or else of thee this I prognosticate: +Thy end is truth's and beauty's doom and date. +When I consider every thing that grows +Holds in perfection but a little moment, +That this huge stage presenteth nought but shows +Whereon the stars in secret influence comment; +When I perceive that men as plants increase, +Cheered and cheque'd even by the self-same sky, +Vaunt in their youthful sap, at height decrease, +And wear their brave state out of memory; +Then the conceit of this inconstant stay +Sets you most rich in youth before my sight, +Where wasteful Time debateth with Decay, +To change your day of youth to sullied night; +And all in war with Time for love of you, +As he takes from you, I engraft you new. +But wherefore do not you a mightier way +Make war upon this bloody tyrant, Time? +And fortify yourself in your decay +With means more blessed than my barren rhyme? +Now stand you on the top of happy hours, +And many maiden gardens yet unset +With virtuous wish would bear your living flowers, +Much liker than your painted counterfeit: +So should the lines of life that life repair, +Which this, Time's pencil, or my pupil pen, +Neither in inward worth nor outward fair, +Can make you live yourself in eyes of men. +To give away yourself keeps yourself still, +And you must live, drawn by your own sweet skill. +Who will believe my verse in time to come, +If it were fill'd with your most high deserts? +Though yet, heaven knows, it is but as a tomb +Which hides your life and shows not half your parts. +If I could write the beauty of your eyes +And in fresh numbers number all your graces, +The age to come would say 'This poet lies: +Such heavenly touches ne'er touch'd earthly faces.' +So should my papers yellow'd with their age +Be scorn'd like old men of less truth than tongue, +And your true rights be term'd a poet's rage +And stretched metre of an antique song: +But were some child of yours alive that time, +You should live twice; in it and in my rhyme. +Shall I compare thee to a summer's day? +Thou art more lovely and more temperate: +Rough winds do shake the darling buds of May, +And summer's lease hath all too short a date: +Sometime too hot the eye of heaven shines, +And often is his gold complexion dimm'd; +And every fair from fair sometime declines, +By chance or nature's changing course untrimm'd; +But thy eternal summer shall not fade +Nor lose possession of that fair thou owest; +Nor shall Death brag thou wander'st in his shade, +When in eternal lines to time thou growest: +So long as men can breathe or eyes can see, +So long lives this and this gives life to thee. +Devouring Time, blunt thou the lion's paws, +And make the earth devour her own sweet brood; +Pluck the keen teeth from the fierce tiger's jaws, +And burn the long-lived phoenix in her blood; +Make glad and sorry seasons as thou fleets, +And do whate'er thou wilt, swift-footed Time, +To the wide world and all her fading sweets; +But I forbid thee one most heinous crime: +O, carve not with thy hours my love's fair brow, +Nor draw no lines there with thine antique pen; +Him in thy course untainted do allow +For beauty's pattern to succeeding men. +Yet, do thy worst, old Time: despite thy wrong, +My love shall in my verse ever live young. +A woman's face with Nature's own hand painted +Hast thou, the master-mistress of my passion; +A woman's gentle heart, but not acquainted +With shifting change, as is false women's fashion; +An eye more bright than theirs, less false in rolling, +Gilding the object whereupon it gazeth; +A man in hue, all 'hues' in his controlling, +Much steals men's eyes and women's souls amazeth. +And for a woman wert thou first created; +Till Nature, as she wrought thee, fell a-doting, +And by addition me of thee defeated, +By adding one thing to my purpose nothing. +But since she prick'd thee out for women's pleasure, +Mine be thy love and thy love's use their treasure. +So is it not with me as with that Muse +Stirr'd by a painted beauty to his verse, +Who heaven itself for ornament doth use +And every fair with his fair doth rehearse +Making a couplement of proud compare, +With sun and moon, with earth and sea's rich gems, +With April's first-born flowers, and all things rare +That heaven's air in this huge rondure hems. +O' let me, true in love, but truly write, +And then believe me, my love is as fair +As any mother's child, though not so bright +As those gold candles fix'd in heaven's air: +Let them say more than like of hearsay well; +I will not praise that purpose not to sell. +My glass shall not persuade me I am old, +So long as youth and thou are of one date; +But when in thee time's furrows I behold, +Then look I death my days should expiate. +For all that beauty that doth cover thee +Is but the seemly raiment of my heart, +Which in thy breast doth live, as thine in me: +How can I then be elder than thou art? +O, therefore, love, be of thyself so wary +As I, not for myself, but for thee will; +Bearing thy heart, which I will keep so chary +As tender nurse her babe from faring ill. +Presume not on thy heart when mine is slain; +Thou gavest me thine, not to give back again. +As an unperfect actor on the stage +Who with his fear is put besides his part, +Or some fierce thing replete with too much rage, +Whose strength's abundance weakens his own heart. +So I, for fear of trust, forget to say +The perfect ceremony of love's rite, +And in mine own love's strength seem to decay, +O'ercharged with burden of mine own love's might. +O, let my books be then the eloquence +And dumb presagers of my speaking breast, +Who plead for love and look for recompense +More than that tongue that more hath more express'd. +O, learn to read what silent love hath writ: +To hear with eyes belongs to love's fine wit. +Mine eye hath play'd the painter and hath stell'd +Thy beauty's form in table of my heart; +My body is the frame wherein 'tis held, +And perspective it is the painter's art. +For through the painter must you see his skill, +To find where your true image pictured lies; +Which in my bosom's shop is hanging still, +That hath his windows glazed with thine eyes. +Now see what good turns eyes for eyes have done: +Mine eyes have drawn thy shape, and thine for me +Are windows to my breast, where-through the sun +Delights to peep, to gaze therein on thee; +Yet eyes this cunning want to grace their art; +They draw but what they see, know not the heart. +Let those who are in favour with their stars +Of public honour and proud titles boast, +Whilst I, whom fortune of such triumph bars, +Unlook'd for joy in that I honour most. +Great princes' favourites their fair leaves spread +But as the marigold at the sun's eye, +And in themselves their pride lies buried, +For at a frown they in their glory die. +The painful warrior famoused for fight, +After a thousand victories once foil'd, +Is from the book of honour razed quite, +And all the rest forgot for which he toil'd: +Then happy I, that love and am beloved +Where I may not remove nor be removed. +Lord of my love, to whom in vassalage +Thy merit hath my duty strongly knit, +To thee I send this written embassage, +To witness duty, not to show my wit: +Duty so great, which wit so poor as mine +May make seem bare, in wanting words to show it, +But that I hope some good conceit of thine +In thy soul's thought, all naked, will bestow it; +Till whatsoever star that guides my moving +Points on me graciously with fair aspect +And puts apparel on my tatter'd loving, +To show me worthy of thy sweet respect: +Then may I dare to boast how I do love thee; +Till then not show my head where thou mayst prove me. +Weary with toil, I haste me to my bed, +The dear repose for limbs with travel tired; +But then begins a journey in my head, +To work my mind, when body's work's expired: +For then my thoughts, from far where I abide, +Intend a zealous pilgrimage to thee, +And keep my drooping eyelids open wide, +Looking on darkness which the blind do see +Save that my soul's imaginary sight +Presents thy shadow to my sightless view, +Which, like a jewel hung in ghastly night, +Makes black night beauteous and her old face new. +Lo! thus, by day my limbs, by night my mind, +For thee and for myself no quiet find. +How can I then return in happy plight, +That am debarr'd the benefit of rest? +When day's oppression is not eased by night, +But day by night, and night by day, oppress'd? +And each, though enemies to either's reign, +Do in consent shake hands to torture me; +The one by toil, the other to complain +How far I toil, still farther off from thee. +I tell the day, to please them thou art bright +And dost him grace when clouds do blot the heaven: +So flatter I the swart-complexion'd night, +When sparkling stars twire not thou gild'st the even. +But day doth daily draw my sorrows longer +And night doth nightly make grief's strength seem stronger. +When, in disgrace with fortune and men's eyes, +I all alone beweep my outcast state +And trouble deal heaven with my bootless cries +And look upon myself and curse my fate, +Wishing me like to one more rich in hope, +Featured like him, like him with friends possess'd, +Desiring this man's art and that man's scope, +With what I most enjoy contented least; +Yet in these thoughts myself almost despising, +Haply I think on thee, and then my state, +Like to the lark at break of day arising +From sullen earth, sings hymns at heaven's gate; +For thy sweet love remember'd such wealth brings +That then I scorn to change my state with kings. +When to the sessions of sweet silent thought +I summon up remembrance of things past, +I sigh the lack of many a thing I sought, +And with old woes new wail my dear time's waste: +Then can I drown an eye, unused to flow, +For precious friends hid in death's dateless night, +And weep afresh love's long since cancell'd woe, +And moan the expense of many a vanish'd sight: +Then can I grieve at grievances foregone, +And heavily from woe to woe tell o'er +The sad account of fore-bemoaned moan, +Which I new pay as if not paid before. +But if the while I think on thee, dear friend, +All losses are restored and sorrows end. +Thy bosom is endeared with all hearts, +Which I by lacking have supposed dead, +And there reigns love and all love's loving parts, +And all those friends which I thought buried. +How many a holy and obsequious tear +Hath dear religious love stol'n from mine eye +As interest of the dead, which now appear +But things removed that hidden in thee lie! +Thou art the grave where buried love doth live, +Hung with the trophies of my lovers gone, +Who all their parts of me to thee did give; +That due of many now is thine alone: +Their images I loved I view in thee, +And thou, all they, hast all the all of me. +If thou survive my well-contented day, +When that churl Death my bones with dust shall cover, +And shalt by fortune once more re-survey +These poor rude lines of thy deceased lover, +Compare them with the bettering of the time, +And though they be outstripp'd by every pen, +Reserve them for my love, not for their rhyme, +Exceeded by the height of happier men. +O, then vouchsafe me but this loving thought: +'Had my friend's Muse grown with this growing age, +A dearer birth than this his love had brought, +To march in ranks of better equipage: +But since he died and poets better prove, +Theirs for their style I'll read, his for his love.' +Full many a glorious morning have I seen +Flatter the mountain-tops with sovereign eye, +Kissing with golden face the meadows green, +Gilding pale streams with heavenly alchemy; +Anon permit the basest clouds to ride +With ugly rack on his celestial face, +And from the forlorn world his visage hide, +Stealing unseen to west with this disgrace: +Even so my sun one early morn did shine +With all triumphant splendor on my brow; +But out, alack! he was but one hour mine; +The region cloud hath mask'd him from me now. +Yet him for this my love no whit disdaineth; +Suns of the world may stain when heaven's sun staineth. +Why didst thou promise such a beauteous day, +And make me travel forth without my cloak, +To let base clouds o'ertake me in my way, +Hiding thy bravery in their rotten smoke? +'Tis not enough that through the cloud thou break, +To dry the rain on my storm-beaten face, +For no man well of such a salve can speak +That heals the wound and cures not the disgrace: +Nor can thy shame give physic to my grief; +Though thou repent, yet I have still the loss: +The offender's sorrow lends but weak relief +To him that bears the strong offence's cross. +Ah! but those tears are pearl which thy love sheds, +And they are rich and ransom all ill deeds. +No more be grieved at that which thou hast done: +Roses have thorns, and silver fountains mud; +Clouds and eclipses stain both moon and sun, +And loathsome canker lives in sweetest bud. +All men make faults, and even I in this, +Authorizing thy trespass with compare, +Myself corrupting, salving thy amiss, +Excusing thy sins more than thy sins are; +For to thy sensual fault I bring in sense-- +Thy adverse party is thy advocate-- +And 'gainst myself a lawful plea commence: +Such civil war is in my love and hate +That I an accessary needs must be +To that sweet thief which sourly robs from me. +Let me confess that we two must be twain, +Although our undivided loves are one: +So shall those blots that do with me remain +Without thy help by me be borne alone. +In our two loves there is but one respect, +Though in our lives a separable spite, +Which though it alter not love's sole effect, +Yet doth it steal sweet hours from love's delight. +I may not evermore acknowledge thee, +Lest my bewailed guilt should do thee shame, +Nor thou with public kindness honour me, +Unless thou take that honour from thy name: +But do not so; I love thee in such sort +As, thou being mine, mine is thy good report. +As a decrepit father takes delight +To see his active child do deeds of youth, +So I, made lame by fortune's dearest spite, +Take all my comfort of thy worth and truth. +For whether beauty, birth, or wealth, or wit, +Or any of these all, or all, or more, +Entitled in thy parts do crowned sit, +I make my love engrafted to this store: +So then I am not lame, poor, nor despised, +Whilst that this shadow doth such substance give +That I in thy abundance am sufficed +And by a part of all thy glory live. +Look, what is best, that best I wish in thee: +This wish I have; then ten times happy me! \ No newline at end of file diff --git a/rust/src/bench/src/datasets/speed_bench.rs b/rust/src/bench/src/datasets/speed_bench.rs new file mode 100644 index 000000000000..a6490b76e851 --- /dev/null +++ b/rust/src/bench/src/datasets/speed_bench.rs @@ -0,0 +1,303 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; + +use super::SampleRequest; +use crate::cli::SpeedBenchConfig; +use crate::error::{BenchError, Result}; +use crate::tokenizer::TokenizerKind; + +/// Marker text for masked entries that need external fetch. +const MASKED_PREFIX: &str = "FULL BENCHMARK DATA SHOULD BE FETCHED"; + +/// Cache directory for downloaded SPEED-Bench datasets. +fn cache_dir() -> std::path::PathBuf { + dirs::cache_dir() + .unwrap_or_else(|| std::path::PathBuf::from("/tmp")) + .join("vllm-bench") + .join("datasets") +} + +/// Download SPEED-Bench dataset from HuggingFace datasets-server API. +/// Results are cached as JSON locally for subsequent runs. +pub fn download_speed_bench(config: SpeedBenchConfig) -> Result { + let config_name = config.as_str(); + + let dir = cache_dir(); + std::fs::create_dir_all(&dir)?; + let cache_path = dir.join(format!("speed-bench-{config_name}.json")); + + // Return cached file if it exists + if cache_path.exists() { + let path_str = cache_path.to_string_lossy().to_string(); + println!("SPEED-Bench ({config_name}) cached: {path_str}"); + return Ok(path_str); + } + + println!("Downloading SPEED-Bench ({config_name}) from HuggingFace datasets-server..."); + + let client = reqwest::blocking::Client::builder() + .timeout(std::time::Duration::from_secs(120)) + .build() + .map_err(|e| BenchError::Config(format!("Failed to build HTTP client: {e}")))?; + + let mut all_rows: Vec = Vec::new(); + let mut offset = 0usize; + let page_size = 100usize; + + loop { + let url = format!( + "https://datasets-server.huggingface.co/rows\ + ?dataset=nvidia/SPEED-Bench\ + &config={config_name}\ + &split=test\ + &offset={offset}\ + &length={page_size}" + ); + + // Retry on transient errors (502, 503, timeouts) + let max_retries = 3; + let mut data: Option = None; + for attempt in 0..=max_retries { + let resp = match client.get(&url).send() { + Ok(r) => r, + Err(e) => { + if attempt < max_retries { + std::thread::sleep(std::time::Duration::from_secs( + 2 * (attempt as u64 + 1), + )); + continue; + } + return Err(BenchError::Config(format!( + "SPEED-Bench download failed after {max_retries} retries: {e}" + ))); + } + }; + + if resp.status().is_server_error() && attempt < max_retries { + std::thread::sleep(std::time::Duration::from_secs(2 * (attempt as u64 + 1))); + continue; + } + + if !resp.status().is_success() { + return Err(BenchError::Config(format!( + "SPEED-Bench API returned HTTP {}", + resp.status() + ))); + } + + data = Some(resp.json().map_err(|e| { + BenchError::Config(format!("Failed to parse SPEED-Bench API response: {e}")) + })?); + break; + } + + let data = data.unwrap(); + + let rows = data["rows"] + .as_array() + .ok_or_else(|| BenchError::Config("No 'rows' in API response".into()))?; + + if rows.is_empty() { + break; + } + + for row in rows { + if let Some(row_data) = row.get("row") { + all_rows.push(row_data.clone()); + } + } + + let fetched = rows.len(); + offset += fetched; + + // Print progress + let total = data["num_rows_total"].as_u64().unwrap_or(0); + eprint!("\r Fetched {offset}/{total} rows..."); + + if fetched < page_size { + break; + } + } + eprintln!(); // newline after progress + + if all_rows.is_empty() { + return Err(BenchError::Config( + "SPEED-Bench download returned no rows".into(), + )); + } + + // Save to cache + let json_str = serde_json::to_string(&all_rows)?; + std::fs::write(&cache_path, &json_str)?; + + let path_str = cache_path.to_string_lossy().to_string(); + println!( + "SPEED-Bench ({config_name}): {} rows saved to {path_str}", + all_rows.len() + ); + Ok(path_str) +} + +/// Load SPEED-Bench dataset and convert to SampleRequests. +/// +/// Filters out masked entries and optionally filters by category. +/// Requires an output length override since SPEED-Bench has no reference outputs. +pub fn load_speed_bench_dataset( + tokenizer: &TokenizerKind, + dataset_path: &str, + num_requests: usize, + output_len: usize, + seed: u64, + request_id_prefix: &str, + category_filter: Option<&str>, + no_oversample: bool, + disable_shuffle: bool, + max_input_len: Option, +) -> Result> { + let content = std::fs::read_to_string(dataset_path).map_err(|e| { + BenchError::Config(format!( + "Failed to read SPEED-Bench file '{dataset_path}': {e}" + )) + })?; + + let entries: Vec = serde_json::from_str(&content) + .map_err(|e| BenchError::Config(format!("Invalid JSON in SPEED-Bench file: {e}")))?; + + // Filter entries + let mut filtered: Vec<&serde_json::Value> = entries + .iter() + .filter(|entry| { + // Must have turns array with at least one non-empty entry + let turns = match entry.get("turns").and_then(|t| t.as_array()) { + Some(t) if !t.is_empty() => t, + _ => return false, + }; + + // Skip masked entries + let first_turn = turns[0].as_str().unwrap_or(""); + if first_turn.starts_with(MASKED_PREFIX) || first_turn.is_empty() { + return false; + } + + // Single-turn only: skip multi-turn entries + let is_multiturn = entry.get("multiturn").and_then(|m| m.as_bool()).unwrap_or(false); + if is_multiturn { + return false; + } + + // Category filter + if let Some(cat) = category_filter { + let entry_cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or(""); + if entry_cat != cat { + return false; + } + } + + true + }) + .collect(); + + if filtered.is_empty() { + let cat_msg = category_filter.map(|c| format!(" with category '{c}'")).unwrap_or_default(); + return Err(BenchError::Config(format!( + "No valid single-turn entries in SPEED-Bench{cat_msg}. \ + Try a different --speed-bench-config or remove --speed-bench-category filter." + ))); + } + + // Shuffle + let mut rng = StdRng::seed_from_u64(seed); + if !disable_shuffle { + filtered.shuffle(&mut rng); + } + + // Build SampleRequests + let mut samples = Vec::new(); + let mut idx = 0; + + for entry in &filtered { + if samples.len() >= num_requests { + break; + } + + let turns = entry["turns"].as_array().unwrap(); + let prompt = turns[0].as_str().unwrap_or(""); + + // Tokenize to get prompt length + let prompt_ids = tokenizer.encode(prompt, false)?; + let prompt_len = prompt_ids.len(); + + if prompt_len < 4 { + continue; + } + + // Truncate if max_input_len is set + let (final_prompt, final_len) = if let Some(max_len) = max_input_len { + if prompt_len > max_len { + let truncated_ids = &prompt_ids[..max_len]; + let truncated_text = tokenizer.decode(truncated_ids, true)?; + (Arc::from(truncated_text.as_str()), max_len) + } else { + (Arc::from(prompt), prompt_len) + } + } else { + (Arc::from(prompt), prompt_len) + }; + + samples.push(SampleRequest { + prompt: final_prompt, + prompt_len: final_len, + expected_output_len: output_len, + request_id: Some(format!("{request_id_prefix}{idx}")), + ..Default::default() + }); + idx += 1; + } + + // Oversample if needed + if samples.len() < num_requests { + if no_oversample { + println!( + "Skipping oversampling. Total samples: {} (requested: {num_requests})", + samples.len() + ); + } else if !samples.is_empty() { + let original_len = samples.len(); + let needed = num_requests - original_len; + for i in 0..needed { + let mut req = samples[rng.random_range(0..original_len)].clone(); + req.request_id = Some(format!("{request_id_prefix}{}", original_len + i)); + samples.push(req); + } + println!( + "Oversampled SPEED-Bench from {original_len} to {} total samples.", + samples.len() + ); + } + } + + if samples.is_empty() { + return Err(BenchError::Config( + "No valid samples after filtering SPEED-Bench dataset.".into(), + )); + } + + // Print category distribution + let mut cat_counts: std::collections::HashMap<&str, usize> = std::collections::HashMap::new(); + for entry in &filtered[..filtered.len().min(samples.len())] { + let cat = entry.get("category").and_then(|c| c.as_str()).unwrap_or("unknown"); + *cat_counts.entry(cat).or_insert(0) += 1; + } + let mut cats: Vec<_> = cat_counts.into_iter().collect(); + cats.sort_by_key(|b| std::cmp::Reverse(b.1)); + let cat_str: Vec = cats.iter().map(|(k, v)| format!("{k}:{v}")).collect(); + println!("SPEED-Bench categories: {}", cat_str.join(", ")); + + Ok(samples) +} diff --git a/rust/src/bench/src/error.rs b/rust/src/bench/src/error.rs new file mode 100644 index 000000000000..77299a4d3996 --- /dev/null +++ b/rust/src/bench/src/error.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum BenchError { + #[error("HTTP request failed: {0}")] + Http(#[from] reqwest::Error), + + #[error("JSON error: {0}")] + Json(#[from] serde_json::Error), + + #[error("Tokenizer error: {0}")] + Tokenizer(String), + + /// The server's /tokenize//detokenize endpoint is not usable (4xx status: + /// not exposed, or rejected by a gateway such as LLM-d/EPP that returns + /// 400 instead of 404). Callers treat this as "skip verification", unlike + /// `Tokenizer` errors which are genuine failures. + #[error("tokenize endpoint unavailable: {0}")] + TokenizeUnavailable(String), + + #[error("Configuration error: {0}")] + Config(String), + + #[error("Endpoint not ready after {0}s: {1}")] + EndpointTimeout(u64, String), + + #[error("Backend error: {0}")] + Backend(String), + + #[error("IO error: {0}")] + Io(#[from] std::io::Error), +} + +pub type Result = std::result::Result; diff --git a/rust/src/bench/src/hub.rs b/rust/src/bench/src/hub.rs new file mode 100644 index 000000000000..0e1b3e962c8e --- /dev/null +++ b/rust/src/bench/src/hub.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Sync facade over the async `hf_hub` API. +//! +//! The workspace bans rustls (`rust/deny.toml`), but hf-hub's sync `ureq` +//! backend unconditionally pulls ureq's default rustls feature. So we use the +//! reqwest/native-tls tokio API instead, and bridge blocking callers (dataset +//! loaders, tokenizer fallback in rayon threads) by running each download on a +//! dedicated thread with its own single-threaded runtime. + +use std::path::PathBuf; + +/// A handle to a HuggingFace Hub repo, downloading via hf-hub's on-disk cache. +pub struct HubRepo { + repo: hf_hub::Repo, +} + +impl HubRepo { + pub fn model(model_id: String) -> Self { + Self { + repo: hf_hub::Repo::model(model_id), + } + } + + pub fn dataset(repo_id: String) -> Self { + Self { + repo: hf_hub::Repo::dataset(repo_id), + } + } + + /// Download (or fetch from cache) a single file from the repo. + /// Auth is handled by hf-hub via HF_TOKEN / the cached login token. + pub fn get(&self, filename: &str) -> Result { + let repo = self.repo.clone(); + let filename = filename.to_string(); + std::thread::spawn(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|e| format!("Failed to build download runtime: {e}"))?; + rt.block_on(async move { + let api = hf_hub::api::tokio::Api::new() + .map_err(|e| format!("Failed to init HF API: {e}"))?; + api.repo(repo).get(&filename).await.map_err(|e| format!("{e}")) + }) + }) + .join() + .map_err(|_| "HF Hub download thread panicked".to_string())? + } +} diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs new file mode 100644 index 000000000000..dd37ec56db38 --- /dev/null +++ b/rust/src/bench/src/main.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +mod backends; +mod benchmark; +mod cli; +mod compare; +mod config; +mod datasets; +mod error; +mod hub; +mod metrics; +mod multi_run; +mod multi_turn; +mod output; +mod rate_control; +mod ready_checker; +mod sweep; +mod tiktoken; +mod tokenizer; + +#[cfg(not(target_env = "msvc"))] +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + +use anyhow::Context; +use clap::Parser; +use cli::Cli; +use config::BenchConfig; + +fn main() -> anyhow::Result<()> { + // Raise the open-file soft limit to the hard limit. High-concurrency + // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit. + if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX) + && new > 1024 + { + eprintln!("Open-file limit: {new}"); + } + + let cli = Cli::parse(); + + // --- Compare mode: no server needed, just diff two JSON files --- + if let Some(ref files) = cli.compare { + return compare::compare_results(&files[0], &files[1]).context("Comparison failed"); + } + + let config = BenchConfig::from_cli(&cli).context("Configuration error")?; + + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("Failed to build tokio runtime"); + + runtime + .block_on(async { + if config.multi_turn { + if let Some(ref sweep_mc) = cli.sweep_max_concurrency { + // --- Sweep over concurrency in multi-turn mode --- + let values = sweep::parse_concurrency_values(sweep_mc) + .context("Invalid --sweep-max-concurrency")?; + sweep::run_multi_turn_concurrency_sweep( + &config, + &values, + cli.sweep_num_prompts_factor, + ) + .await?; + } else { + // --- Single multi-turn conversation benchmark --- + multi_turn::run_multi_turn_benchmark(&config).await?; + } + } else if let Some(ref sweep_mc) = cli.sweep_max_concurrency { + // --- Sweep over max-concurrency --- + let values = sweep::parse_concurrency_values(sweep_mc) + .context("Invalid --sweep-max-concurrency")?; + sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor) + .await?; + } else if let Some(ref sweep_rate) = cli.sweep_request_rate { + // --- Sweep over request-rate --- + let values = + sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?; + sweep::run_rate_sweep(&config, &values).await?; + } else if cli.num_runs > 1 { + // --- Multi-run with statistical aggregation --- + multi_run::run_multi(&config, cli.num_runs).await?; + } else { + // --- Normal single benchmark --- + benchmark::run_benchmark(&config).await?; + } + anyhow::Ok(()) + }) + .context("Benchmark failed") +} diff --git a/rust/src/bench/src/metrics/calculator.rs b/rust/src/bench/src/metrics/calculator.rs new file mode 100644 index 000000000000..195dad93de3e --- /dev/null +++ b/rust/src/bench/src/metrics/calculator.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use crate::backends::RequestFuncOutput; +use crate::config::GoodputConfig; +use crate::datasets::SampleRequest; +use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics}; +use crate::multi_turn::ConversationOutput; + +/// Calculate benchmark metrics from request outputs. +/// +/// Mirrors Python's `calculate_metrics()` from serve.py:392-599. +pub fn calculate_metrics( + input_requests: &[SampleRequest], + outputs: &[RequestFuncOutput], + dur_s: f64, + selected_percentiles: &[f64], + _has_tokenizer: bool, + goodput_config: &GoodputConfig, +) -> (BenchmarkMetrics, Vec) { + let mut actual_output_lens: Vec = Vec::with_capacity(outputs.len()); + let mut total_input: usize = 0; + let mut completed: usize = 0; + let mut itls: Vec = Vec::new(); + let mut tpots: Vec = Vec::new(); + // Per-request TPOT for goodput SLO checking (parallel to ttfts/e2els). + let mut all_tpots: Vec = Vec::new(); + let mut ttfts: Vec = Vec::new(); + let mut e2els: Vec = Vec::new(); + + for (i, output) in outputs.iter().enumerate() { + if output.success { + let output_len = if output.output_tokens > 0 { + output.output_tokens + } else { + // Fallback: first token + ITL entries. + // Python re-encodes generated_text when tokenizer is available, + // but with vLLM's stream_options.include_usage=true, + // output_tokens is always set so this path is rarely hit. + 1 + output.itl.len() + }; + + actual_output_lens.push(output_len); + total_input += input_requests[i].prompt_len; + + if output_len > 1 { + let latency_minus_ttft = output.latency - output.ttft; + let tpot = latency_minus_ttft / (output_len as f64 - 1.0); + tpots.push(tpot); + all_tpots.push(tpot); + } else { + all_tpots.push(0.0); + } + + itls.extend_from_slice(&output.itl); + ttfts.push(output.ttft); + e2els.push(output.latency); + completed += 1; + } else { + actual_output_lens.push(0); + } + } + + let failed = outputs.len() - completed; + + // Print failed request errors (capped to 10) + let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect(); + if !failed_outputs.is_empty() { + eprintln!("Failed requests during benchmark run detected (capping to 10):"); + for (i, err) in failed_outputs.iter().take(10).enumerate() { + eprintln!("Error {i}: {}", err.error); + } + } + + // Calculate max output tokens per second and max concurrent requests + let mut max_output_tokens_per_s = 0.0_f64; + let mut max_concurrent_requests: usize = 0; + + let successful_outputs: Vec<&RequestFuncOutput> = + outputs.iter().filter(|o| o.success).collect(); + + if !successful_outputs.is_empty() { + let min_start_time = + successful_outputs.iter().map(|o| o.start_time).fold(f64::INFINITY, f64::min); + let max_end_time = successful_outputs + .iter() + .map(|o| o.start_time + o.latency) + .fold(f64::NEG_INFINITY, f64::max); + + let raw_duration = (max_end_time - min_start_time).ceil() as usize + 1; + // Cap at 24 hours to prevent OOM from corrupted timing data + let duration_seconds = raw_duration.min(86_400); + let mut tokens_per_second = vec![0.0_f64; duration_seconds]; + let mut concurrent_per_second = vec![0usize; duration_seconds]; + + for output in &successful_outputs { + // Calculate token generation timestamps + let mut token_times = vec![output.start_time + output.ttft]; + let mut current_time = token_times[0]; + for itl_value in &output.itl { + current_time += itl_value; + token_times.push(current_time); + } + + // Add tokens to second buckets + for token_time in &token_times { + let bucket = (token_time - min_start_time) as usize; + if bucket < duration_seconds { + tokens_per_second[bucket] += 1.0; + } + } + + // Track concurrent requests + let start_second = (output.start_time - min_start_time) as usize; + let end_second = ((output.start_time + output.latency) - min_start_time) as usize; + for slot in concurrent_per_second + .iter_mut() + .take(end_second.min(duration_seconds - 1) + 1) + .skip(start_second) + { + *slot += 1; + } + } + + max_output_tokens_per_s = tokens_per_second.iter().cloned().fold(0.0_f64, f64::max); + max_concurrent_requests = *concurrent_per_second.iter().max().unwrap_or(&0); + } + + let total_output: usize = actual_output_lens.iter().sum(); + + // Compute goodput: count requests meeting ALL specified SLOs. + // Mirrors Python serve.py:458-481. + let good_completed = if !goodput_config.is_empty() { + let mut good = 0usize; + // ttfts, all_tpots, e2els are parallel (one per successful request) + for i in 0..ttfts.len() { + let mut is_good = true; + if let Some(slo_ms) = goodput_config.ttft_ms + && ttfts[i] * 1000.0 > slo_ms + { + is_good = false; + } + if let Some(slo_ms) = goodput_config.tpot_ms + && all_tpots[i] * 1000.0 > slo_ms + { + is_good = false; + } + if let Some(slo_ms) = goodput_config.e2el_ms + && e2els[i] * 1000.0 > slo_ms + { + is_good = false; + } + if is_good { + good += 1; + } + } + good + } else { + 0 + }; + + // Sort each metric array once, then compute median + percentiles from sorted data + let sorted_ttfts = sort_clone(&ttfts); + let sorted_tpots = sort_clone(&tpots); + let sorted_itls = sort_clone(&itls); + let sorted_e2els = sort_clone(&e2els); + + let request_goodput = if !goodput_config.is_empty() { + good_completed as f64 / dur_s + } else { + 0.0 + }; + + let metrics = BenchmarkMetrics { + completed, + failed, + total_input, + total_output, + request_throughput: completed as f64 / dur_s, + request_goodput, + input_throughput: total_input as f64 / dur_s, + output_throughput: total_output as f64 / dur_s, + total_token_throughput: (total_input + total_output) as f64 / dur_s, + mean_ttft_ms: mean(&ttfts) * 1000.0, + median_ttft_ms: median_sorted(&sorted_ttfts) * 1000.0, + std_ttft_ms: std_dev(&ttfts) * 1000.0, + percentiles_ttft_ms: percentiles_from_sorted(&sorted_ttfts, selected_percentiles), + mean_tpot_ms: mean(&tpots) * 1000.0, + median_tpot_ms: median_sorted(&sorted_tpots) * 1000.0, + std_tpot_ms: std_dev(&tpots) * 1000.0, + percentiles_tpot_ms: percentiles_from_sorted(&sorted_tpots, selected_percentiles), + mean_itl_ms: mean(&itls) * 1000.0, + median_itl_ms: median_sorted(&sorted_itls) * 1000.0, + std_itl_ms: std_dev(&itls) * 1000.0, + percentiles_itl_ms: percentiles_from_sorted(&sorted_itls, selected_percentiles), + mean_e2el_ms: mean(&e2els) * 1000.0, + median_e2el_ms: median_sorted(&sorted_e2els) * 1000.0, + std_e2el_ms: std_dev(&e2els) * 1000.0, + percentiles_e2el_ms: percentiles_from_sorted(&sorted_e2els, selected_percentiles), + max_output_tokens_per_s, + max_concurrent_requests, + steady_state: None, + }; + + (metrics, actual_output_lens) +} + +fn mean(data: &[f64]) -> f64 { + if data.is_empty() { + return 0.0; + } + data.iter().sum::() / data.len() as f64 +} + +fn sort_clone(data: &[f64]) -> Vec { + let mut sorted = data.to_vec(); + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + sorted +} + +fn median_sorted(sorted: &[f64]) -> f64 { + if sorted.is_empty() { + return 0.0; + } + let mid = sorted.len() / 2; + if sorted.len().is_multiple_of(2) { + (sorted[mid - 1] + sorted[mid]) / 2.0 + } else { + sorted[mid] + } +} + +fn std_dev(data: &[f64]) -> f64 { + if data.is_empty() { + return 0.0; + } + let m = mean(data); + let variance = data.iter().map(|x| (x - m).powi(2)).sum::() / data.len() as f64; + variance.sqrt() +} + +fn percentile_sorted(sorted_data: &[f64], p: f64) -> f64 { + if sorted_data.is_empty() { + return 0.0; + } + if sorted_data.len() == 1 { + return sorted_data[0]; + } + // Use the same interpolation as numpy (linear) + let idx = p / 100.0 * (sorted_data.len() - 1) as f64; + let lo = idx.floor() as usize; + let hi = idx.ceil() as usize; + let frac = idx - lo as f64; + if lo == hi { + sorted_data[lo] + } else { + sorted_data[lo] * (1.0 - frac) + sorted_data[hi] * frac + } +} + +fn percentiles_from_sorted(sorted: &[f64], percentiles: &[f64]) -> Vec<(f64, f64)> { + if sorted.is_empty() { + return percentiles.iter().map(|&p| (p, 0.0)).collect(); + } + percentiles + .iter() + .map(|&p| (p, percentile_sorted(sorted, p) * 1000.0)) + .collect() +} + +/// Calculate benchmark metrics for embedding/pooling requests. +/// +/// Mirrors Python's `calculate_metrics_for_embeddings()` from serve.py. +/// Key differences from `calculate_metrics()`: +/// - Uses `output.prompt_len` (server-reported) instead of `input_requests[i].prompt_len` +/// - Only computes E2EL (no TTFT/TPOT/ITL since pooling is non-streaming) +/// - `total_output` is 0 (pooling produces no output tokens) +/// - `total_token_throughput` = `total_input / dur_s` (input tokens only) +pub fn calculate_embedding_metrics( + outputs: &[RequestFuncOutput], + dur_s: f64, + selected_percentiles: &[f64], +) -> BenchmarkMetrics { + let mut total_input: usize = 0; + let mut completed: usize = 0; + let mut e2els: Vec = Vec::new(); + + for output in outputs { + if output.success { + e2els.push(output.latency); + completed += 1; + total_input += output.prompt_len; + } + } + + let failed = outputs.len() - completed; + + // Print failed request errors (capped to 10) + let failed_outputs: Vec<&RequestFuncOutput> = outputs.iter().filter(|o| !o.success).collect(); + if !failed_outputs.is_empty() { + eprintln!("Failed requests during benchmark run detected (capping to 10):"); + for (i, err) in failed_outputs.iter().take(10).enumerate() { + eprintln!("Error {i}: {}", err.error); + } + } + + // Compute peak concurrent requests from start_time + latency windows + let successful_outputs: Vec<&RequestFuncOutput> = + outputs.iter().filter(|o| o.success).collect(); + let max_concurrent_requests = if !successful_outputs.is_empty() { + let min_start = + successful_outputs.iter().map(|o| o.start_time).fold(f64::INFINITY, f64::min); + let max_end = successful_outputs + .iter() + .map(|o| o.start_time + o.latency) + .fold(f64::NEG_INFINITY, f64::max); + + let raw_duration = (max_end - min_start).ceil() as usize + 1; + let duration_seconds = raw_duration.min(86_400); + let mut concurrent_per_second = vec![0usize; duration_seconds]; + + for output in &successful_outputs { + let start_second = (output.start_time - min_start) as usize; + let end_second = ((output.start_time + output.latency) - min_start) as usize; + for slot in concurrent_per_second + .iter_mut() + .take(end_second.min(duration_seconds - 1) + 1) + .skip(start_second) + { + *slot += 1; + } + } + + *concurrent_per_second.iter().max().unwrap_or(&0) + } else { + 0 + }; + + let sorted_e2els = sort_clone(&e2els); + + BenchmarkMetrics { + completed, + failed, + total_input, + total_output: 0, + request_throughput: completed as f64 / dur_s, + request_goodput: 0.0, + input_throughput: total_input as f64 / dur_s, + output_throughput: 0.0, + total_token_throughput: total_input as f64 / dur_s, + mean_ttft_ms: 0.0, + median_ttft_ms: 0.0, + std_ttft_ms: 0.0, + percentiles_ttft_ms: Vec::new(), + mean_tpot_ms: 0.0, + median_tpot_ms: 0.0, + std_tpot_ms: 0.0, + percentiles_tpot_ms: Vec::new(), + mean_itl_ms: 0.0, + median_itl_ms: 0.0, + std_itl_ms: 0.0, + percentiles_itl_ms: Vec::new(), + mean_e2el_ms: mean(&e2els) * 1000.0, + median_e2el_ms: median_sorted(&sorted_e2els) * 1000.0, + std_e2el_ms: std_dev(&e2els) * 1000.0, + percentiles_e2el_ms: percentiles_from_sorted(&sorted_e2els, selected_percentiles), + max_output_tokens_per_s: 0.0, + max_concurrent_requests, + steady_state: None, + } +} + +/// Calculate multi-turn benchmark metrics with per-turn breakdown. +/// +/// 1. Flatten all turns into SampleRequest/RequestFuncOutput pairs → overall metrics +/// 2. Group by turn_index → per-turn metrics +/// 3. Conversation-level stats +pub fn calculate_multi_turn_metrics( + conversation_outputs: &[ConversationOutput], + dur_s: f64, + selected_percentiles: &[f64], + goodput_config: &GoodputConfig, + max_turn_count: usize, +) -> MultiTurnMetrics { + // Flatten all turns for overall metrics + let mut all_requests: Vec = Vec::new(); + let mut all_outputs: Vec = Vec::new(); + + for conv in conversation_outputs { + for turn in &conv.turns { + all_requests.push(SampleRequest { + prompt: std::sync::Arc::from(""), + prompt_len: turn.cumulative_input_tokens, + expected_output_len: turn.request_output.output_tokens, + request_id: None, + ..Default::default() + }); + all_outputs.push(turn.request_output.clone()); + } + } + + let (overall, _) = calculate_metrics( + &all_requests, + &all_outputs, + dur_s, + selected_percentiles, + true, + goodput_config, + ); + + // Per-turn breakdown + let mut per_turn: Vec = Vec::new(); + for turn_idx in 0..max_turn_count { + let mut turn_requests: Vec = Vec::new(); + let mut turn_outputs: Vec = Vec::new(); + + for conv in conversation_outputs { + if let Some(turn) = conv.turns.get(turn_idx) { + turn_requests.push(SampleRequest { + prompt: std::sync::Arc::from(""), + prompt_len: turn.cumulative_input_tokens, + expected_output_len: turn.request_output.output_tokens, + request_id: None, + ..Default::default() + }); + turn_outputs.push(turn.request_output.clone()); + } + } + + if !turn_outputs.is_empty() { + let (turn_metrics, _) = calculate_metrics( + &turn_requests, + &turn_outputs, + dur_s, + selected_percentiles, + true, + goodput_config, + ); + per_turn.push(turn_metrics); + } + } + + // Conversation-level stats + let conversations_completed = conversation_outputs.iter().filter(|c| c.all_success).count(); + let conversations_failed = conversation_outputs.len() - conversations_completed; + + let avg_turns_completed = if conversation_outputs.is_empty() { + 0.0 + } else { + conversation_outputs.iter().map(|c| c.turns.len() as f64).sum::() + / conversation_outputs.len() as f64 + }; + + let avg_conversation_duration_ms = if conversation_outputs.is_empty() { + 0.0 + } else { + conversation_outputs.iter().map(|c| c.total_duration_ms).sum::() + / conversation_outputs.len() as f64 + }; + + MultiTurnMetrics { + overall, + per_turn, + conversations_completed, + conversations_failed, + avg_turns_completed, + avg_conversation_duration_ms, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_mean_empty() { + assert_eq!(mean(&[]), 0.0); + } + + #[test] + fn test_mean_values() { + assert!((mean(&[1.0, 2.0, 3.0]) - 2.0).abs() < 1e-10); + } + + #[test] + fn test_median_odd() { + let sorted = sort_clone(&[3.0, 1.0, 2.0]); + assert!((median_sorted(&sorted) - 2.0).abs() < 1e-10); + } + + #[test] + fn test_median_even() { + let sorted = sort_clone(&[1.0, 2.0, 3.0, 4.0]); + assert!((median_sorted(&sorted) - 2.5).abs() < 1e-10); + } + + #[test] + fn test_percentile_99() { + let data: Vec = (0..100).map(|i| i as f64).collect(); + let p99 = percentile_sorted(&data, 99.0); + assert!((p99 - 98.01).abs() < 0.1); + } +} diff --git a/rust/src/bench/src/metrics/mod.rs b/rust/src/bench/src/metrics/mod.rs new file mode 100644 index 000000000000..829291acc84e --- /dev/null +++ b/rust/src/bench/src/metrics/mod.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +pub mod calculator; +pub mod steady_state; + +use serde::{Deserialize, Serialize}; +pub use steady_state::SteadyStateMetrics; + +/// Multi-turn benchmark metrics with overall and per-turn breakdown. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MultiTurnMetrics { + pub overall: BenchmarkMetrics, + pub per_turn: Vec, + pub conversations_completed: usize, + pub conversations_failed: usize, + pub avg_turns_completed: f64, + pub avg_conversation_duration_ms: f64, +} + +/// Full benchmark metrics for generation tasks. +/// Matches Python's BenchmarkMetrics dataclass. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BenchmarkMetrics { + pub completed: usize, + pub failed: usize, + pub total_input: usize, + pub total_output: usize, + pub request_throughput: f64, + pub request_goodput: f64, + pub input_throughput: f64, + pub output_throughput: f64, + pub total_token_throughput: f64, + pub mean_ttft_ms: f64, + pub median_ttft_ms: f64, + pub std_ttft_ms: f64, + pub percentiles_ttft_ms: Vec<(f64, f64)>, + pub mean_tpot_ms: f64, + pub median_tpot_ms: f64, + pub std_tpot_ms: f64, + pub percentiles_tpot_ms: Vec<(f64, f64)>, + pub mean_itl_ms: f64, + pub median_itl_ms: f64, + pub std_itl_ms: f64, + pub percentiles_itl_ms: Vec<(f64, f64)>, + pub mean_e2el_ms: f64, + pub median_e2el_ms: f64, + pub std_e2el_ms: f64, + pub percentiles_e2el_ms: Vec<(f64, f64)>, + pub max_output_tokens_per_s: f64, + pub max_concurrent_requests: usize, + #[serde(default)] + pub steady_state: Option, +} diff --git a/rust/src/bench/src/metrics/steady_state.rs b/rust/src/bench/src/metrics/steady_state.rs new file mode 100644 index 000000000000..a403275b7ac6 --- /dev/null +++ b/rust/src/bench/src/metrics/steady_state.rs @@ -0,0 +1,629 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Steady-state metrics: throughput and TTFT measured over the window +//! during which client-side in-flight concurrency is at or above a +//! configurable fraction of `--max-concurrency`. Excludes ramp-up and +//! drain phases to reduce run-to-run variance at very high concurrency. + +use serde::{Deserialize, Serialize}; + +/// Bounds and metadata of the detected steady-state window. +/// +/// `observed_peak` is event-exact and may differ by 1 from the +/// bucket-approximated `max_concurrent_requests` on `BenchmarkMetrics`; +/// this is expected. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SteadyStateWindow { + pub start_s: f64, + pub end_s: f64, + pub duration_s: f64, + pub target_concurrency: usize, + pub threshold: f64, + pub threshold_abs: usize, + pub observed_peak: usize, + pub requests_started_in_window: usize, + pub requests_completed_in_window: usize, + pub requests_total: usize, + pub warning: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SteadyStateMetrics { + pub window: SteadyStateWindow, + pub request_throughput: f64, + pub input_throughput: f64, + pub output_throughput: f64, + pub total_token_throughput: f64, + pub mean_ttft_ms: f64, + pub median_ttft_ms: f64, + pub percentiles_ttft_ms: Vec<(f64, f64)>, + pub mean_tpot_ms: f64, + pub median_tpot_ms: f64, + pub p90_tpot_ms: f64, + pub p99_tpot_ms: f64, +} + +use crate::backends::RequestFuncOutput; + +/// Detect the steady-state window over `outputs` under the given target concurrency. +/// +/// Returns `None` when closed-loop gate is not met (no target), when concurrency +/// never crosses `threshold_abs`, or when both started and completed in-window +/// request counts are zero. +/// +/// Timestamps in the returned window are normalized so `start_s` / `end_s` are +/// "seconds from the earliest successful request's start_time." +/// +/// `user_min_window_s`: if `Some`, sets the minimum window duration below which +/// a `warning` is attached. If `None`, the caller is expected to have pre-resolved +/// the default `max(10.0, 0.1 * total_run_duration_s)`. +pub fn detect_window( + outputs: &[RequestFuncOutput], + target_concurrency: Option, + threshold: f64, + min_window_s: f64, + total_run_duration_s: f64, +) -> Option { + let target = target_concurrency?; + let threshold_abs = ((threshold * target as f64).ceil() as usize).max(1); + + // Collect sorted start / end timestamps from successful requests only. + let mut starts: Vec = outputs.iter().filter(|o| o.success).map(|o| o.start_time).collect(); + let mut ends: Vec = + outputs.iter().filter(|o| o.success).map(|o| o.start_time + o.latency).collect(); + if starts.is_empty() { + return None; + } + starts.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + ends.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + + let min_start = starts[0]; + let requests_total = outputs.iter().filter(|o| o.success).count(); + + // Two-pointer merge. Tie rule: process `end` before `start` at equal timestamps. + let mut i = 0usize; + let mut j = 0usize; + let mut concurrency: usize = 0; + let mut observed_peak: usize = 0; + let mut up_crossing: Option = None; + let mut last_down_crossing: Option = None; + + while i < starts.len() || j < ends.len() { + let next_is_end = match (starts.get(i), ends.get(j)) { + (Some(&s), Some(&e)) => e <= s, // tie -> end first + (None, Some(_)) => true, + (Some(_), None) => false, + (None, None) => break, + }; + + if next_is_end { + let t = ends[j]; + j += 1; + let before = concurrency; + concurrency = concurrency.saturating_sub(1); + if before >= threshold_abs && concurrency < threshold_abs { + last_down_crossing = Some(t); + } + } else { + let t = starts[i]; + i += 1; + let before = concurrency; + concurrency += 1; + if concurrency > observed_peak { + observed_peak = concurrency; + } + if before < threshold_abs && concurrency >= threshold_abs && up_crossing.is_none() { + up_crossing = Some(t); + } + } + } + + let start_abs = up_crossing?; + // A down-crossing always exists once an up-crossing exists — the final ends + // drain concurrency to 0. Fall back defensively to the last end timestamp. + let end_abs = last_down_crossing.or_else(|| ends.last().copied())?; + + let start_s = start_abs - min_start; + let end_s = end_abs - min_start; + let duration_s = end_s - start_s; + + // Count requests whose start / completion events land in [start_abs, end_abs). + let mut requests_started_in_window = 0usize; + let mut requests_completed_in_window = 0usize; + for o in outputs.iter().filter(|o| o.success) { + if o.start_time >= start_abs && o.start_time < end_abs { + requests_started_in_window += 1; + } + let end_t = o.start_time + o.latency; + if end_t >= start_abs && end_t < end_abs { + requests_completed_in_window += 1; + } + } + + if requests_started_in_window == 0 && requests_completed_in_window == 0 { + return None; + } + + let warning = if duration_s < min_window_s { + Some(format!( + "steady-state window is {:.1}s ({:.1}% of run); may not be representative", + duration_s, + 100.0 * duration_s / total_run_duration_s.max(1e-9) + )) + } else { + None + }; + + Some(SteadyStateWindow { + start_s, + end_s, + duration_s, + target_concurrency: target, + threshold, + threshold_abs, + observed_peak, + requests_started_in_window, + requests_completed_in_window, + requests_total, + warning, + }) +} + +use crate::datasets::SampleRequest; + +/// Compute steady-state metrics from per-request outputs, the detected window, +/// and the original request list (for input-token attribution). +/// +/// `is_pooling` gates off TTFT because pooling backends write `ttft = latency` +/// as a placeholder; reporting TTFT over that would silently be E2EL. +pub fn compute( + outputs: &[RequestFuncOutput], + input_requests: &[SampleRequest], + window: &SteadyStateWindow, + percentiles: &[f64], + is_pooling: bool, +) -> SteadyStateMetrics { + // Re-derive min_start to convert normalized window bounds back to absolute times. + let min_start = outputs + .iter() + .filter(|o| o.success) + .map(|o| o.start_time) + .fold(f64::INFINITY, f64::min); + let start_abs = window.start_s + min_start; + let end_abs = window.end_s + min_start; + let duration = window.duration_s.max(1e-9); + + let mut completed_in_window = 0usize; + let mut input_tokens_in_window: usize = 0; + let mut output_tokens_in_window: usize = 0; + let mut ttfts_in_window: Vec = Vec::new(); + let mut tpots_in_window: Vec = Vec::new(); + + for (idx, o) in outputs.iter().enumerate() { + if !o.success { + continue; + } + let started_in = o.start_time >= start_abs && o.start_time < end_abs; + let end_t = o.start_time + o.latency; + let completed_in = end_t >= start_abs && end_t < end_abs; + + if completed_in { + completed_in_window += 1; + } + + if started_in { + input_tokens_in_window += input_requests[idx].prompt_len; + if !is_pooling { + ttfts_in_window.push(o.ttft); + // Per-request TPOT: (latency - ttft) / (output_tokens - 1). + // Matches calculator.rs:46-50; skip requests with <= 1 token. + if o.output_tokens > 1 { + let tpot = (o.latency - o.ttft) / (o.output_tokens as f64 - 1.0); + tpots_in_window.push(tpot); + } + } + } + + // Per-token emission attribution (same approach as calculator.rs:98-113). + // Reconstruct: first token at start_time + ttft, then cumulative itl. + if o.output_tokens == 0 && o.itl.is_empty() { + continue; + } + let first_token_t = o.start_time + o.ttft; + if first_token_t >= start_abs && first_token_t < end_abs { + output_tokens_in_window += 1; + } + let mut t = first_token_t; + for &dt in &o.itl { + t += dt; + if t >= start_abs && t < end_abs { + output_tokens_in_window += 1; + } + if t >= end_abs { + break; // tokens only move forward in time + } + } + } + + let request_throughput = completed_in_window as f64 / duration; + let input_throughput = input_tokens_in_window as f64 / duration; + let output_throughput = output_tokens_in_window as f64 / duration; + let total_token_throughput = input_throughput + output_throughput; + + let (mean_ttft, median_ttft, pct_ttft) = if is_pooling { + (0.0, 0.0, Vec::new()) + } else { + dist_stats(&ttfts_in_window, percentiles) + }; + + let (mean_tpot, median_tpot, pct_tpot) = if is_pooling { + (0.0, 0.0, Vec::new()) + } else { + dist_stats(&tpots_in_window, &[90.0, 99.0]) + }; + let p90_tpot = pct_tpot.first().map(|(_, v)| *v).unwrap_or(0.0); + let p99_tpot = pct_tpot.get(1).map(|(_, v)| *v).unwrap_or(0.0); + + SteadyStateMetrics { + window: window.clone(), + request_throughput, + input_throughput, + output_throughput, + total_token_throughput, + mean_ttft_ms: mean_ttft, + median_ttft_ms: median_ttft, + percentiles_ttft_ms: pct_ttft, + mean_tpot_ms: mean_tpot, + median_tpot_ms: median_tpot, + p90_tpot_ms: p90_tpot, + p99_tpot_ms: p99_tpot, + } +} + +/// Returns (mean_ms, median_ms, percentiles_ms). Values are seconds in; output +/// is milliseconds. Empty input yields zeros with a zero-filled percentile vec +/// aligned to the requested percentile list. +fn dist_stats(values: &[f64], percentiles: &[f64]) -> (f64, f64, Vec<(f64, f64)>) { + if values.is_empty() { + let pct = percentiles.iter().map(|&p| (p, 0.0)).collect(); + return (0.0, 0.0, pct); + } + let mut sorted = values.to_vec(); + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let n = sorted.len(); + let mean = sorted.iter().sum::() / n as f64; + let median = if n.is_multiple_of(2) { + (sorted[n / 2 - 1] + sorted[n / 2]) / 2.0 + } else { + sorted[n / 2] + }; + let pct = percentiles + .iter() + .map(|&p| { + let idx = (p / 100.0) * (n - 1) as f64; + let lo = idx.floor() as usize; + let hi = idx.ceil() as usize; + let frac = idx - lo as f64; + let v = if lo == hi { + sorted[lo] + } else { + sorted[lo] * (1.0 - frac) + sorted[hi] * frac + }; + (p, v * 1000.0) + }) + .collect(); + (mean * 1000.0, median * 1000.0, pct) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::backends::RequestFuncOutput; + + /// Build a successful synthetic output with given `start_time` and `latency`. + fn mk(start: f64, latency: f64) -> RequestFuncOutput { + RequestFuncOutput { + success: true, + start_time: start, + latency, + ttft: 0.0, + itl: Vec::new(), + output_tokens: 0, + ..Default::default() + } + } + + #[test] + fn window_plateau_trapezoid() { + // 10 requests ramp in at t=0..10 each lasting 30s. + // Concurrency reaches 10 at t=10 and stays until t=30 when requests start ending. + // threshold = 1.0 means exact-target; threshold_abs = 10. + let mut outputs: Vec = (0..10).map(|i| mk(i as f64, 30.0)).collect(); + // Add one request fully inside the plateau so it can pick up as "in window" + outputs.push(mk(15.0, 5.0)); + + let w = detect_window(&outputs, Some(10), 1.0, 5.0, 30.0).unwrap(); + assert_eq!(w.target_concurrency, 10); + assert_eq!(w.threshold_abs, 10); + // Up-crossing at t=9 (the 10th start) — concurrency transitions 9→10. + // Window start reported in normalized time (min_start = 0.0). + assert!((w.start_s - 9.0).abs() < 1e-9, "start_s = {}", w.start_s); + // Last down-crossing: the first end at t=30 (concurrency 10→9, falling below 10). + assert!((w.end_s - 30.0).abs() < 1e-9, "end_s = {}", w.end_s); + assert!(w.duration_s > 0.0); + assert_eq!(w.observed_peak, 11); // 10 ramp-in + 1 extra = 11 peak briefly + assert!(w.warning.is_none()); + } + + #[test] + fn window_none_when_target_none() { + let outputs = vec![mk(0.0, 1.0)]; + assert!(detect_window(&outputs, None, 0.95, 1.0, 10.0).is_none()); + } + + #[test] + fn window_none_when_never_reaches_threshold() { + // Serial requests, concurrency never exceeds 1. + let outputs: Vec = (0..5).map(|i| mk(i as f64 * 10.0, 1.0)).collect(); + assert!(detect_window(&outputs, Some(5), 0.95, 1.0, 100.0).is_none()); + } + + #[test] + fn window_none_when_all_fail() { + let mut o = mk(0.0, 1.0); + o.success = false; + assert!(detect_window(&[o], Some(1), 1.0, 0.1, 10.0).is_none()); + } + + #[test] + fn window_threshold_one() { + // threshold=1.0, target=2 -> threshold_abs=2; needs both requests in flight. + let outputs = vec![mk(0.0, 5.0), mk(1.0, 5.0)]; + let w = detect_window(&outputs, Some(2), 1.0, 0.1, 10.0).unwrap(); + assert_eq!(w.threshold_abs, 2); + assert!((w.start_s - 1.0).abs() < 1e-9); + assert!((w.end_s - 5.0).abs() < 1e-9); + } + + #[test] + fn window_ties_net_to_zero() { + // A request finishes and another starts at exactly the same instant; + // tie rule (end before start) prevents concurrency > 2 blip. + let outputs = vec![ + mk(0.0, 5.0), // ends at 5.0 + mk(2.0, 3.0), // ends at 5.0 + mk(5.0, 2.0), // starts at 5.0 — ties with the two ends + ]; + let w = detect_window(&outputs, Some(2), 1.0, 0.1, 20.0).unwrap(); + assert_eq!(w.observed_peak, 2); + } + + #[test] + fn window_warning_when_short() { + let outputs: Vec = (0..10).map(|i| mk(i as f64 * 0.1, 1.0)).collect(); + // Plateau is ~0.1s; min_window_s = 5s => warning. + let w = detect_window(&outputs, Some(10), 1.0, 5.0, 10.0).unwrap(); + assert!(w.warning.is_some(), "expected warning"); + } + + #[test] + fn window_constant_concurrency_equals_full_run() { + // Concurrency is at target throughout — steady-state window covers the + // entire middle of the run where all 3 requests overlap. + let outputs = vec![mk(0.0, 10.0), mk(0.0, 10.0), mk(0.0, 10.0)]; + let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); + assert!((w.start_s - 0.0).abs() < 1e-9); + assert!((w.end_s - 10.0).abs() < 1e-9); + } + + use std::sync::Arc; + + use crate::datasets::SampleRequest; + + fn mk_full( + start: f64, + latency: f64, + ttft: f64, + itl: Vec, + output_tokens: usize, + ) -> RequestFuncOutput { + RequestFuncOutput { + success: true, + start_time: start, + latency, + ttft, + itl, + output_tokens, + prompt_len: 10, + ..Default::default() + } + } + + fn mk_sample(prompt_len: usize) -> SampleRequest { + SampleRequest { + prompt: Arc::from(""), + prompt_len, + expected_output_len: 0, + request_id: None, + ..Default::default() + } + } + + #[test] + fn compute_constant_concurrency_matches_full_run() { + // 3 identical requests, each 10s, each emitting 5 tokens uniformly via itl=2s. + let itl = vec![2.0, 2.0, 2.0, 2.0]; // 5 tokens: first at t=start+ttft, then 4 more + let outputs = vec![ + mk_full(0.0, 10.0, 2.0, itl.clone(), 5), + mk_full(0.0, 10.0, 2.0, itl.clone(), 5), + mk_full(0.0, 10.0, 2.0, itl.clone(), 5), + ]; + let requests: Vec = + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); + + let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); + let m = compute( + &outputs, + &requests, + &w, + &[99.0], + false, // is_pooling + ); + + // Window is [0.0, 10.0). All 3 requests start in window (3). + // Tokens per request: first token at ttft=2s, then 4 more at itl=2s each + // -> t = 2, 4, 6, 8, 10. The 10.0 token is AT end_s and is half-open excluded. + // -> 4 tokens per request land in window. + assert_eq!(w.requests_started_in_window, 3); + // Half-open window [0.0, 10.0) excludes the end_t=10.0 completions, so + // request_throughput is exactly 0.0 here — tightens the half-open invariant. + assert!((m.request_throughput - 0.0).abs() < 1e-9); + // output tokens in window = 4 per request * 3 = 12, over 10s = 1.2 tok/s + assert!( + (m.output_throughput - 1.2).abs() < 1e-6, + "got {}", + m.output_throughput + ); + // input throughput = 30 input tokens / 10s = 3.0 + assert!( + (m.input_throughput - 3.0).abs() < 1e-9, + "got {}", + m.input_throughput + ); + } + + #[test] + fn compute_drain_only_tokens_not_counted() { + // 1 "steady" request defining the window, + 1 late-admitted request whose + // 1000 output tokens land after end_s — must not pump output_throughput. + let mut outputs = vec![mk_full(0.0, 10.0, 0.0, vec![1.0; 9], 10)]; + // Late request: starts at 9.9 (within window), long output that drains past end_s. + let late_itl = vec![0.02; 999]; // 1000 tokens; emission starts after end_s + let late = mk_full(9.9, 30.0, 0.05, late_itl, 1000); + outputs.push(late); + let requests: Vec = + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); + + // Target = 2 so threshold_abs = 2; window opens when both are in flight. + let w = detect_window(&outputs, Some(2), 1.0, 0.1, 40.0).unwrap(); + let m = compute(&outputs, &requests, &w, &[99.0], false); + + // Late request started in window but most of its tokens are emitted AFTER end_s. + // Only a handful (≤3) should count. + // Window end is 10.0 (the first request's completion). + // Late request first token at 9.9 + 0.05 = 9.95 (in window). + // Next tokens at 9.97, 9.99, 10.01... -> only ~3 fit before end_s=10.0. + let expected_max_late_tokens = 6.0; // generous bound + let window_dur = w.duration_s; + // Upper bound on output throughput: 10 (from req 1) + ≤6 (from late) / window_dur. + let upper_bound = (10.0 + expected_max_late_tokens) / window_dur; + assert!( + m.output_throughput <= upper_bound, + "output_throughput {} exceeds upper bound {} — drain tokens are being counted", + m.output_throughput, + upper_bound + ); + } + + #[test] + fn compute_pooling_skips_ttft_and_tpot() { + // For pooling backends, ttft == latency is a placeholder; we must NOT emit TTFT stats. + // Pooling has no output tokens so TPOT is also zeroed. + let outputs = vec![ + mk_full(0.0, 5.0, 5.0, Vec::new(), 0), + mk_full(0.0, 5.0, 5.0, Vec::new(), 0), + ]; + let requests: Vec = + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); + let w = detect_window(&outputs, Some(2), 1.0, 0.1, 10.0).unwrap(); + let m = compute(&outputs, &requests, &w, &[99.0], /* is_pooling */ true); + assert_eq!(m.mean_ttft_ms, 0.0); + assert_eq!(m.median_ttft_ms, 0.0); + assert!(m.percentiles_ttft_ms.is_empty()); + assert_eq!(m.mean_tpot_ms, 0.0); + assert_eq!(m.median_tpot_ms, 0.0); + assert_eq!(m.p90_tpot_ms, 0.0); + assert_eq!(m.p99_tpot_ms, 0.0); + // Throughput fields still populated. + assert!(m.input_throughput > 0.0); + } + + #[test] + fn compute_tpot_matches_per_request_definition() { + // 3 requests, each 10s latency with ttft=2s, 5 output tokens. + // Per-request TPOT = (10 - 2) / (5 - 1) = 2s = 2000ms. + // All requests start at t=0 so all are in-window; window duration = 10s. + let itl = vec![2.0; 4]; + let outputs = vec![ + mk_full(0.0, 10.0, 2.0, itl.clone(), 5), + mk_full(0.0, 10.0, 2.0, itl.clone(), 5), + mk_full(0.0, 10.0, 2.0, itl.clone(), 5), + ]; + let requests: Vec = + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); + let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); + let m = compute(&outputs, &requests, &w, &[99.0], false); + + assert!((m.mean_tpot_ms - 2000.0).abs() < 1e-6); + assert!((m.median_tpot_ms - 2000.0).abs() < 1e-6); + assert!((m.p90_tpot_ms - 2000.0).abs() < 1e-6); + assert!((m.p99_tpot_ms - 2000.0).abs() < 1e-6); + } + + #[test] + fn json_roundtrip_some() { + let outputs = vec![mk_full(0.0, 10.0, 2.0, vec![2.0; 4], 5); 3]; + let requests: Vec = + outputs.iter().map(|o| mk_sample(o.prompt_len)).collect(); + let w = detect_window(&outputs, Some(3), 1.0, 0.1, 10.0).unwrap(); + let m = compute(&outputs, &requests, &w, &[99.0], false); + let s = serde_json::to_string(&m).unwrap(); + let parsed: SteadyStateMetrics = serde_json::from_str(&s).unwrap(); + assert!((parsed.output_throughput - m.output_throughput).abs() < 1e-12); + assert_eq!(parsed.window.target_concurrency, 3); + } + + #[test] + fn json_deserialize_missing_key_as_none() { + // Pre-feature JSON lacks `steady_state` — deserializing a wrapper that + // contains an Option field with #[serde(default)] should yield None. + #[derive(serde::Deserialize)] + struct Wrap { + #[serde(default)] + steady_state: Option, + } + let old_json = r#"{"other":"value"}"#; + let w: Wrap = serde_json::from_str(old_json).unwrap(); + assert!(w.steady_state.is_none()); + } + + #[test] + fn end_to_end_trapezoid_run() { + // 5 requests ramp in at t=0..5, each lasts 20s, each emits 10 tokens + // at 1s intervals starting at ttft=1s. + let itl = vec![1.0; 9]; + let outputs: Vec = + (0..5).map(|i| mk_full(i as f64, 20.0, 1.0, itl.clone(), 10)).collect(); + let requests: Vec = (0..5).map(|_| mk_sample(100)).collect(); + + let w = detect_window(&outputs, Some(5), 1.0, 1.0, 25.0).unwrap(); + assert_eq!(w.threshold_abs, 5); + assert!((w.start_s - 4.0).abs() < 1e-9); + assert!((w.end_s - 20.0).abs() < 1e-9); + + let m = compute(&outputs, &requests, &w, &[50.0, 99.0], false); + + // Round-trip through JSON. + let s = serde_json::to_string(&m).unwrap(); + let back: SteadyStateMetrics = serde_json::from_str(&s).unwrap(); + assert!((back.request_throughput - m.request_throughput).abs() < 1e-12); + assert!((back.output_throughput - m.output_throughput).abs() < 1e-12); + assert_eq!(back.window.target_concurrency, 5); + + // TTFT percentiles should have 2 entries matching the input percentiles. + assert_eq!(m.percentiles_ttft_ms.len(), 2); + assert!((m.percentiles_ttft_ms[0].0 - 50.0).abs() < 1e-9); + assert!((m.percentiles_ttft_ms[1].0 - 99.0).abs() < 1e-9); + } +} diff --git a/rust/src/bench/src/multi_run.rs b/rust/src/bench/src/multi_run.rs new file mode 100644 index 000000000000..18ff7753b949 --- /dev/null +++ b/rust/src/bench/src/multi_run.rs @@ -0,0 +1,357 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use crate::config::BenchConfig; +use crate::error::Result; + +/// Key metrics extracted from a single run's JSON result. +struct RunMetrics { + request_throughput: f64, + output_throughput: f64, + total_token_throughput: f64, + mean_ttft_ms: f64, + median_ttft_ms: f64, + p99_ttft_ms: f64, + mean_tpot_ms: f64, + median_tpot_ms: f64, + p99_tpot_ms: f64, + mean_itl_ms: f64, + mean_e2el_ms: f64, + p99_e2el_ms: f64, + max_output_tokens_per_s: f64, + completed: f64, + failed: f64, + duration: f64, + ss_request_throughput: Option, + ss_output_throughput: Option, + ss_input_throughput: Option, + ss_total_token_throughput: Option, + ss_mean_ttft_ms: Option, + ss_median_ttft_ms: Option, + ss_mean_tpot_ms: Option, + ss_median_tpot_ms: Option, + ss_p90_tpot_ms: Option, + ss_p99_tpot_ms: Option, +} + +impl RunMetrics { + fn from_json(json: &serde_json::Value) -> Self { + Self { + request_throughput: get(json, "request_throughput"), + output_throughput: get(json, "output_throughput"), + total_token_throughput: get(json, "total_token_throughput"), + mean_ttft_ms: get(json, "mean_ttft_ms"), + median_ttft_ms: get(json, "median_ttft_ms"), + p99_ttft_ms: get(json, "p99_ttft_ms"), + mean_tpot_ms: get(json, "mean_tpot_ms"), + median_tpot_ms: get(json, "median_tpot_ms"), + p99_tpot_ms: get(json, "p99_tpot_ms"), + mean_itl_ms: get(json, "mean_itl_ms"), + mean_e2el_ms: get(json, "mean_e2el_ms"), + p99_e2el_ms: get(json, "p99_e2el_ms"), + max_output_tokens_per_s: get(json, "max_output_tokens_per_s"), + completed: get(json, "completed"), + failed: get(json, "failed"), + duration: get(json, "duration"), + ss_request_throughput: get_ss_opt(json, "request_throughput"), + ss_output_throughput: get_ss_opt(json, "output_throughput"), + ss_input_throughput: get_ss_opt(json, "input_throughput"), + ss_total_token_throughput: get_ss_opt(json, "total_token_throughput"), + ss_mean_ttft_ms: get_ss_opt(json, "mean_ttft_ms"), + ss_median_ttft_ms: get_ss_opt(json, "median_ttft_ms"), + ss_mean_tpot_ms: get_ss_opt(json, "mean_tpot_ms"), + ss_median_tpot_ms: get_ss_opt(json, "median_tpot_ms"), + ss_p90_tpot_ms: get_ss_opt(json, "p90_tpot_ms"), + ss_p99_tpot_ms: get_ss_opt(json, "p99_tpot_ms"), + } + } +} + +fn get(json: &serde_json::Value, key: &str) -> f64 { + json.get(key).and_then(|v| v.as_f64()).unwrap_or(0.0) +} + +fn get_ss_opt(json: &serde_json::Value, key: &str) -> Option { + json.get("steady_state") + .and_then(|ss| if ss.is_null() { None } else { Some(ss) }) + .and_then(|ss| ss.get(key)) + .and_then(|v| v.as_f64()) +} + +/// Aggregated statistics for a single metric across N runs. +struct MetricStats { + label: &'static str, + mean: f64, + std: f64, + min: f64, + max: f64, +} + +/// Run the benchmark N times and report aggregated statistics. +pub async fn run_multi(config: &BenchConfig, num_runs: usize) -> Result<()> { + println!( + "{:=^70}", + format!(" Multi-Run Benchmark ({num_runs} runs) ") + ); + println!(); + + let mut all_runs: Vec = Vec::with_capacity(num_runs); + + for i in 0..num_runs { + println!("{:-^70}", format!(" Run {}/{} ", i + 1, num_runs)); + + let mut run_config = config.clone(); + // Suppress per-run save + run_config.save_result = false; + run_config.append_result = false; + + let result = crate::benchmark::run_benchmark(&run_config).await?; + all_runs.push(RunMetrics::from_json(&result)); + + println!(); + } + + print_multi_run_summary(&all_runs); + Ok(()) +} + +fn print_multi_run_summary(runs: &[RunMetrics]) { + let n = runs.len(); + + // Collect each metric into a series, compute stats + let stats = vec![ + compute_stats("Request throughput (req/s)", runs, |r| r.request_throughput), + compute_stats("Output throughput (tok/s)", runs, |r| r.output_throughput), + compute_stats("Total token throughput (tok/s)", runs, |r| { + r.total_token_throughput + }), + compute_stats("Peak output tokens/s", runs, |r| r.max_output_tokens_per_s), + compute_stats("Mean TTFT (ms)", runs, |r| r.mean_ttft_ms), + compute_stats("Median TTFT (ms)", runs, |r| r.median_ttft_ms), + compute_stats("P99 TTFT (ms)", runs, |r| r.p99_ttft_ms), + compute_stats("Mean TPOT (ms)", runs, |r| r.mean_tpot_ms), + compute_stats("Median TPOT (ms)", runs, |r| r.median_tpot_ms), + compute_stats("P99 TPOT (ms)", runs, |r| r.p99_tpot_ms), + compute_stats("Mean ITL (ms)", runs, |r| r.mean_itl_ms), + compute_stats("Mean E2EL (ms)", runs, |r| r.mean_e2el_ms), + compute_stats("P99 E2EL (ms)", runs, |r| r.p99_e2el_ms), + compute_stats("Completed requests", runs, |r| r.completed), + compute_stats("Failed requests", runs, |r| r.failed), + compute_stats("Duration (s)", runs, |r| r.duration), + ]; + + println!("{:=^80}", format!(" Multi-Run Summary ({n} runs) ")); + println!( + "{:<35} {:>10} {:>10} {:>10} {:>10}", + "Metric", "Mean", "Std", "Min", "Max" + ); + println!( + "{:-<35} {:->10} {:->10} {:->10} {:->10}", + "", "", "", "", "" + ); + + for s in &stats { + println!( + "{:<35} {:>10} {:>10} {:>10} {:>10}", + s.label, + fmt(s.mean), + fmt(s.std), + fmt(s.min), + fmt(s.max), + ); + } + + println!("{:=<80}", ""); + + // Print coefficient of variation for throughput + let tp = &stats[0]; // request_throughput + if tp.mean > 0.0 { + let cv = (tp.std / tp.mean) * 100.0; + println!( + "Throughput CV: {:.1}% (lower = more stable across runs)", + cv + ); + } + + // Steady-state summary — aggregate only runs that have a window. + let ss_stats = vec![ + compute_stats_opt("Request throughput (req/s)", runs, |r| { + r.ss_request_throughput + }), + compute_stats_opt("Output throughput (tok/s)", runs, |r| { + r.ss_output_throughput + }), + compute_stats_opt("Input throughput (tok/s)", runs, |r| r.ss_input_throughput), + compute_stats_opt("Total token throughput (tok/s)", runs, |r| { + r.ss_total_token_throughput + }), + compute_stats_opt("Mean TTFT (ms)", runs, |r| r.ss_mean_ttft_ms), + compute_stats_opt("Median TTFT (ms)", runs, |r| r.ss_median_ttft_ms), + compute_stats_opt("Mean TPOT (ms)", runs, |r| r.ss_mean_tpot_ms), + compute_stats_opt("Median TPOT (ms)", runs, |r| r.ss_median_tpot_ms), + compute_stats_opt("P90 TPOT (ms)", runs, |r| r.ss_p90_tpot_ms), + compute_stats_opt("P99 TPOT (ms)", runs, |r| r.ss_p99_tpot_ms), + ]; + + let k_present = ss_stats[0].n_present; + let n_total = ss_stats[0].n_total; + + if k_present == 0 { + println!(); + println!("Steady-state: no runs had a valid window"); + } else { + println!(); + println!( + "{:=^80}", + format!(" Steady-State Summary ({k_present}/{n_total} runs) ") + ); + if k_present < n_total { + println!( + "Aggregated over {k_present}/{n_total} runs ({} runs had no valid window)", + n_total - k_present + ); + } + println!( + "{:<35} {:>10} {:>10} {:>10} {:>10}", + "Metric", "Mean", "Std", "Min", "Max" + ); + for s in &ss_stats { + println!( + "{:<35} {:>10.2} {:>10.2} {:>10.2} {:>10.2}", + s.label, s.mean, s.std, s.min, s.max + ); + } + } +} + +fn compute_stats(label: &'static str, runs: &[RunMetrics], f: F) -> MetricStats +where + F: Fn(&RunMetrics) -> f64, +{ + let values: Vec = runs.iter().map(&f).collect(); + let n = values.len() as f64; + let mean = values.iter().sum::() / n; + let variance = values.iter().map(|x| (x - mean).powi(2)).sum::() / n; + let std = variance.sqrt(); + let min = values.iter().cloned().fold(f64::INFINITY, f64::min); + let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + + MetricStats { + label, + mean, + std, + min, + max, + } +} + +fn fmt(v: f64) -> String { + if v == 0.0 { + "0".to_string() + } else if v == v.floor() && v.abs() < 1e9 { + format!("{}", v as i64) + } else { + format!("{:.2}", v) + } +} + +/// Aggregated statistics over only the Some-valued runs. +struct MetricStatsOpt { + label: &'static str, + mean: f64, + std: f64, + min: f64, + max: f64, + n_present: usize, + n_total: usize, +} + +fn compute_stats_opt(label: &'static str, runs: &[RunMetrics], f: F) -> MetricStatsOpt +where + F: Fn(&RunMetrics) -> Option, +{ + let values: Vec = runs.iter().filter_map(&f).collect(); + let n_present = values.len(); + let n_total = runs.len(); + if values.is_empty() { + return MetricStatsOpt { + label, + mean: 0.0, + std: 0.0, + min: 0.0, + max: 0.0, + n_present, + n_total, + }; + } + let mean = values.iter().sum::() / n_present as f64; + let variance = values.iter().map(|x| (x - mean).powi(2)).sum::() / n_present as f64; + let std = variance.sqrt(); + let min = values.iter().cloned().fold(f64::INFINITY, f64::min); + let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max); + MetricStatsOpt { + label, + mean, + std, + min, + max, + n_present, + n_total, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn mk_run(ss: Option) -> RunMetrics { + RunMetrics { + request_throughput: 0.0, + output_throughput: 0.0, + total_token_throughput: 0.0, + mean_ttft_ms: 0.0, + median_ttft_ms: 0.0, + p99_ttft_ms: 0.0, + mean_tpot_ms: 0.0, + median_tpot_ms: 0.0, + p99_tpot_ms: 0.0, + mean_itl_ms: 0.0, + mean_e2el_ms: 0.0, + p99_e2el_ms: 0.0, + max_output_tokens_per_s: 0.0, + completed: 0.0, + failed: 0.0, + duration: 0.0, + ss_request_throughput: ss, + ss_output_throughput: None, + ss_input_throughput: None, + ss_total_token_throughput: None, + ss_mean_ttft_ms: None, + ss_median_ttft_ms: None, + ss_mean_tpot_ms: None, + ss_median_tpot_ms: None, + ss_p90_tpot_ms: None, + ss_p99_tpot_ms: None, + } + } + + #[test] + fn compute_stats_opt_excludes_none() { + let runs = vec![mk_run(Some(10.0)), mk_run(None), mk_run(Some(20.0))]; + let s = compute_stats_opt("x", &runs, |r| r.ss_request_throughput); + assert_eq!(s.n_present, 2); + assert_eq!(s.n_total, 3); + assert!((s.mean - 15.0).abs() < 1e-9); + assert_eq!(s.min, 10.0); + assert_eq!(s.max, 20.0); + } + + #[test] + fn compute_stats_opt_all_none() { + let runs = vec![mk_run(None), mk_run(None)]; + let s = compute_stats_opt("x", &runs, |r| r.ss_request_throughput); + assert_eq!(s.n_present, 0); + assert_eq!(s.n_total, 2); + assert_eq!(s.mean, 0.0); + } +} diff --git a/rust/src/bench/src/multi_turn.rs b/rust/src/bench/src/multi_turn.rs new file mode 100644 index 000000000000..8f0aedf31335 --- /dev/null +++ b/rust/src/bench/src/multi_turn.rs @@ -0,0 +1,827 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Instant; + +use indicatif::{ProgressBar, ProgressStyle}; +use tokio::sync::Semaphore; + +use crate::backends::{Backend, RequestFuncInput, RequestFuncOutput, get_backend}; +use crate::benchmark::{ + assign_lora_modules, compute_spec_decode_stats, fetch_spec_decode_metrics, pre_resolve_dns, + profile_on_batch_threshold, start_profiler_immediate, stop_profiler_immediate, +}; +use crate::cli::DatasetName; +use crate::config::BenchConfig; +use crate::datasets::MultiTurnConversation; +use crate::error::{BenchError, Result}; +use crate::metrics::calculator::calculate_multi_turn_metrics; +use crate::output::console::print_multi_turn_results; +use crate::output::json::{ + append_result, build_multi_turn_result_json, compute_result_filename, save_result, +}; + +/// Output from a single turn within a conversation. +#[derive(Debug, Clone)] +pub struct TurnOutput { + pub turn_index: usize, + pub request_output: RequestFuncOutput, + pub cumulative_input_tokens: usize, +} + +/// Output from an entire conversation. +#[derive(Debug, Clone)] +pub struct ConversationOutput { + pub conversation_id: String, + pub turns: Vec, + pub total_duration_ms: f64, + pub all_success: bool, +} + +/// Run the multi-turn conversation benchmark. +pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result { + let _ = get_backend(config.backend)?; + + // Build HTTP client + let concurrency = config + .multi_turn_concurrency + .unwrap_or(config.max_concurrency.unwrap_or(config.num_prompts)); + + let mut client_builder = reqwest::Client::builder() + .pool_max_idle_per_host(concurrency.max(256)) + .timeout(std::time::Duration::from_secs(6 * 60 * 60)) + .connect_timeout(std::time::Duration::from_secs(30)) + .tcp_keepalive(std::time::Duration::from_secs(60)) + .tcp_nodelay(true) + .http1_only() + .no_proxy(); + + if config.insecure { + client_builder = client_builder.danger_accept_invalid_certs(true); + } + + client_builder = pre_resolve_dns(&config.base_url, client_builder); + + let client = client_builder + .build() + .map_err(|e| BenchError::Backend(format!("Failed to build HTTP client: {e}")))?; + + // Resolve model + let (model_id, model_name) = if let Some(ref m) = config.model { + (m.clone(), config.model_name.clone()) + } else { + println!("Model not specified, fetching first model from server..."); + let (name, id) = get_first_model(&config.base_url, &client, &config.extra_headers).await?; + println!("First model name: {name}, first model id: {id}"); + (id, Some(name)) + }; + + // Load tokenizer + let tokenizer = if config.skip_tokenizer_init { + None + } else { + let tid = config.tokenizer_id.as_deref().unwrap_or(&model_id); + println!("Loading tokenizer: {tid}"); + let server_info = Some((config.base_url.as_str(), model_id.as_str())); + let t = crate::tokenizer::load_tokenizer(tid, config.trust_remote_code, server_info)?; + println!("Tokenizer loaded successfully."); + Some(t) + }; + + // Generate/load conversations + println!("Generating multi-turn conversations..."); + let gen_start = Instant::now(); + + let mut conversations = match config.dataset_name { + DatasetName::Random => { + let tok = tokenizer + .as_ref() + .ok_or_else(|| BenchError::Config("Random dataset requires a tokenizer".into()))?; + let prefix_sharing_config = if config.multi_turn_prefix_global_ratio > 0.0 + || config.multi_turn_prefix_conversation_ratio > 0.0 + { + Some(crate::datasets::multi_turn::PrefixSharingConfig { + global_ratio: config.multi_turn_prefix_global_ratio, + conversation_ratio: config.multi_turn_prefix_conversation_ratio, + }) + } else { + None + }; + let random_cfg = crate::datasets::multi_turn::MultiTurnRandomConfig { + num_conversations: config.num_prompts, + min_turns: config.multi_turn_min_turns, + max_turns: config.multi_turn_max_turns, + prefix_len: config.random_prefix_len, + input_len: config.random_input_len, + per_turn_input_len: config.per_turn_input_len, + output_len: config.random_output_len, + seed: config.seed, + request_id_prefix: config.request_id_prefix.clone(), + prefix_sharing_config, + }; + crate::datasets::multi_turn::generate_multi_turn_random(tok, &random_cfg)? + } + DatasetName::ShareGpt => { + let tok = tokenizer.as_ref().ok_or_else(|| { + BenchError::Config("ShareGPT dataset requires a tokenizer".into()) + })?; + let downloaded; + let path = match config.dataset_path.as_deref() { + Some(p) => p, + None => { + downloaded = crate::datasets::sharegpt::download_sharegpt_dataset()?; + downloaded.as_str() + } + }; + crate::datasets::multi_turn::load_sharegpt_multi_turn( + tok, + path, + config.num_prompts, + config.sharegpt_output_len, + config.sharegpt_multi_turn_max_turns, + config.seed, + &config.request_id_prefix, + )? + } + DatasetName::RandomMm => { + return Err(BenchError::Config( + "Random-MM multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), + )); + } + DatasetName::Sonnet => { + return Err(BenchError::Config( + "Sonnet multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), + )); + } + DatasetName::SpeedBench => { + return Err(BenchError::Config( + "SPEED-Bench multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), + )); + } + DatasetName::Custom | DatasetName::PrefixRepetition | DatasetName::RandomRerank => { + return Err(BenchError::Config( + "This dataset does not support multi-turn. Use 'random' or 'sharegpt' with --multi-turn.".into(), + )); + } + DatasetName::Hf => { + return Err(BenchError::Config( + "HF dataset multi-turn is not yet supported. Use 'random' or 'sharegpt' with --multi-turn.".into(), + )); + } + }; + + let no_history = config.multi_turn_prefix_global_ratio > 0.0 + || config.multi_turn_prefix_conversation_ratio > 0.0; + + if let Some(max_model_len) = config.max_model_len { + let (filtered_conversations, filtered_turns) = + filter_turns_by_max_model_len(&mut conversations, max_model_len, no_history); + if filtered_turns > 0 || filtered_conversations > 0 { + println!( + "Filtered {filtered_turns} turn(s) and {filtered_conversations} conversation(s) above --max-model-len {max_model_len}." + ); + } + if conversations.is_empty() { + return Err(BenchError::Config( + "No conversations remain after applying --max-model-len".into(), + )); + } + } + + let gen_elapsed = gen_start.elapsed(); + let total_turns: usize = conversations.iter().map(|c| c.turns.len()).sum(); + println!( + "Generated {} conversations ({} total turns) in {:.2}s", + conversations.len(), + total_turns, + gen_elapsed.as_secs_f64() + ); + + // Log prefix sharing info + if no_history { + let num_special = tokenizer.as_ref().map(|t| t.num_special_tokens_to_add()).unwrap_or(0); + let real_input_len = config.random_input_len.saturating_sub(num_special); + let global_tokens = + (real_input_len as f64 * config.multi_turn_prefix_global_ratio).floor() as usize; + let conv_tokens = + (real_input_len as f64 * config.multi_turn_prefix_conversation_ratio).floor() as usize; + let unique_tokens = real_input_len.saturating_sub(global_tokens + conv_tokens); + println!( + "User message prefix sharing: {:.0}% global ({} tokens), {:.0}% per-conversation ({} tokens), {:.0}% unique ({} tokens)", + config.multi_turn_prefix_global_ratio * 100.0, + global_tokens, + config.multi_turn_prefix_conversation_ratio * 100.0, + conv_tokens, + (1.0 - config.multi_turn_prefix_global_ratio + - config.multi_turn_prefix_conversation_ratio) + * 100.0, + unique_tokens, + ); + println!("No history accumulation: each turn sends fixed-length prompt only."); + } + + if config.dry_run { + let total_user_tokens: usize = conversations + .iter() + .flat_map(|c| c.turns.iter()) + .map(|t| t.user_message_len) + .sum(); + println!("Dry run stats:"); + println!(" Total conversations: {}", conversations.len()); + println!(" Total turns: {total_turns}"); + println!(" Total user message tokens: {total_user_tokens}"); + return Ok(serde_json::json!({"dry_run": true, "mode": "multi_turn"})); + } + + // Ready check with a simple single request + if config.ready_check_timeout_sec > 0 { + let first_turn = &conversations[0].turns[0]; + let test_input = RequestFuncInput { + prompt: first_turn.user_message.clone(), + api_url: config.api_url.clone(), + prompt_len: first_turn.user_message_len, + output_len: first_turn.expected_output_len, + model: model_id.clone(), + model_name: model_name.clone(), + logprobs: config.logprobs, + extra_headers: config.extra_headers.clone(), + extra_body: config.extra_body.clone(), + ignore_eos: config.ignore_eos, + request_id: None, + ..Default::default() + }; + + println!("Starting initial single prompt test run..."); + let test_output = crate::ready_checker::wait_for_endpoint( + config.backend, + &client, + &test_input, + config.ready_check_timeout_sec, + 5, + ) + .await?; + if !test_output.success { + return Err(BenchError::Backend(format!( + "Initial test run failed: {}", + test_output.error + ))); + } + println!("Initial test run completed."); + } + + // For random datasets in multi-turn mode, auto-set min_tokens to enforce + // output length without using ignore_eos (which causes unbounded context growth). + // min_tokens + max_completion_tokens together control output length precisely. + let extra_body = if config.dataset_name == DatasetName::Random && !config.ignore_eos { + let mut body = config.extra_body.clone().unwrap_or_else(|| serde_json::json!({})); + if let serde_json::Value::Object(ref mut map) = body + && !map.contains_key("min_tokens") + { + map.insert( + "min_tokens".to_string(), + serde_json::json!(config.random_output_len), + ); + println!( + "Auto-setting min_tokens={} for multi-turn random dataset (use --extra-body to override)", + config.random_output_len + ); + } + Some(body) + } else { + config.extra_body.clone() + }; + + // Fetch speculative decoding metrics before benchmark + let spec_decode_before = + fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await; + if spec_decode_before.is_some() { + println!("Speculative decoding detected, will collect metrics."); + } + + // Start profiler if requested (immediate mode — no batch threshold) + if config.profile && config.profile_batch_threshold.is_none() { + start_profiler_immediate(&client, &config.base_url, &config.extra_headers).await; + } + + // Threshold-based profiling: spawn background task that polls /metrics + // and triggers start/stop profile when batch size is reached. + let profile_task = if let Some(threshold) = config.profile_batch_threshold { + let poll_client = client.clone(); + let base_url = config.base_url.clone(); + let extra_headers = config.extra_headers.clone(); + let duration_secs = config.profile_duration; + let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel::<()>(); + let handle = tokio::spawn(async move { + profile_on_batch_threshold( + &poll_client, + &base_url, + &extra_headers, + threshold, + duration_secs, + cancel_rx, + ) + .await; + }); + Some((cancel_tx, handle)) + } else { + None + }; + + // Main benchmark + println!("Starting multi-turn benchmark..."); + println!("Conversations: {}", conversations.len()); + println!("Concurrency: {concurrency}"); + println!("Inter-turn delay: {} ms", config.multi_turn_delay_ms); + + let max_turn_count = conversations.iter().map(|c| c.turns.len()).max().unwrap_or(0); + + // Progress bar counts total turns + let pb = if config.disable_tqdm { + None + } else { + let bar = ProgressBar::new(total_turns as u64); + bar.set_style( + ProgressStyle::with_template( + "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} turns ({eta})", + ) + .unwrap() + .progress_chars("#>-"), + ); + Some(bar) + }; + + let benchmark_start = Instant::now(); + + // Per-conversation LoRA assignment (sticky: all turns of a conversation + // share the same adapter). None when --lora-modules not set. + let lora_assignments = assign_lora_modules( + &config.lora_modules, + config.lora_assignment, + conversations.len(), + config.seed, + ); + if let Some(modules) = config.lora_modules.as_ref() { + let names: Vec<&str> = modules.iter().map(|s| s.as_ref()).collect(); + println!( + "LoRA adapters ({}): {:?} [assignment={:?}, scope=conversation]", + modules.len(), + names, + config.lora_assignment + ); + } + + // Shared state + let backend = get_backend(config.backend)?; + let api_url = Arc::new(config.api_url.clone()); + let model = Arc::new(model_id.clone()); + let model_name_arc = Arc::new(model_name.clone()); + let extra_body = Arc::new(extra_body); + let base_extra_headers = Arc::new(config.extra_headers.clone()); + let ignore_eos = config.ignore_eos; + let logprobs = config.logprobs; + let delay_ms = config.multi_turn_delay_ms; + + // Semaphore controls max in-flight requests (not conversations). + // Each conversation acquires the permit only during the HTTP request, + // releasing it before the inter-turn delay so other conversations can + // immediately fill the slot. + let semaphore = Arc::new(Semaphore::new(concurrency)); + + // Spawn one task per conversation (all at once) + let mut handles = Vec::with_capacity(conversations.len()); + for (i, conv) in conversations.into_iter().enumerate() { + let client = client.clone(); + let backend = backend.clone(); + let api_url = api_url.clone(); + let model = model.clone(); + let model_name = model_name_arc.clone(); + let extra_body = extra_body.clone(); + let base_extra_headers = base_extra_headers.clone(); + let pb = pb.clone(); + let bench_start = benchmark_start; + let semaphore = semaphore.clone(); + let lora_name = lora_assignments.as_ref().map(|v| v[i].clone()); + + handles.push(tokio::spawn(async move { + run_conversation( + &conv, + &backend, + &client, + &api_url, + &model, + &model_name, + lora_name.as_deref(), + &extra_body, + &base_extra_headers, + ignore_eos, + logprobs, + delay_ms, + no_history, + bench_start, + pb.as_ref(), + &semaphore, + ) + .await + })); + } + + // Collect all conversation outputs + let mut all_outputs: Vec = Vec::with_capacity(handles.len()); + for handle in handles { + match handle.await { + Ok(output) => all_outputs.push(output), + Err(e) => { + eprintln!("Conversation task panicked: {e}"); + } + } + } + + if let Some(ref pb) = pb { + pb.finish_and_clear(); + } + + let benchmark_duration = benchmark_start.elapsed().as_secs_f64(); + + // Stop profiler if requested (immediate mode — no batch threshold) + if config.profile && config.profile_batch_threshold.is_none() { + stop_profiler_immediate(&client, &config.base_url, &config.extra_headers).await; + } + + // Signal the threshold-based profile task that the benchmark is done, then wait + if let Some((cancel_tx, task)) = profile_task { + let _ = cancel_tx.send(()); + if let Err(e) = task.await { + eprintln!("WARNING: Profile background task failed: {e}"); + } + } + + // Fetch speculative decoding metrics after benchmark and compute stats + let spec_decode_stats = if spec_decode_before.is_some() { + let spec_decode_after = + fetch_spec_decode_metrics(&config.base_url, &client, &config.extra_headers).await; + match (spec_decode_before.as_ref(), spec_decode_after.as_ref()) { + (Some(before), Some(after)) => compute_spec_decode_stats(before, after), + _ => None, + } + } else { + None + }; + + // Calculate metrics + let mt_metrics = calculate_multi_turn_metrics( + &all_outputs, + benchmark_duration, + &config.selected_percentiles, + &config.goodput, + max_turn_count, + ); + + // Print console output + print_multi_turn_results( + &mt_metrics, + benchmark_duration, + config, + spec_decode_stats.as_ref(), + ); + + // Build result JSON + let date_iso = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%SZ").to_string(); + let dt_filename = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string(); + let result_json = build_multi_turn_result_json( + config, + &mt_metrics, + &all_outputs, + benchmark_duration, + &date_iso, + spec_decode_stats.as_ref(), + ); + + // Save if requested + if config.save_result || config.append_result { + let model_for_filename = config.model.as_deref().unwrap_or(&model_id); + let file_name = compute_result_filename(config, model_for_filename, &dt_filename); + + if let Some(ref dir) = config.result_dir { + std::fs::create_dir_all(dir)?; + } + + if config.append_result { + append_result(&result_json, &file_name)?; + } else { + save_result(&result_json, &file_name)?; + } + } + + Ok(result_json) +} + +fn filter_turns_by_max_model_len( + conversations: &mut Vec, + max_model_len: usize, + no_history: bool, +) -> (usize, usize) { + let before_conversations = conversations.len(); + let before_turns: usize = conversations.iter().map(|c| c.turns.len()).sum(); + + for conversation in conversations.iter_mut() { + if no_history { + conversation.turns.retain(|turn| { + turn.user_message_len.saturating_add(turn.expected_output_len) <= max_model_len + }); + continue; + } + + let keep_turns = valid_prefix_len_for_max_model_len(conversation, max_model_len); + conversation.turns.truncate(keep_turns); + } + + conversations.retain(|conversation| conversation.turns.len() >= 2); + + let after_conversations = conversations.len(); + let after_turns: usize = conversations.iter().map(|c| c.turns.len()).sum(); + + ( + before_conversations - after_conversations, + before_turns - after_turns, + ) +} + +fn valid_prefix_len_for_max_model_len( + conversation: &MultiTurnConversation, + max_model_len: usize, +) -> usize { + let mut cumulative_input_tokens = 0usize; + let mut keep_turns = 0usize; + + for turn in &conversation.turns { + cumulative_input_tokens = cumulative_input_tokens.saturating_add(turn.user_message_len); + + if cumulative_input_tokens.saturating_add(turn.expected_output_len) > max_model_len { + break; + } + + cumulative_input_tokens = cumulative_input_tokens.saturating_add(turn.expected_output_len); + keep_turns += 1; + } + + keep_turns +} + +/// Run a single conversation: sequential turns, building up message history. +/// +/// `lora_name`, when set, replaces both `model` and `model_name` in every +/// turn's request payload — sticky for the whole conversation so prefix-cache +/// reuse across turns isn't broken by mid-conversation adapter switches. +async fn run_conversation( + conversation: &MultiTurnConversation, + backend: &Backend, + client: &reqwest::Client, + api_url: &str, + model: &str, + model_name: &Option, + lora_name: Option<&str>, + extra_body: &Option, + base_extra_headers: &Option>, + ignore_eos: bool, + logprobs: Option, + delay_ms: u64, + no_history: bool, + bench_start: Instant, + pb: Option<&ProgressBar>, + semaphore: &Semaphore, +) -> ConversationOutput { + let conv_start = Instant::now(); + let mut messages: Vec = Vec::new(); + let mut cumulative_tokens: usize = 0; + let mut turn_outputs: Vec = Vec::new(); + let mut all_success = true; + + // Router affinity: all turns share same X-Session-ID + let mut extra_headers = base_extra_headers.clone().unwrap_or_default(); + extra_headers.insert( + "X-Session-ID".to_string(), + conversation.conversation_id.clone(), + ); + + for (turn_idx, turn) in conversation.turns.iter().enumerate() { + if no_history { + // Prefix sharing mode: reset messages each turn, send only this turn's message + messages.clear(); + messages.push(serde_json::json!({ + "role": "user", + "content": [{"type": "text", "text": &*turn.user_message}] + })); + cumulative_tokens = turn.user_message_len; + } else { + // Normal mode: accumulate history + messages.push(serde_json::json!({ + "role": "user", + "content": [{"type": "text", "text": &*turn.user_message}] + })); + cumulative_tokens += turn.user_message_len; + } + + // Set min_tokens per-request so each turn's output length matches + // the dataset's expected length. Skip if ignore_eos is set (unbounded + // generation) or the user already provided min_tokens via --extra-body. + let turn_extra_body = if !ignore_eos { + let mut body = extra_body.clone().unwrap_or_else(|| serde_json::json!({})); + if let serde_json::Value::Object(ref mut map) = body + && !map.contains_key("min_tokens") + { + map.insert( + "min_tokens".to_string(), + serde_json::json!(turn.expected_output_len), + ); + } + Some(body) + } else { + extra_body.clone() + }; + + let (req_model, req_model_name) = match lora_name { + Some(name) => (name.to_string(), Some(name.to_string())), + None => (model.to_string(), model_name.clone()), + }; + + let input = RequestFuncInput { + prompt: turn.user_message.clone(), + api_url: api_url.to_string(), + prompt_len: cumulative_tokens, + output_len: turn.expected_output_len, + model: req_model, + model_name: req_model_name, + logprobs, + extra_headers: Some(extra_headers.clone()), + extra_body: turn_extra_body, + ignore_eos, + request_id: Some(format!("{}-turn{}", conversation.conversation_id, turn_idx)), + messages: Some(serde_json::json!(messages)), + ..Default::default() + }; + + // Acquire semaphore permit before sending the request. + // Released after the response so the slot is free during inter-turn delay. + let _permit = match semaphore.acquire().await { + Ok(p) => p, + Err(_) => { + all_success = false; + break; + } + }; + let request_instant = Instant::now(); + let result = backend.send_request(&input, client).await; + drop(_permit); + + let output = match result { + Ok(mut o) => { + o.start_time = request_instant.duration_since(bench_start).as_secs_f64(); + o + } + Err(e) => RequestFuncOutput { + success: false, + error: e.to_string(), + prompt_len: cumulative_tokens, + start_time: request_instant.duration_since(bench_start).as_secs_f64(), + ..Default::default() + }, + }; + + if let Some(pb) = pb { + pb.inc(1); + } + + if output.success { + if !no_history { + // Add assistant response to history for next turn + messages.push(serde_json::json!({ + "role": "assistant", + "content": &output.generated_text + })); + cumulative_tokens += output.output_tokens; + } + + turn_outputs.push(TurnOutput { + turn_index: turn_idx, + request_output: output, + cumulative_input_tokens: cumulative_tokens, + }); + } else { + all_success = false; + turn_outputs.push(TurnOutput { + turn_index: turn_idx, + request_output: output, + cumulative_input_tokens: cumulative_tokens, + }); + break; // Conversation stops on failure + } + + // Inter-turn delay (semaphore NOT held — slot is free for others) + if delay_ms > 0 && turn_idx + 1 < conversation.turns.len() { + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + } + } + + let total_duration_ms = conv_start.elapsed().as_secs_f64() * 1000.0; + + ConversationOutput { + conversation_id: conversation.conversation_id.clone(), + turns: turn_outputs, + total_duration_ms, + all_success, + } +} + +/// Fetch the first model from the server's /v1/models endpoint. +async fn get_first_model( + base_url: &str, + client: &reqwest::Client, + extra_headers: &Option>, +) -> Result<(String, String)> { + let url = format!("{base_url}/v1/models"); + let mut request = client.get(&url); + if let Some(headers) = extra_headers { + for (k, v) in headers { + request = request.header(k, v); + } + } + if let Ok(api_key) = std::env::var("OPENAI_API_KEY") { + request = request.header("Authorization", format!("Bearer {api_key}")); + } + + let response = request.send().await?; + let data: serde_json::Value = response.json().await?; + + if let Some(models) = data.get("data").and_then(|d| d.as_array()) + && let Some(first) = models.first() + { + let id = first.get("id").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + let root = first.get("root").and_then(|v| v.as_str()).unwrap_or(&id).to_string(); + return Ok((id, root)); + } + + Err(BenchError::Config(format!( + "No models found on the server at {base_url}" + ))) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::{filter_turns_by_max_model_len, valid_prefix_len_for_max_model_len}; + use crate::datasets::{ConversationTurn, MultiTurnConversation}; + + fn conversation(turns: &[(usize, usize)]) -> MultiTurnConversation { + MultiTurnConversation { + conversation_id: "conv-0".to_string(), + turns: turns + .iter() + .map(|(user_message_len, expected_output_len)| ConversationTurn { + user_message: Arc::from("hello"), + user_message_len: *user_message_len, + expected_output_len: *expected_output_len, + }) + .collect(), + } + } + + #[test] + fn test_valid_prefix_len_for_max_model_len_with_history() { + let conv = conversation(&[(40, 10), (45, 10)]); + + assert_eq!(valid_prefix_len_for_max_model_len(&conv, 105), 2); + assert_eq!(valid_prefix_len_for_max_model_len(&conv, 104), 1); + } + + #[test] + fn test_filter_turns_by_max_model_len_truncates_history_conversations() { + let mut conversations = vec![ + conversation(&[(40, 10), (45, 10), (1, 1)]), + conversation(&[(100, 10), (1, 1)]), + ]; + + let (filtered_conversations, filtered_turns) = + filter_turns_by_max_model_len(&mut conversations, 105, false); + + assert_eq!(filtered_conversations, 1); + assert_eq!(filtered_turns, 3); + assert_eq!(conversations.len(), 1); + assert_eq!(conversations[0].turns.len(), 2); + } + + #[test] + fn test_filter_turns_by_max_model_len_retains_independent_no_history_turns() { + let mut conversations = vec![conversation(&[(40, 10), (95, 10), (1, 1)])]; + + let (filtered_conversations, filtered_turns) = + filter_turns_by_max_model_len(&mut conversations, 104, true); + + assert_eq!(filtered_conversations, 0); + assert_eq!(filtered_turns, 1); + assert_eq!(conversations.len(), 1); + assert_eq!(conversations[0].turns.len(), 2); + } +} diff --git a/rust/src/bench/src/output/console.rs b/rust/src/bench/src/output/console.rs new file mode 100644 index 000000000000..ae0184e61541 --- /dev/null +++ b/rust/src/bench/src/output/console.rs @@ -0,0 +1,400 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use crate::benchmark::SpecDecodeStats; +use crate::config::BenchConfig; +use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics}; + +/// Print benchmark results to console in the same format as Python. +/// +/// Mirrors the output format from serve.py:918-1093. +pub fn print_results( + metrics: &BenchmarkMetrics, + benchmark_duration: f64, + config: &BenchConfig, + has_tokenizer: bool, + spec_decode_stats: Option<&SpecDecodeStats>, +) { + let is_pooling = config.backend.is_pooling(); + + if is_pooling { + println!("{:=^60}", " Embedding/Pooling Benchmark Result "); + } else { + println!("{:=^60}", " Serving Benchmark Result "); + } + println!("{:<40} {:<10}", "Successful requests:", metrics.completed); + println!("{:<40} {:<10}", "Failed requests:", metrics.failed); + + if let Some(mc) = config.max_concurrency { + println!("{:<40} {:<10}", "Maximum request concurrency:", mc); + } + if !config.request_rate.is_infinite() { + println!( + "{:<40} {:<10.2}", + "Request rate configured (RPS):", config.request_rate + ); + } + println!( + "{:<40} {:<10.2}", + "Benchmark duration (s):", benchmark_duration + ); + println!("{:<40} {:<10}", "Total input tokens:", metrics.total_input); + + if has_tokenizer && !is_pooling { + println!( + "{:<40} {:<10}", + "Total generated tokens:", metrics.total_output + ); + } + + println!( + "{:<40} {:<10.2}", + "Request throughput (req/s):", metrics.request_throughput + ); + if metrics.request_goodput > 0.0 { + println!( + "{:<40} {:<10.2}", + "Request goodput (req/s):", metrics.request_goodput + ); + } + + if !is_pooling { + if has_tokenizer { + println!( + "{:<40} {:<10.2}", + "Output token throughput (tok/s):", metrics.output_throughput + ); + println!( + "{:<40} {:<10.2}", + "Peak output token throughput (tok/s):", metrics.max_output_tokens_per_s + ); + } + + println!( + "{:<40} {:<10.2}", + "Peak concurrent requests:", metrics.max_concurrent_requests as f64 + ); + + if has_tokenizer { + println!( + "{:<40} {:<10.2}", + "Total token throughput (tok/s):", metrics.total_token_throughput + ); + } + } else { + println!( + "{:<40} {:<10.2}", + "Input token throughput (tok/s):", metrics.input_throughput + ); + println!( + "{:<40} {:<10.2}", + "Peak concurrent requests:", metrics.max_concurrent_requests as f64 + ); + } + + // Print per-metric percentiles + if has_tokenizer && !is_pooling { + print_metric_section( + "ttft", + "TTFT", + "Time to First Token", + &config.selected_percentile_metrics, + metrics.mean_ttft_ms, + metrics.median_ttft_ms, + &metrics.percentiles_ttft_ms, + ); + print_metric_section( + "tpot", + "TPOT", + "Time per Output Token (excl. 1st token)", + &config.selected_percentile_metrics, + metrics.mean_tpot_ms, + metrics.median_tpot_ms, + &metrics.percentiles_tpot_ms, + ); + print_metric_section( + "itl", + "ITL", + "Inter-token Latency", + &config.selected_percentile_metrics, + metrics.mean_itl_ms, + metrics.median_itl_ms, + &metrics.percentiles_itl_ms, + ); + } + print_metric_section( + "e2el", + "E2EL", + "End-to-end Latency", + &config.selected_percentile_metrics, + metrics.mean_e2el_ms, + metrics.median_e2el_ms, + &metrics.percentiles_e2el_ms, + ); + + if let Some(stats) = spec_decode_stats { + print_spec_decode_section(stats); + } + + println!("{:=<60}", ""); + + print_steady_state(metrics); +} + +/// Print multi-turn benchmark results to console. +pub fn print_multi_turn_results( + mt_metrics: &MultiTurnMetrics, + benchmark_duration: f64, + config: &BenchConfig, + spec_decode_stats: Option<&SpecDecodeStats>, +) { + println!("{:=^60}", " Multi-Turn Benchmark Result "); + println!( + "{:<45} {}/{}", + "Conversations completed/total:", + mt_metrics.conversations_completed, + mt_metrics.conversations_completed + mt_metrics.conversations_failed, + ); + println!( + "{:<45} {}", + "Turns per conversation (configured):", config.multi_turn_num_turns, + ); + println!( + "{:<45} {:.1}", + "Avg turns completed:", mt_metrics.avg_turns_completed, + ); + println!( + "{:<45} {:.0}", + "Avg conversation duration (ms):", mt_metrics.avg_conversation_duration_ms, + ); + println!( + "{:<45} {}", + "Concurrency:", + config + .multi_turn_concurrency + .unwrap_or(config.max_concurrency.unwrap_or(config.num_prompts)), + ); + println!( + "{:<45} {}", + "Inter-turn delay (ms):", config.multi_turn_delay_ms, + ); + println!( + "{:<45} {:.2}", + "Benchmark duration (s):", benchmark_duration, + ); + + // Overall metrics + println!("{:-^60}", " Overall (All Turns) "); + print_metrics_block(&mt_metrics.overall, config); + + // Per-turn breakdown + for (i, turn_metrics) in mt_metrics.per_turn.iter().enumerate() { + let samples = turn_metrics.completed + turn_metrics.failed; + println!("{:-^60}", format!(" Turn {} ({} samples) ", i + 1, samples)); + print_metrics_block(turn_metrics, config); + } + + if let Some(stats) = spec_decode_stats { + print_spec_decode_section(stats); + } + + println!("{:=<60}", ""); +} + +/// Print a metrics block (reused for overall and per-turn). +fn print_metrics_block(metrics: &BenchmarkMetrics, config: &BenchConfig) { + println!("{:<45} {:<10}", "Successful requests:", metrics.completed); + println!("{:<45} {:<10}", "Failed requests:", metrics.failed); + println!("{:<45} {:<10}", "Total input tokens:", metrics.total_input); + println!( + "{:<45} {:<10}", + "Total generated tokens:", metrics.total_output + ); + println!( + "{:<45} {:<10.2}", + "Request throughput (req/s):", metrics.request_throughput + ); + println!( + "{:<45} {:<10.2}", + "Input token throughput (tok/s):", metrics.input_throughput + ); + println!( + "{:<45} {:<10.2}", + "Output token throughput (tok/s):", metrics.output_throughput + ); + println!( + "{:<45} {:<10.2}", + "Total token throughput (tok/s):", metrics.total_token_throughput + ); + + print_metric_section( + "ttft", + "TTFT", + "Time to First Token", + &config.selected_percentile_metrics, + metrics.mean_ttft_ms, + metrics.median_ttft_ms, + &metrics.percentiles_ttft_ms, + ); + print_metric_section( + "tpot", + "TPOT", + "Time per Output Token (excl. 1st token)", + &config.selected_percentile_metrics, + metrics.mean_tpot_ms, + metrics.median_tpot_ms, + &metrics.percentiles_tpot_ms, + ); + print_metric_section( + "itl", + "ITL", + "Inter-token Latency", + &config.selected_percentile_metrics, + metrics.mean_itl_ms, + metrics.median_itl_ms, + &metrics.percentiles_itl_ms, + ); + print_metric_section( + "e2el", + "E2EL", + "End-to-end Latency", + &config.selected_percentile_metrics, + metrics.mean_e2el_ms, + metrics.median_e2el_ms, + &metrics.percentiles_e2el_ms, + ); +} + +fn print_metric_section( + attr_name: &str, + short_name: &str, + header: &str, + selected: &[String], + mean_ms: f64, + median_ms: f64, + percentiles: &[(f64, f64)], +) { + if !selected.iter().any(|s| s == attr_name) { + return; + } + println!("{:-^60}", header); + println!( + "{:<40} {:<10.2}", + format!("Mean {short_name} (ms):"), + mean_ms + ); + println!( + "{:<40} {:<10.2}", + format!("Median {short_name} (ms):"), + median_ms + ); + for (p, value) in percentiles { + let p_str = if *p == p.floor() { + format!("{}", *p as i64) + } else { + format!("{p}") + }; + println!( + "{:<40} {:<10.2}", + format!("P{p_str} {short_name} (ms):"), + value + ); + } +} + +/// Print speculative decoding metrics section. +fn print_spec_decode_section(stats: &SpecDecodeStats) { + println!("{:-^50}", "Speculative Decoding"); + println!( + "{:<40} {:<10.2}", + "Acceptance rate (%):", stats.acceptance_rate + ); + println!( + "{:<40} {:<10.2}", + "Acceptance length:", stats.acceptance_length + ); + println!("{:<40} {:<10}", "Drafts:", stats.num_drafts); + println!("{:<40} {:<10}", "Draft tokens:", stats.draft_tokens); + println!("{:<40} {:<10}", "Accepted tokens:", stats.accepted_tokens); + if !stats.per_position_acceptance_rates.is_empty() { + println!("Per-position acceptance (%):"); + for (i, rate) in stats.per_position_acceptance_rates.iter().enumerate() { + println!("{:<40} {:<10.2}", format!(" Position {i}:"), rate * 100.0); + } + } +} + +fn print_steady_state(metrics: &crate::metrics::BenchmarkMetrics) { + let Some(ss) = metrics.steady_state.as_ref() else { + return; + }; + + if let Some(warning) = ss.window.warning.as_ref() { + println!("Warning: {warning}"); + } + + let total = ss.window.requests_total.max(1); + let started_pct = 100.0 * ss.window.requests_started_in_window as f64 / total as f64; + + println!("{:=^47}", " Steady-State Metrics "); + println!( + "{:<35} >= {:.2} * {} = {}", + "Concurrency threshold:", + ss.window.threshold, + ss.window.target_concurrency, + ss.window.threshold_abs, + ); + println!( + "{:<35} {:.1}s -> {:.1}s ({:.1}s)", + "Window:", ss.window.start_s, ss.window.end_s, ss.window.duration_s + ); + println!( + "{:<35} {}", + "Observed peak concurrency:", ss.window.observed_peak + ); + println!( + "{:<35} {} / {} ({:.1}%)", + "Requests started in window:", + ss.window.requests_started_in_window, + ss.window.requests_total, + started_pct + ); + println!( + "{:<35} {}", + "Requests completed in window:", ss.window.requests_completed_in_window + ); + println!("{:-<47}", ""); + println!( + "{:<35} {:.2}", + "Request throughput (req/s):", ss.request_throughput + ); + println!( + "{:<35} {:.2}", + "Output token throughput (tok/s):", ss.output_throughput + ); + println!( + "{:<35} {:.2}", + "Input token throughput (tok/s):", ss.input_throughput + ); + println!( + "{:<35} {:.2}", + "Total token throughput (tok/s):", ss.total_token_throughput + ); + + if !ss.percentiles_ttft_ms.is_empty() || ss.mean_ttft_ms > 0.0 { + println!("{:-<47}", ""); + println!("{:<35} {:.2}", "Mean TTFT (ms):", ss.mean_ttft_ms); + println!("{:<35} {:.2}", "Median TTFT (ms):", ss.median_ttft_ms); + for (p, v) in &ss.percentiles_ttft_ms { + println!("{:<35} {:.2}", format!("P{} TTFT (ms):", *p as u32), v); + } + } + if ss.mean_tpot_ms > 0.0 || ss.median_tpot_ms > 0.0 { + println!("{:-<47}", ""); + println!("{:<35} {:.2}", "Mean TPOT (ms):", ss.mean_tpot_ms); + println!("{:<35} {:.2}", "Median TPOT (ms):", ss.median_tpot_ms); + println!("{:<35} {:.2}", "P90 TPOT (ms):", ss.p90_tpot_ms); + println!("{:<35} {:.2}", "P99 TPOT (ms):", ss.p99_tpot_ms); + } + println!("{:=^47}", ""); +} diff --git a/rust/src/bench/src/output/json.rs b/rust/src/bench/src/output/json.rs new file mode 100644 index 000000000000..36ec07f65f10 --- /dev/null +++ b/rust/src/bench/src/output/json.rs @@ -0,0 +1,776 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use serde_json::Value; + +use crate::backends::RequestFuncOutput; +use crate::benchmark::SpecDecodeStats; +use crate::config::BenchConfig; +use crate::error::Result; +use crate::metrics::{BenchmarkMetrics, MultiTurnMetrics}; +use crate::multi_turn::ConversationOutput; + +/// Build the result JSON object matching the Python output schema exactly. +/// +/// Mirrors serve.py:1801-1947 result_json construction. +pub fn build_result_json( + config: &BenchConfig, + metrics: &BenchmarkMetrics, + actual_output_lens: &[usize], + outputs: &[RequestFuncOutput], + benchmark_duration: f64, + current_dt: &str, + spec_decode_stats: Option<&SpecDecodeStats>, +) -> serde_json::Value { + let mut result = serde_json::Map::new(); + + // Setup + result.insert("date".into(), Value::String(current_dt.to_string())); + result.insert( + "endpoint_type".into(), + Value::String(config.backend.to_string()), + ); + result.insert("backend".into(), Value::String(config.backend.to_string())); + result.insert( + "label".into(), + config.label.as_ref().map(|l| Value::String(l.clone())).unwrap_or(Value::Null), + ); + result.insert( + "model_id".into(), + Value::String(config.model.clone().unwrap_or_default()), + ); + result.insert( + "tokenizer_id".into(), + config + .tokenizer_id + .as_ref() + .map(|t| Value::String(t.clone())) + .unwrap_or(Value::Null), + ); + result.insert( + "num_prompts".into(), + Value::Number(config.num_prompts.into()), + ); + result.insert( + "max_model_len".into(), + config.max_model_len.map(|v| Value::Number(v.into())).unwrap_or(Value::Null), + ); + + // Metadata + if let Some(ref metadata) = config.metadata { + for (k, v) in metadata { + result.insert(k.clone(), Value::String(v.clone())); + } + } + + // Traffic + if config.request_rate.is_infinite() { + result.insert("request_rate".into(), Value::String("inf".to_string())); + } else { + result.insert( + "request_rate".into(), + serde_json::json!(config.request_rate), + ); + } + result.insert("burstiness".into(), serde_json::json!(config.burstiness)); + result.insert( + "max_concurrency".into(), + config.max_concurrency.map(|v| Value::Number(v.into())).unwrap_or(Value::Null), + ); + + let is_pooling = config.backend.is_pooling(); + + // Benchmark results + result.insert("duration".into(), serde_json::json!(benchmark_duration)); + result.insert("completed".into(), serde_json::json!(metrics.completed)); + result.insert("failed".into(), serde_json::json!(metrics.failed)); + result.insert( + "total_input_tokens".into(), + serde_json::json!(metrics.total_input), + ); + if !is_pooling { + result.insert( + "total_output_tokens".into(), + serde_json::json!(metrics.total_output), + ); + } + result.insert( + "request_throughput".into(), + serde_json::json!(metrics.request_throughput), + ); + if config.goodput.is_empty() { + result.insert("request_goodput".into(), Value::Null); + } else { + result.insert( + "request_goodput".into(), + serde_json::json!(metrics.request_goodput), + ); + } + if is_pooling { + result.insert( + "total_token_throughput".into(), + serde_json::json!(metrics.total_token_throughput), + ); + } else { + result.insert( + "input_throughput".into(), + serde_json::json!(metrics.input_throughput), + ); + result.insert( + "output_throughput".into(), + serde_json::json!(metrics.output_throughput), + ); + result.insert( + "total_token_throughput".into(), + serde_json::json!(metrics.total_token_throughput), + ); + result.insert( + "max_output_tokens_per_s".into(), + serde_json::json!(metrics.max_output_tokens_per_s), + ); + result.insert( + "max_concurrent_requests".into(), + serde_json::json!(metrics.max_concurrent_requests), + ); + // Inverse Real-Time Factor (for ASR benchmarks; 0.0 for generation) + result.insert("rtfx".into(), serde_json::json!(0.0)); + } + + // Per-failure log (always emitted when there are failures, regardless of --save-detailed). + // ttft/itl/output_tokens/start_time/latency are best-effort — zero/empty for failures + // that occurred before any tokens were received. Non-zero values indicate a partial + // failure (some tokens streamed before the error). + if metrics.failed > 0 { + result.insert( + "failed_requests".into(), + Value::Array(collect_failed_requests(outputs)), + ); + } + + // Per-request data. + // For pooling: Python always includes input_lens and errors (serve.py:1004-1013). + // For generation: all per-request arrays are gated by --save-detailed. + if is_pooling { + result.insert( + "input_lens".into(), + serde_json::json!(outputs.iter().map(|o| o.prompt_len).collect::>()), + ); + result.insert( + "errors".into(), + serde_json::json!(outputs.iter().map(|o| &o.error).collect::>()), + ); + } + if config.save_detailed { + if !is_pooling { + result.insert( + "input_lens".into(), + serde_json::json!(outputs.iter().map(|o| o.prompt_len).collect::>()), + ); + result.insert("output_lens".into(), serde_json::json!(actual_output_lens)); + result.insert( + "ttfts".into(), + serde_json::json!(outputs.iter().map(|o| o.ttft).collect::>()), + ); + result.insert( + "itls".into(), + serde_json::json!(outputs.iter().map(|o| &o.itl).collect::>()), + ); + result.insert( + "generated_texts".into(), + serde_json::json!(outputs.iter().map(|o| &o.generated_text).collect::>()), + ); + result.insert( + "errors".into(), + serde_json::json!(outputs.iter().map(|o| &o.error).collect::>()), + ); + } + result.insert( + "latencies".into(), + serde_json::json!(outputs.iter().map(|o| o.latency).collect::>()), + ); + result.insert( + "start_times".into(), + serde_json::json!(outputs.iter().map(|o| o.start_time).collect::>()), + ); + } + + // Speculative decoding stats + if let Some(stats) = spec_decode_stats { + insert_spec_decode_stats(&mut result, stats); + } + + // Per-metric stats (pooling only has e2el) + if !is_pooling { + add_metric_stats( + &mut result, + "ttft", + metrics.mean_ttft_ms, + metrics.median_ttft_ms, + metrics.std_ttft_ms, + &metrics.percentiles_ttft_ms, + &config.selected_percentile_metrics, + ); + add_metric_stats( + &mut result, + "tpot", + metrics.mean_tpot_ms, + metrics.median_tpot_ms, + metrics.std_tpot_ms, + &metrics.percentiles_tpot_ms, + &config.selected_percentile_metrics, + ); + add_metric_stats( + &mut result, + "itl", + metrics.mean_itl_ms, + metrics.median_itl_ms, + metrics.std_itl_ms, + &metrics.percentiles_itl_ms, + &config.selected_percentile_metrics, + ); + } + add_metric_stats( + &mut result, + "e2el", + metrics.mean_e2el_ms, + metrics.median_e2el_ms, + metrics.std_e2el_ms, + &metrics.percentiles_e2el_ms, + &config.selected_percentile_metrics, + ); + + // Steady-state metrics (null when not computed or scope gate failed). + result.insert( + "steady_state".into(), + serde_json::to_value(&metrics.steady_state).unwrap_or(Value::Null), + ); + + Value::Object(result) +} + +/// Build multi-turn result JSON. +pub fn build_multi_turn_result_json( + config: &BenchConfig, + mt_metrics: &MultiTurnMetrics, + conversation_outputs: &[ConversationOutput], + benchmark_duration: f64, + current_dt: &str, + spec_decode_stats: Option<&SpecDecodeStats>, +) -> serde_json::Value { + let mut result = serde_json::Map::new(); + + // Setup + result.insert("date".into(), Value::String(current_dt.to_string())); + result.insert("mode".into(), Value::String("multi_turn".to_string())); + result.insert("backend".into(), Value::String(config.backend.to_string())); + result.insert( + "model_id".into(), + Value::String(config.model.clone().unwrap_or_default()), + ); + result.insert( + "num_conversations".into(), + Value::Number(config.num_prompts.into()), + ); + result.insert( + "max_model_len".into(), + config.max_model_len.map(|v| Value::Number(v.into())).unwrap_or(Value::Null), + ); + result.insert( + "turns_per_conversation".into(), + Value::Number(config.multi_turn_num_turns.into()), + ); + result.insert( + "multi_turn_concurrency".into(), + serde_json::json!( + config + .multi_turn_concurrency + .unwrap_or(config.max_concurrency.unwrap_or(config.num_prompts)) + ), + ); + result.insert( + "inter_turn_delay_ms".into(), + Value::Number(config.multi_turn_delay_ms.into()), + ); + if config.multi_turn_prefix_global_ratio > 0.0 + || config.multi_turn_prefix_conversation_ratio > 0.0 + { + result.insert( + "prefix_global_ratio".into(), + serde_json::json!(config.multi_turn_prefix_global_ratio), + ); + result.insert( + "prefix_conversation_ratio".into(), + serde_json::json!(config.multi_turn_prefix_conversation_ratio), + ); + result.insert( + "prefix_unique_ratio".into(), + serde_json::json!( + (1.0 - (config.multi_turn_prefix_global_ratio + + config.multi_turn_prefix_conversation_ratio)) + .max(0.0) + ), + ); + result.insert("no_history_accumulation".into(), serde_json::json!(true)); + } + + // Metadata + if let Some(ref metadata) = config.metadata { + for (k, v) in metadata { + result.insert(k.clone(), Value::String(v.clone())); + } + } + + // Conversation-level stats + result.insert("duration".into(), serde_json::json!(benchmark_duration)); + result.insert( + "conversations_completed".into(), + serde_json::json!(mt_metrics.conversations_completed), + ); + result.insert( + "conversations_failed".into(), + serde_json::json!(mt_metrics.conversations_failed), + ); + result.insert( + "avg_turns_completed".into(), + serde_json::json!(mt_metrics.avg_turns_completed), + ); + result.insert( + "avg_conversation_duration_ms".into(), + serde_json::json!(mt_metrics.avg_conversation_duration_ms), + ); + + // Overall metrics + let overall = &mt_metrics.overall; + result.insert("completed".into(), serde_json::json!(overall.completed)); + result.insert("failed".into(), serde_json::json!(overall.failed)); + + // Per-failure log (always emitted when there are failures). Each entry is a single failed turn. + if overall.failed > 0 { + result.insert( + "failed_requests".into(), + Value::Array(collect_failed_turns(conversation_outputs)), + ); + } + result.insert( + "total_input_tokens".into(), + serde_json::json!(overall.total_input), + ); + result.insert( + "total_output_tokens".into(), + serde_json::json!(overall.total_output), + ); + result.insert( + "request_throughput".into(), + serde_json::json!(overall.request_throughput), + ); + result.insert( + "input_throughput".into(), + serde_json::json!(overall.input_throughput), + ); + result.insert( + "output_throughput".into(), + serde_json::json!(overall.output_throughput), + ); + result.insert( + "total_token_throughput".into(), + serde_json::json!(overall.total_token_throughput), + ); + + // Overall per-metric stats + add_metric_stats( + &mut result, + "ttft", + overall.mean_ttft_ms, + overall.median_ttft_ms, + overall.std_ttft_ms, + &overall.percentiles_ttft_ms, + &config.selected_percentile_metrics, + ); + add_metric_stats( + &mut result, + "tpot", + overall.mean_tpot_ms, + overall.median_tpot_ms, + overall.std_tpot_ms, + &overall.percentiles_tpot_ms, + &config.selected_percentile_metrics, + ); + add_metric_stats( + &mut result, + "itl", + overall.mean_itl_ms, + overall.median_itl_ms, + overall.std_itl_ms, + &overall.percentiles_itl_ms, + &config.selected_percentile_metrics, + ); + add_metric_stats( + &mut result, + "e2el", + overall.mean_e2el_ms, + overall.median_e2el_ms, + overall.std_e2el_ms, + &overall.percentiles_e2el_ms, + &config.selected_percentile_metrics, + ); + + // Speculative decoding stats + if let Some(stats) = spec_decode_stats { + insert_spec_decode_stats(&mut result, stats); + } + + // Per-turn metrics array + let per_turn_json: Vec = mt_metrics + .per_turn + .iter() + .enumerate() + .map(|(i, m)| { + let mut turn = serde_json::Map::new(); + turn.insert("turn_index".into(), serde_json::json!(i)); + turn.insert( + "num_samples".into(), + serde_json::json!(m.completed + m.failed), + ); + turn.insert("completed".into(), serde_json::json!(m.completed)); + turn.insert("failed".into(), serde_json::json!(m.failed)); + turn.insert( + "total_input_tokens".into(), + serde_json::json!(m.total_input), + ); + turn.insert( + "total_output_tokens".into(), + serde_json::json!(m.total_output), + ); + turn.insert( + "request_throughput".into(), + serde_json::json!(m.request_throughput), + ); + turn.insert( + "input_throughput".into(), + serde_json::json!(m.input_throughput), + ); + turn.insert( + "output_throughput".into(), + serde_json::json!(m.output_throughput), + ); + turn.insert( + "total_token_throughput".into(), + serde_json::json!(m.total_token_throughput), + ); + + add_metric_stats( + &mut turn, + "ttft", + m.mean_ttft_ms, + m.median_ttft_ms, + m.std_ttft_ms, + &m.percentiles_ttft_ms, + &config.selected_percentile_metrics, + ); + add_metric_stats( + &mut turn, + "tpot", + m.mean_tpot_ms, + m.median_tpot_ms, + m.std_tpot_ms, + &m.percentiles_tpot_ms, + &config.selected_percentile_metrics, + ); + add_metric_stats( + &mut turn, + "itl", + m.mean_itl_ms, + m.median_itl_ms, + m.std_itl_ms, + &m.percentiles_itl_ms, + &config.selected_percentile_metrics, + ); + add_metric_stats( + &mut turn, + "e2el", + m.mean_e2el_ms, + m.median_e2el_ms, + m.std_e2el_ms, + &m.percentiles_e2el_ms, + &config.selected_percentile_metrics, + ); + + Value::Object(turn) + }) + .collect(); + + result.insert("per_turn_metrics".into(), Value::Array(per_turn_json)); + + Value::Object(result) +} + +fn collect_failed_requests(outputs: &[RequestFuncOutput]) -> Vec { + outputs + .iter() + .enumerate() + .filter(|(_, o)| !o.success) + .map(|(i, o)| { + serde_json::json!({ + "index": i, + "error": o.error, + "prompt_len": o.prompt_len, + "start_time": o.start_time, + "latency": o.latency, + "ttft": o.ttft, + "output_tokens": o.output_tokens, + "itl": o.itl, + }) + }) + .collect() +} + +fn collect_failed_turns(conversation_outputs: &[ConversationOutput]) -> Vec { + let mut failed = Vec::new(); + for conv in conversation_outputs { + for turn in &conv.turns { + if !turn.request_output.success { + let o = &turn.request_output; + failed.push(serde_json::json!({ + "conversation_id": conv.conversation_id, + "turn_index": turn.turn_index, + "error": o.error, + "prompt_len": o.prompt_len, + "start_time": o.start_time, + "latency": o.latency, + "ttft": o.ttft, + "output_tokens": o.output_tokens, + "itl": o.itl, + })); + } + } + } + failed +} + +fn insert_spec_decode_stats(result: &mut serde_json::Map, stats: &SpecDecodeStats) { + result.insert( + "spec_decode_acceptance_rate".into(), + serde_json::json!(stats.acceptance_rate), + ); + result.insert( + "spec_decode_acceptance_length".into(), + serde_json::json!(stats.acceptance_length), + ); + result.insert( + "spec_decode_num_drafts".into(), + serde_json::json!(stats.num_drafts), + ); + result.insert( + "spec_decode_draft_tokens".into(), + serde_json::json!(stats.draft_tokens), + ); + result.insert( + "spec_decode_accepted_tokens".into(), + serde_json::json!(stats.accepted_tokens), + ); + result.insert( + "spec_decode_per_position_acceptance_rates".into(), + serde_json::json!(stats.per_position_acceptance_rates), + ); +} + +fn add_metric_stats( + result: &mut serde_json::Map, + name: &str, + mean: f64, + median: f64, + std: f64, + percentiles: &[(f64, f64)], + selected: &[String], +) { + if !selected.iter().any(|s| s == name) { + return; + } + result.insert(format!("mean_{name}_ms"), serde_json::json!(mean)); + result.insert(format!("median_{name}_ms"), serde_json::json!(median)); + result.insert(format!("std_{name}_ms"), serde_json::json!(std)); + for (p, value) in percentiles { + let p_str = if *p == p.floor() { + format!("{}", *p as i64) + } else { + format!("{p}") + }; + result.insert(format!("p{p_str}_{name}_ms"), serde_json::json!(value)); + } +} + +/// Save result JSON to file (overwrite mode). +pub fn save_result(json: &Value, file_path: &str) -> Result<()> { + let content = serde_json::to_string(json)?; + std::fs::write(file_path, content)?; + println!("Results saved to {file_path}"); + Ok(()) +} + +/// Append result JSON to file (JSONL format, matching Python's --append-result). +pub fn append_result(json: &Value, file_path: &str) -> Result<()> { + use std::io::Write; + let content = serde_json::to_string(json)?; + let mut file = std::fs::OpenOptions::new().create(true).append(true).open(file_path)?; + // If file is non-empty, prepend a newline + let meta = file.metadata()?; + if meta.len() > 0 { + file.write_all(b"\n")?; + } + file.write_all(content.as_bytes())?; + println!("Results appended to {file_path}"); + Ok(()) +} + +/// Compute the result filename matching Python's logic. +pub fn compute_result_filename(config: &BenchConfig, model_id: &str, current_dt: &str) -> String { + let base_model = model_id.split('/').next_back().unwrap_or(model_id); + let max_conc_str = + config.max_concurrency.map(|mc| format!("-concurrency{mc}")).unwrap_or_default(); + let label = config.label.as_deref().unwrap_or_else(|| config.backend.as_str()); + + let file_name = if config.request_rate.is_infinite() { + format!("{label}-infqps{max_conc_str}-{base_model}-{current_dt}.json") + } else { + format!( + "{label}-{}qps{max_conc_str}-{base_model}-{current_dt}.json", + config.request_rate + ) + }; + + if let Some(ref explicit) = config.result_filename { + if let Some(ref dir) = config.result_dir { + return format!("{dir}/{explicit}"); + } + return explicit.clone(); + } + + if let Some(ref dir) = config.result_dir { + format!("{dir}/{file_name}") + } else { + file_name + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_insert_spec_decode_stats() { + let stats = SpecDecodeStats { + num_drafts: 1000, + draft_tokens: 3000, + accepted_tokens: 2400, + acceptance_rate: 80.0, + acceptance_length: 3.4, + per_position_acceptance_rates: vec![0.9, 0.8, 0.7], + }; + + let mut result = serde_json::Map::new(); + insert_spec_decode_stats(&mut result, &stats); + + assert_eq!(result["spec_decode_acceptance_rate"], 80.0); + assert_eq!(result["spec_decode_acceptance_length"], 3.4); + assert_eq!(result["spec_decode_num_drafts"], 1000); + assert_eq!(result["spec_decode_draft_tokens"], 3000); + assert_eq!(result["spec_decode_accepted_tokens"], 2400); + assert_eq!( + result["spec_decode_per_position_acceptance_rates"], + serde_json::json!([0.9, 0.8, 0.7]) + ); + } + + #[test] + fn test_insert_spec_decode_stats_empty_positions() { + let stats = SpecDecodeStats { + num_drafts: 500, + draft_tokens: 1500, + accepted_tokens: 1200, + acceptance_rate: 80.0, + acceptance_length: 3.4, + per_position_acceptance_rates: vec![], + }; + + let mut result = serde_json::Map::new(); + insert_spec_decode_stats(&mut result, &stats); + + assert_eq!( + result["spec_decode_per_position_acceptance_rates"], + serde_json::json!([]) + ); + assert_eq!(result.len(), 6); + } + + #[test] + fn test_collect_failed_requests_empty_when_all_success() { + let outputs = vec![ + RequestFuncOutput { + success: true, + ..Default::default() + }, + RequestFuncOutput { + success: true, + ..Default::default() + }, + ]; + assert!(collect_failed_requests(&outputs).is_empty()); + } + + #[test] + fn test_collect_failed_requests_preserves_index_and_fields() { + let outputs = vec![ + RequestFuncOutput { + success: true, + ..Default::default() + }, + // Pre-stream failure: no tokens received. + RequestFuncOutput { + success: false, + error: "connection reset".into(), + prompt_len: 128, + start_time: 1.5, + latency: 0.0, + ttft: 0.0, + output_tokens: 0, + itl: vec![], + ..Default::default() + }, + RequestFuncOutput { + success: true, + ..Default::default() + }, + // Partial failure: streamed 3 tokens before timing out. + RequestFuncOutput { + success: false, + error: "timeout".into(), + prompt_len: 64, + start_time: 2.25, + latency: 30.0, + ttft: 0.5, + output_tokens: 3, + itl: vec![0.1, 0.12], + ..Default::default() + }, + ]; + + let failed = collect_failed_requests(&outputs); + assert_eq!(failed.len(), 2); + + assert_eq!(failed[0]["index"], 1); + assert_eq!(failed[0]["error"], "connection reset"); + assert_eq!(failed[0]["prompt_len"], 128); + assert_eq!(failed[0]["start_time"], 1.5); + assert_eq!(failed[0]["latency"], 0.0); + assert_eq!(failed[0]["ttft"], 0.0); + assert_eq!(failed[0]["output_tokens"], 0); + assert_eq!(failed[0]["itl"], serde_json::json!([])); + + assert_eq!(failed[1]["index"], 3); + assert_eq!(failed[1]["error"], "timeout"); + assert_eq!(failed[1]["latency"], 30.0); + assert_eq!(failed[1]["ttft"], 0.5); + assert_eq!(failed[1]["output_tokens"], 3); + assert_eq!(failed[1]["itl"], serde_json::json!([0.1, 0.12])); + } +} diff --git a/rust/src/bench/src/output/mod.rs b/rust/src/bench/src/output/mod.rs new file mode 100644 index 000000000000..54c88f0d6452 --- /dev/null +++ b/rust/src/bench/src/output/mod.rs @@ -0,0 +1,5 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +pub mod console; +pub mod json; diff --git a/rust/src/bench/src/rate_control.rs b/rust/src/bench/src/rate_control.rs new file mode 100644 index 000000000000..aa05f8f8a8df --- /dev/null +++ b/rust/src/bench/src/rate_control.rs @@ -0,0 +1,171 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use rand::SeedableRng; +use rand::rngs::StdRng; +use rand_distr::{Distribution, Gamma}; + +use crate::cli::RampUpStrategy; +use crate::config::RampUpConfig; + +/// Pre-computed request schedule: cumulative delays from start and per-request rates. +pub struct RequestSchedule { + /// Cumulative absolute delay from start (seconds) for each request. + pub delays: Vec, + /// Instantaneous rate for each request (for ramp-up logging). + #[allow(dead_code)] + pub rates: Vec, +} + +/// Compute the per-request rate, accounting for optional ramp-up. +/// +/// Mirrors Python's `_get_current_request_rate()` from serve.py:218-240. +fn get_current_request_rate( + ramp_up: Option<&RampUpConfig>, + request_index: usize, + total_requests: usize, + base_rate: f64, +) -> f64 { + let config = match ramp_up { + Some(c) => c, + None => return base_rate, + }; + + let progress = request_index as f64 / (total_requests - 1).max(1) as f64; + + match config.strategy { + RampUpStrategy::Linear => { + let increase = (config.end_rps - config.start_rps) * progress; + config.start_rps + increase + } + RampUpStrategy::Exponential => { + let ratio = config.end_rps / config.start_rps; + config.start_rps * ratio.powf(progress) + } + } +} + +/// Compute the request schedule for all requests. +/// +/// Ports the Python `get_request()` / `_generate_request_timestamps()` logic +/// from serve.py:243-340. +pub fn compute_schedule( + num_requests: usize, + request_rate: f64, + burstiness: f64, + seed: u64, + ramp_up: Option<&RampUpConfig>, +) -> RequestSchedule { + assert!(burstiness > 0.0, "burstiness must be positive"); + assert!(num_requests > 0, "must have at least one request"); + + let mut rng = StdRng::seed_from_u64(seed); + let mut delay_ts = Vec::with_capacity(num_requests); + let mut rates = Vec::with_capacity(num_requests); + + for i in 0..num_requests { + let current_rate = get_current_request_rate(ramp_up, i, num_requests, request_rate); + rates.push(current_rate); + + if current_rate.is_infinite() { + delay_ts.push(0.0); + } else if burstiness.is_infinite() { + // When burstiness → ∞, delay becomes constant = 1/rate + delay_ts.push(1.0 / current_rate); + } else { + let theta = 1.0 / (current_rate * burstiness); + let gamma = Gamma::new(burstiness, theta).unwrap(); + delay_ts.push(gamma.sample(&mut rng)); + } + } + + // Compute cumulative delays + for i in 1..delay_ts.len() { + delay_ts[i] += delay_ts[i - 1]; + } + + // Normalize: scale cumulative delays so total matches target. + // Only for fixed-rate (no ramp-up) mode, matching Python behavior. + if ramp_up.is_none() + && let Some(&last) = delay_ts.last() + && last > 0.0 + && !request_rate.is_infinite() + { + let target_total = num_requests as f64 / request_rate; + let factor = target_total / last; + for d in &mut delay_ts { + *d *= factor; + } + } + + RequestSchedule { + delays: delay_ts, + rates, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_infinite_rate_all_zero_delays() { + let sched = compute_schedule(100, f64::INFINITY, 1.0, 42, None); + assert_eq!(sched.delays.len(), 100); + for d in &sched.delays { + assert_eq!(*d, 0.0); + } + } + + #[test] + fn test_fixed_rate_monotonic() { + let sched = compute_schedule(50, 10.0, 1.0, 42, None); + for i in 1..sched.delays.len() { + assert!(sched.delays[i] >= sched.delays[i - 1]); + } + // Total should be close to num_requests / rate = 5.0 + let last = *sched.delays.last().unwrap(); + assert!((last - 5.0).abs() < 0.01, "last delay = {last}"); + } + + #[test] + fn test_infinite_burstiness_constant_delay() { + let sched = compute_schedule(10, 5.0, f64::INFINITY, 42, None); + // Each delay should be 0.2s apart (1/5) + for i in 1..sched.delays.len() { + let delta = sched.delays[i] - sched.delays[i - 1]; + assert!((delta - 0.2).abs() < 1e-10); + } + } + + #[test] + fn test_linear_ramp_up() { + let ramp = RampUpConfig { + strategy: RampUpStrategy::Linear, + start_rps: 1.0, + end_rps: 10.0, + }; + let sched = compute_schedule(10, 5.0, 1.0, 42, Some(&ramp)); + // Rates should increase linearly from 1.0 to 10.0 + assert!((sched.rates[0] - 1.0).abs() < 1e-10); + assert!((sched.rates[9] - 10.0).abs() < 1e-10); + // Middle should be ~5.5 + assert!((sched.rates[4] - 5.0).abs() < 0.5); + } + + #[test] + fn test_exponential_ramp_up() { + let ramp = RampUpConfig { + strategy: RampUpStrategy::Exponential, + start_rps: 1.0, + end_rps: 100.0, + }; + let sched = compute_schedule(10, 5.0, 1.0, 42, Some(&ramp)); + assert!((sched.rates[0] - 1.0).abs() < 1e-10); + assert!((sched.rates[9] - 100.0).abs() < 1e-10); + // Exponential: rates should be monotonically increasing + for i in 1..sched.rates.len() { + assert!(sched.rates[i] >= sched.rates[i - 1]); + } + } +} diff --git a/rust/src/bench/src/ready_checker.rs b/rust/src/bench/src/ready_checker.rs new file mode 100644 index 000000000000..86af8663d115 --- /dev/null +++ b/rust/src/bench/src/ready_checker.rs @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::time::Instant; + +use indicatif::{ProgressBar, ProgressStyle}; + +use crate::backends::{RequestFuncInput, RequestFuncOutput, get_backend}; +use crate::cli::BackendKind; +use crate::error::{BenchError, Result}; + +/// Wait for the serving endpoint to become available. +/// +/// Sends test requests with retry until success or timeout. +/// Mirrors Python's `wait_for_endpoint` in ready_checker.py. +pub async fn wait_for_endpoint( + backend: BackendKind, + client: &reqwest::Client, + test_input: &RequestFuncInput, + timeout_seconds: u64, + retry_interval: u64, +) -> Result { + let backend = get_backend(backend)?; + let deadline = Instant::now() + std::time::Duration::from_secs(timeout_seconds); + + println!("Waiting for endpoint to become up in {timeout_seconds}s"); + + let pb = ProgressBar::new(timeout_seconds); + pb.set_style( + ProgressStyle::with_template("{msg} |{bar:40}| {elapsed} elapsed, {eta} remaining") + .unwrap() + .progress_chars("##-"), + ); + + let mut last_error = String::new(); + + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + let elapsed = timeout_seconds.saturating_sub(remaining.as_secs()); + pb.set_position(elapsed); + + if remaining.is_zero() { + pb.finish_and_clear(); + break; + } + + // Ping the endpoint + match backend.send_request(test_input, client).await { + Ok(output) if output.success => { + pb.finish_and_clear(); + return Ok(output); + } + Ok(output) => { + let err = output.error.clone(); + let err_last_line = err.lines().last().unwrap_or(&err); + eprintln!("Endpoint is not ready. Error='{err_last_line}'"); + last_error = err; + } + Err(e) => { + last_error = e.to_string(); + } + } + + // Retry after delay + let sleep_dur = std::cmp::min(std::time::Duration::from_secs(retry_interval), remaining); + if !sleep_dur.is_zero() { + tokio::time::sleep(sleep_dur).await; + } + } + + Err(BenchError::EndpointTimeout(timeout_seconds, last_error)) +} diff --git a/rust/src/bench/src/sweep.rs b/rust/src/bench/src/sweep.rs new file mode 100644 index 000000000000..68cee1582530 --- /dev/null +++ b/rust/src/bench/src/sweep.rs @@ -0,0 +1,561 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use crate::config::BenchConfig; +use crate::error::{BenchError, Result}; +use crate::output::json::{compute_result_filename, save_result}; + +/// Reset the server's prefix cache by calling POST /reset_prefix_cache. +/// Requires VLLM_SERVER_DEV_MODE=1 on the vLLM server. +async fn reset_prefix_cache(base_url: &str) -> Result<()> { + let url = format!("{}/reset_prefix_cache", base_url.trim_end_matches('/')); + let client = reqwest::Client::new(); + let resp = client + .post(&url) + .send() + .await + .map_err(|e| BenchError::Backend(format!("Failed to reset prefix cache: {e}")))?; + if resp.status().is_success() { + println!("Prefix cache reset successfully."); + } else { + let status = resp.status(); + let body = resp.text().await.unwrap_or_default(); + return Err(BenchError::Backend(format!( + "Failed to reset prefix cache: HTTP {status} — {body}. \ + Is VLLM_SERVER_DEV_MODE=1 set on the server?" + ))); + } + Ok(()) +} + +/// Result of a single sweep point. +struct SweepPoint { + label: String, + #[allow(dead_code)] + value: f64, + result_json: serde_json::Value, + ss_request_throughput: Option, + ss_output_throughput: Option, +} + +/// Run a sweep over max-concurrency values. +pub async fn run_concurrency_sweep( + base_config: &BenchConfig, + values: &[usize], + num_prompts_factor: Option, +) -> Result<()> { + println!("{:=^70}", " Concurrency Sweep "); + println!( + "Sweeping --max-concurrency over {} values: {:?}", + values.len(), + values + ); + println!(); + + let mut points = Vec::with_capacity(values.len()); + + let current_dt = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string(); + + for (i, &mc) in values.iter().enumerate() { + println!( + "{:-^70}", + format!(" Run {}/{}: max_concurrency={} ", i + 1, values.len(), mc) + ); + + if base_config.reset_prefix_cache { + reset_prefix_cache(&base_config.base_url).await?; + } + + let mut config = base_config.clone(); + config.max_concurrency = Some(mc); + if let Some(factor) = num_prompts_factor { + config.num_prompts = mc * factor; + } + // Suppress per-run save (we save with concurrency suffix below) + config.save_result = false; + config.append_result = false; + + let result = crate::benchmark::run_benchmark(&config).await?; + + if base_config.save_result { + let model_id = config.model.as_deref().unwrap_or("unknown"); + let file_name = compute_result_filename(&config, model_id, ¤t_dt); + if let Some(ref dir) = config.result_dir { + std::fs::create_dir_all(dir)?; + } + save_result(&result, &file_name)?; + } + + points.push(SweepPoint { + label: format!("concurrency={mc}"), + value: mc as f64, + ss_request_throughput: ss_f64(&result, "request_throughput"), + ss_output_throughput: ss_f64(&result, "output_throughput"), + result_json: result, + }); + + println!(); + } + + print_sweep_summary( + "Max Concurrency", + &points, + &base_config.sweep_summary_percentiles, + ); + Ok(()) +} + +/// Run a sweep over multi-turn concurrency values. +pub async fn run_multi_turn_concurrency_sweep( + base_config: &BenchConfig, + values: &[usize], + num_prompts_factor: Option, +) -> Result<()> { + println!("{:=^70}", " Multi-Turn Concurrency Sweep "); + println!( + "Sweeping --multi-turn-concurrency over {} values: {:?}", + values.len(), + values + ); + println!(); + + let mut points = Vec::with_capacity(values.len()); + let current_dt = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string(); + + for (i, &mc) in values.iter().enumerate() { + println!( + "{:-^70}", + format!( + " Run {}/{}: multi_turn_concurrency={} ", + i + 1, + values.len(), + mc + ) + ); + + if base_config.reset_prefix_cache { + reset_prefix_cache(&base_config.base_url).await?; + } + + let mut config = base_config.clone(); + config.multi_turn_concurrency = Some(mc); + // Set max_concurrency so compute_result_filename includes the + // concurrency suffix — without this every sweep point overwrites + // the same file. + config.max_concurrency = Some(mc); + if let Some(factor) = num_prompts_factor { + config.num_prompts = mc * factor; + } + // Suppress per-run save (we save with concurrency suffix below) + config.save_result = false; + config.append_result = false; + + let result = crate::multi_turn::run_multi_turn_benchmark(&config).await?; + + if base_config.save_result { + let model_id = config.model.as_deref().unwrap_or("unknown"); + let file_name = compute_result_filename(&config, model_id, ¤t_dt); + if let Some(ref dir) = config.result_dir { + std::fs::create_dir_all(dir)?; + } + save_result(&result, &file_name)?; + } + + points.push(SweepPoint { + label: format!("concurrency={mc}"), + value: mc as f64, + ss_request_throughput: ss_f64(&result, "request_throughput"), + ss_output_throughput: ss_f64(&result, "output_throughput"), + result_json: result, + }); + + println!(); + } + + print_sweep_summary( + "MT Concurrency", + &points, + &base_config.sweep_summary_percentiles, + ); + Ok(()) +} + +/// Run a sweep over request-rate values. +pub async fn run_rate_sweep(base_config: &BenchConfig, values: &[f64]) -> Result<()> { + println!("{:=^70}", " Request Rate Sweep "); + println!( + "Sweeping --request-rate over {} values: {:?}", + values.len(), + values + ); + println!(); + + let mut points = Vec::with_capacity(values.len()); + let current_dt = chrono::Local::now().format("%Y%m%d-%H%M%S").to_string(); + + for (i, &rate) in values.iter().enumerate() { + let rate_str = if rate.is_infinite() { + "inf".to_string() + } else { + format!("{rate}") + }; + println!( + "{:-^70}", + format!( + " Run {}/{}: request_rate={} ", + i + 1, + values.len(), + rate_str + ) + ); + + if base_config.reset_prefix_cache { + reset_prefix_cache(&base_config.base_url).await?; + } + + let mut config = base_config.clone(); + config.request_rate = rate; + config.save_result = false; + config.append_result = false; + + let result = crate::benchmark::run_benchmark(&config).await?; + + if base_config.save_result { + let model_id = config.model.as_deref().unwrap_or("unknown"); + let file_name = compute_result_filename(&config, model_id, ¤t_dt); + if let Some(ref dir) = config.result_dir { + std::fs::create_dir_all(dir)?; + } + save_result(&result, &file_name)?; + } + + points.push(SweepPoint { + label: format!("rate={rate_str}"), + value: rate, + ss_request_throughput: ss_f64(&result, "request_throughput"), + ss_output_throughput: ss_f64(&result, "output_throughput"), + result_json: result, + }); + + println!(); + } + + print_sweep_summary( + "Request Rate", + &points, + &base_config.sweep_summary_percentiles, + ); + Ok(()) +} + +/// Print a summary table after all sweep points complete. +fn print_sweep_summary(param_name: &str, points: &[SweepPoint], summary_percentiles: &[f64]) { + let param_width = param_name.len().max(20); + let columns = build_summary_columns(summary_percentiles); + let total_width = param_width + + columns.iter().map(SummaryColumn::render_width).sum::() + + columns.len(); + + println!("{:=^width$}", " Sweep Summary ", width = total_width); + + print!("{param_name:width$}", column.header, width = column.render_width()); + } + println!(); + + print!("{:- Vec { + let mut columns = vec![ + SummaryColumn::new("Req/s", "request_throughput", 10), + SummaryColumn::new("Tok/s", "output_throughput", 10), + SummaryColumn::new("Total tok/s", "total_token_throughput", 12), + SummaryColumn::new("SS req/s", SS_REQUEST_THROUGHPUT_KEY, 10), + SummaryColumn::new("SS out tok/s", SS_OUTPUT_THROUGHPUT_KEY, 12), + SummaryColumn::new("P50 TTFT(ms)", "median_ttft_ms", 12), + SummaryColumn::new("P50 TPOT(ms)", "median_tpot_ms", 12), + SummaryColumn::new("P90 TTFT(ms)", "p90_ttft_ms", 12), + SummaryColumn::new("P90 TPOT(ms)", "p90_tpot_ms", 12), + ]; + + for &percentile in summary_percentiles { + if percentile == 50.0 || percentile == 90.0 { + continue; + } + let p_str = format_percentile(percentile); + columns.push(SummaryColumn::new( + &format!("P{p_str} TTFT(ms)"), + &format!("p{p_str}_ttft_ms"), + 12, + )); + columns.push(SummaryColumn::new( + &format!("P{p_str} TPOT(ms)"), + &format!("p{p_str}_tpot_ms"), + 12, + )); + } + + columns +} + +fn render_summary_row(point: &SweepPoint, param_width: usize, columns: &[SummaryColumn]) -> String { + let mut row = format!("{: point.ss_request_throughput, + SS_OUTPUT_THROUGHPUT_KEY => point.ss_output_throughput, + key => get_f64(&point.result_json, key), + }; + row.push(' '); + row.push_str(&format!( + "{:>width$}", + fmt_f64(value), + width = column.render_width() + )); + } + row +} + +fn format_percentile(percentile: f64) -> String { + if percentile == percentile.floor() { + format!("{}", percentile as i64) + } else { + format!("{percentile}") + } +} + +struct SummaryColumn { + header: String, + key: String, + width: usize, +} + +impl SummaryColumn { + fn new(header: &str, key: &str, width: usize) -> Self { + Self { + header: header.to_string(), + key: key.to_string(), + width, + } + } + + fn render_width(&self) -> usize { + self.width.max(self.header.len()) + } +} + +/// Parse a comma-separated list of concurrency values. +pub fn parse_concurrency_values(s: &str) -> Result> { + s.split(',') + .map(|v| { + v.trim().parse::().map_err(|_| { + BenchError::Config(format!("Invalid concurrency value: '{}'", v.trim())) + }) + }) + .collect() +} + +/// Parse a comma-separated list of request rate values (supports "inf"). +pub fn parse_rate_values(s: &str) -> Result> { + s.split(',') + .map(|v| { + let v = v.trim(); + if v == "inf" { + Ok(f64::INFINITY) + } else { + v.parse::() + .map_err(|_| BenchError::Config(format!("Invalid request rate value: '{v}'"))) + } + }) + .collect() +} + +fn get_f64(json: &serde_json::Value, key: &str) -> Option { + json.get(key).and_then(|v| v.as_f64()) +} + +/// Sentinel keys used on `SummaryColumn` to mark steady-state columns whose +/// values come from `SweepPoint::ss_request_throughput` / `ss_output_throughput` +/// rather than from the result JSON top level. +const SS_REQUEST_THROUGHPUT_KEY: &str = "__ss_request_throughput"; +const SS_OUTPUT_THROUGHPUT_KEY: &str = "__ss_output_throughput"; + +/// Extract a numeric field from the `steady_state` sub-object of a result JSON. +/// Returns `None` when `steady_state` is missing or null, or when the field is +/// absent / non-numeric. +fn ss_f64(result: &serde_json::Value, key: &str) -> Option { + result + .get("steady_state") + .and_then(|ss| if ss.is_null() { None } else { Some(ss) }) + .and_then(|ss| ss.get(key)) + .and_then(|v| v.as_f64()) +} + +fn fmt_f64(v: Option) -> String { + match v { + Some(f) => format!("{:.2}", f), + None => "-".to_string(), + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn point(result_json: serde_json::Value) -> SweepPoint { + let ss_request_throughput = ss_f64(&result_json, "request_throughput"); + let ss_output_throughput = ss_f64(&result_json, "output_throughput"); + SweepPoint { + label: "point".to_string(), + value: 1.0, + result_json, + ss_request_throughput, + ss_output_throughput, + } + } + + #[test] + fn summary_columns_follow_requested_percentile_order() { + let columns = build_summary_columns(&[90.0, 95.0]); + let headers: Vec<&str> = columns.iter().map(|column| column.header.as_str()).collect(); + + assert_eq!( + headers, + vec![ + "Req/s", + "Tok/s", + "Total tok/s", + "SS req/s", + "SS out tok/s", + "P50 TTFT(ms)", + "P50 TPOT(ms)", + "P90 TTFT(ms)", + "P90 TPOT(ms)", + "P95 TTFT(ms)", + "P95 TPOT(ms)", + ] + ); + } + + #[test] + fn summary_row_uses_dash_for_missing_requested_percentile_values() { + let columns = build_summary_columns(&[90.0]); + let row = render_summary_row( + &point(json!({ + "request_throughput": 1.0, + "output_throughput": 2.0, + "total_token_throughput": 3.0, + "median_ttft_ms": 10.0, + "median_tpot_ms": 20.0, + "p90_ttft_ms": 40.0, + })), + 20, + &columns, + ); + + assert!(row.contains("40.00")); + assert!(row.contains(" -")); + } + + #[test] + fn summary_columns_skip_duplicate_p50_and_p90_percentiles() { + let columns = build_summary_columns(&[50.0, 90.0]); + let headers: Vec<&str> = columns.iter().map(|column| column.header.as_str()).collect(); + + assert_eq!( + headers, + vec![ + "Req/s", + "Tok/s", + "Total tok/s", + "SS req/s", + "SS out tok/s", + "P50 TTFT(ms)", + "P50 TPOT(ms)", + "P90 TTFT(ms)", + "P90 TPOT(ms)", + ] + ); + } + + #[test] + fn summary_row_renders_steady_state_columns() { + let columns = build_summary_columns(&[]); + let row = render_summary_row( + &point(json!({ + "request_throughput": 1.0, + "output_throughput": 2.0, + "total_token_throughput": 3.0, + "median_ttft_ms": 10.0, + "median_tpot_ms": 20.0, + "p90_ttft_ms": 40.0, + "p90_tpot_ms": 50.0, + "steady_state": { + "request_throughput": 0.77, + "output_throughput": 88.88, + }, + })), + 20, + &columns, + ); + + assert!(row.contains("0.77")); + assert!(row.contains("88.88")); + } + + #[test] + fn summary_row_renders_dash_when_steady_state_missing() { + let columns = build_summary_columns(&[]); + let row = render_summary_row( + &point(json!({ + "request_throughput": 1.0, + "output_throughput": 2.0, + "total_token_throughput": 3.0, + "median_ttft_ms": 10.0, + "median_tpot_ms": 20.0, + "p90_ttft_ms": 40.0, + "p90_tpot_ms": 50.0, + "steady_state": null, + })), + 20, + &columns, + ); + + // Two dashes for the two SS columns (SS req/s, SS out tok/s). Other + // columns have values so they won't emit dashes. + let dash_count = row.matches(" -").count(); + assert!( + dash_count >= 2, + "expected at least 2 dashes, got {dash_count}: {row}" + ); + } +} diff --git a/rust/src/bench/src/tiktoken.rs b/rust/src/bench/src/tiktoken.rs new file mode 100644 index 000000000000..b0cfbfeae499 --- /dev/null +++ b/rust/src/bench/src/tiktoken.rs @@ -0,0 +1,542 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::HashMap; +use std::path::Path; + +use rustc_hash::FxHashMap; + +use crate::error::{BenchError, Result}; + +/// Default regex pattern matching cl100k_base (GPT-4, Qwen, etc.) +const DEFAULT_TIKTOKEN_PATTERN: &str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"; + +/// Number of reserved special token slots (matches Python's +/// TikTokenTokenizer.num_reserved_special_tokens) +const NUM_RESERVED_SPECIAL_TOKENS: u32 = 256; + +/// Tiktoken-based tokenizer for models that use tiktoken format (Kimi, Qwen, etc.) +pub struct TiktokenTokenizer { + bpe: tiktoken_rs::CoreBPE, + vocab_size: u32, + #[allow(dead_code)] + num_base_tokens: u32, + /// IDs of all special tokens (for filtering from allowed_tokens) + special_token_ids: Vec, + /// Reverse mapping: token_id -> byte sequence, for lossy UTF-8 decoding. + /// Empty for built-in encodings (use bpe.decode instead). + decoder: Vec>, + /// True for built-in encodings loaded via tiktoken_rs (o200k_base, cl100k_base, etc.) + is_builtin: bool, +} + +impl TiktokenTokenizer { + /// Load from a tiktoken .model file. + /// + /// `special_tokens_from_config`: tokens from tokenizer_config.json's added_tokens_decoder + /// `all_special_tokens`: ALL special tokens including the full 256 reserved slots + pub fn from_file( + model_path: &Path, + all_special_tokens: HashMap, + special_token_ids: Vec, + pattern: Option<&str>, + ) -> Result { + let content = std::fs::read_to_string(model_path) + .map_err(|e| BenchError::Tokenizer(format!("Failed to read tiktoken model: {e}")))?; + + let mut encoder: FxHashMap, u32> = FxHashMap::default(); + + for line in content.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + let mut parts = line.split_whitespace(); + let b64 = match parts.next() { + Some(s) => s, + None => continue, + }; + let rank_str = match parts.next() { + Some(s) => s, + None => continue, + }; + + use base64::Engine; + let token_bytes = base64::engine::general_purpose::STANDARD + .decode(b64) + .map_err(|e| BenchError::Tokenizer(format!("Invalid base64 in model file: {e}")))?; + let rank: u32 = rank_str + .parse() + .map_err(|e| BenchError::Tokenizer(format!("Invalid rank: {e}")))?; + encoder.insert(token_bytes, rank); + } + + if encoder.is_empty() { + return Err(BenchError::Tokenizer("Empty tiktoken model file".into())); + } + + let num_base_tokens = encoder.len() as u32; + + // Convert special tokens to FxHashMap for tiktoken-rs + let special_fx: FxHashMap = + all_special_tokens.iter().map(|(k, &v)| (k.clone(), v)).collect(); + + // vocab_size = base tokens + all reserved special token slots + let vocab_size = num_base_tokens + NUM_RESERVED_SPECIAL_TOKENS; + + // Build reverse mapping for lossy decode + let mut decoder = vec![Vec::new(); vocab_size as usize]; + for (bytes, &rank) in &encoder { + if (rank as usize) < decoder.len() { + decoder[rank as usize] = bytes.clone(); + } + } + for (text, &rank) in &all_special_tokens { + if (rank as usize) < decoder.len() { + decoder[rank as usize] = text.as_bytes().to_vec(); + } + } + + let pat = pattern.unwrap_or(DEFAULT_TIKTOKEN_PATTERN); + + let bpe = tiktoken_rs::CoreBPE::new(encoder, special_fx, pat) + .map_err(|e| BenchError::Tokenizer(format!("Failed to build tiktoken BPE: {e}")))?; + + Ok(Self { + bpe, + vocab_size, + num_base_tokens, + special_token_ids, + decoder, + is_builtin: false, + }) + } + + /// Create from a built-in tiktoken encoding (o200k_base, cl100k_base, etc.) + pub fn from_builtin_bpe(bpe: tiktoken_rs::CoreBPE, vocab_size: u32) -> Self { + Self { + bpe, + vocab_size, + num_base_tokens: vocab_size, + special_token_ids: Vec::new(), + decoder: Vec::new(), + is_builtin: true, + } + } + + pub fn encode(&self, text: &str) -> Vec { + // Use encode_with_special_tokens to match Python's encode(allowed_special="all") + self.bpe.encode_with_special_tokens(text) + } + + /// Decode token IDs to text with lossy UTF-8 handling. + pub fn decode(&self, ids: &[u32]) -> Result { + if self.is_builtin { + // Byte-level decode + lossy UTF-8, matching the file-based path below. + // (CoreBPE::decode errors on invalid UTF-8, which random token + // sequences routinely produce with byte-level BPE vocabularies.) + let bytes: Vec = + self.bpe._decode_native_and_split(ids.to_vec()).flatten().collect(); + return Ok(String::from_utf8_lossy(&bytes).into_owned()); + } + let mut bytes = Vec::new(); + for &id in ids { + if let Some(token_bytes) = self.decoder.get(id as usize) { + bytes.extend_from_slice(token_bytes); + } + } + Ok(String::from_utf8_lossy(&bytes).into_owned()) + } + + pub fn vocab_size(&self) -> u32 { + self.vocab_size + } + + /// Get non-special token IDs whose byte representation is valid UTF-8. + /// Excludes ALL special tokens (like Python's `set(all_tokens) - set(prohibited_tokens)`). + pub fn get_allowed_tokens(&self) -> Vec { + if self.is_builtin { + // For built-in encodings, return full token range. + // Note: when used with random dataset, these IDs are sent to vLLM as-is; + // ensure the model's tokenizer is compatible (e.g. GPT-4o for o200k_base). + return (0..self.vocab_size).collect(); + } + let special_set: std::collections::HashSet = + self.special_token_ids.iter().copied().collect(); + + self.decoder + .iter() + .enumerate() + .filter(|(id, bytes)| { + !bytes.is_empty() + && std::str::from_utf8(bytes).is_ok() + && !special_set.contains(&(*id as u32)) + }) + .map(|(id, _)| id as u32) + .collect() + } +} + +/// Load a built-in tiktoken encoding by name. +/// +/// Supported names: `o200k_base` (GPT-4o), `cl100k_base` (GPT-4/3.5-turbo), +/// `p50k_base`, `r50k_base`, `gpt2`. +/// +/// These encodings are bundled with tiktoken-rs — no network download required. +/// Useful for consistent cross-model token counting (e.g. Artificial Analysis methodology). +pub fn load_builtin_tiktoken(encoding: &str) -> Result { + let (bpe, vocab_size) = match encoding { + "o200k_base" => (tiktoken_rs::o200k_base(), 200_275u32), + "cl100k_base" => (tiktoken_rs::cl100k_base(), 100_277u32), + "p50k_base" => (tiktoken_rs::p50k_base(), 50_281u32), + "p50k_edit" => (tiktoken_rs::p50k_edit(), 50_281u32), + "r50k_base" | "gpt2" => (tiktoken_rs::r50k_base(), 50_257u32), + _ => { + return Err(BenchError::Tokenizer(format!( + "Unknown built-in tiktoken encoding: '{encoding}'. \ + Supported: o200k_base, cl100k_base, p50k_base, r50k_base, gpt2" + ))); + } + }; + let bpe = bpe.map_err(|e| BenchError::Tokenizer(format!("Failed to load {encoding}: {e}")))?; + println!("Tokenizer: Built-in tiktoken {encoding} (vocab_size={vocab_size})"); + Ok(TiktokenTokenizer::from_builtin_bpe(bpe, vocab_size)) +} + +/// Try to load a tiktoken tokenizer from a local directory or HuggingFace model repo. +pub fn try_load_tiktoken(model_id: &str) -> Result { + // Phase 1: If model_id is a local directory, look for tiktoken files there + let local_dir = Path::new(model_id); + if local_dir.is_dir() { + return try_load_tiktoken_from_dir(local_dir, model_id); + } + + // Phase 2: Fall back to HuggingFace Hub download + try_load_tiktoken_from_hf(model_id) +} + +/// Common tiktoken model filenames to search for. +const TIKTOKEN_MODEL_FILENAMES: &[&str] = &["tiktoken.model", "qwen.tiktoken", "vocab.tiktoken"]; + +/// Load a tiktoken tokenizer from a local directory. +fn try_load_tiktoken_from_dir(dir: &Path, model_id: &str) -> Result { + let model_path = TIKTOKEN_MODEL_FILENAMES + .iter() + .map(|f| dir.join(f)) + .find(|p| p.exists()) + .ok_or_else(|| { + BenchError::Tokenizer(format!( + "No tiktoken model file found in local directory '{model_id}' \ + (looked for: {})", + TIKTOKEN_MODEL_FILENAMES.join(", ") + )) + })?; + + let num_base_tokens = count_base_tokens(&model_path)?; + + let config_path = dir.join("tokenizer_config.json"); + let config = if config_path.exists() { + read_tokenizer_config(&config_path) + } else { + None + }; + + let pattern = extract_pat_str_from_local_dir(dir); + + build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens) +} + +/// Load a tiktoken tokenizer from a HuggingFace model repo. +fn try_load_tiktoken_from_hf(model_id: &str) -> Result { + let repo = crate::hub::HubRepo::model(model_id.to_string()); + + let model_path = repo + .get("tiktoken.model") + .or_else(|_| repo.get("qwen.tiktoken")) + .or_else(|_| repo.get("vocab.tiktoken")) + .map_err(|_| { + BenchError::Tokenizer(format!("No tiktoken model file found for '{model_id}'")) + })?; + + let num_base_tokens = count_base_tokens(&model_path)?; + + let config = match repo.get("tokenizer_config.json") { + Ok(config_path) => read_tokenizer_config(&config_path), + Err(_) => None, + }; + + let pattern = extract_pat_str_from_repo(&repo); + + build_tiktoken(model_id, &model_path, config, pattern, num_base_tokens) +} + +/// Build a TiktokenTokenizer from discovered model file, config, and pattern. +fn build_tiktoken( + model_id: &str, + model_path: &Path, + config: Option, + pattern: Option, + num_base_tokens: u32, +) -> Result { + // Build the full 256 reserved special tokens map. + // Python: {special_tokens_mapping.get(i, f"<|reserved_token_{i}|>"): i + // for i in range(num_base_tokens, num_base_tokens + 256)} + let config_special = config.as_ref().map(|c| &c.added_tokens).cloned().unwrap_or_default(); + + // Invert config_special: id -> content + let id_to_content: HashMap = + config_special.iter().map(|(content, &id)| (id, content.clone())).collect(); + + let mut all_special_tokens: HashMap = HashMap::new(); + let mut special_token_ids: Vec = Vec::new(); + + for i in 0..NUM_RESERVED_SPECIAL_TOKENS { + let token_id = num_base_tokens + i; + let content = id_to_content + .get(&token_id) + .cloned() + .unwrap_or_else(|| format!("<|reserved_token_{i}|>")); + all_special_tokens.insert(content, token_id); + special_token_ids.push(token_id); + } + + // Also collect special IDs that Python marks as prohibited + // (those with "special": true in added_tokens_decoder) + let prohibited_ids = config.as_ref().map(|c| c.special_ids.clone()).unwrap_or_default(); + for id in &prohibited_ids { + if !special_token_ids.contains(id) { + special_token_ids.push(*id); + } + } + + println!( + "Loading tiktoken model for '{model_id}' (base={}, special={}, pat={})...", + num_base_tokens, + all_special_tokens.len(), + if pattern.is_some() { + "custom" + } else { + "default" + }, + ); + + TiktokenTokenizer::from_file( + model_path, + all_special_tokens, + special_token_ids, + pattern.as_deref(), + ) +} + +/// Count base tokens in a tiktoken .model file without fully parsing it. +fn count_base_tokens(model_path: &Path) -> Result { + let content = std::fs::read_to_string(model_path) + .map_err(|e| BenchError::Tokenizer(format!("Failed to read tiktoken model: {e}")))?; + let count = content + .lines() + .filter(|l| { + let l = l.trim(); + !l.is_empty() && l.split_whitespace().count() >= 2 + }) + .count(); + Ok(count as u32) +} + +/// Parsed tokenizer config +struct TokenizerConfig { + /// All tokens from added_tokens_decoder (content -> id) + added_tokens: HashMap, + /// IDs of tokens marked "special": true + special_ids: Vec, +} + +/// Read and parse tokenizer_config.json +fn read_tokenizer_config(config_path: &Path) -> Option { + let content = std::fs::read_to_string(config_path).ok()?; + let config: serde_json::Value = serde_json::from_str(&content).ok()?; + + let mut added_tokens = HashMap::new(); + let mut special_ids = Vec::new(); + + if let Some(added) = config.get("added_tokens_decoder").and_then(|v| v.as_object()) { + for (id_str, token_info) in added { + if let (Ok(id), Some(content)) = ( + id_str.parse::(), + token_info.get("content").and_then(|c| c.as_str()), + ) { + added_tokens.insert(content.to_string(), id); + + let is_special = + token_info.get("special").and_then(|s| s.as_bool()).unwrap_or(false); + if is_special { + special_ids.push(id); + } + } + } + } + + Some(TokenizerConfig { + added_tokens, + special_ids, + }) +} + +/// Try to extract pat_str from Python tokenizer source files in a local directory. +/// Returns None if unavailable or unparsable. +fn extract_pat_str_from_local_dir(dir: &Path) -> Option { + ["tokenization_kimi.py", "tokenizer.py"] + .iter() + .map(|f| dir.join(f)) + .filter(|p| p.exists()) + .find_map(|p| { + std::fs::read_to_string(p) + .ok() + .and_then(|source| extract_pat_str_from_source(&source)) + }) +} + +/// Try to download the Python tokenizer source file and extract pat_str via regex. +/// Returns None if unavailable or unparsable. +fn extract_pat_str_from_repo(repo: &crate::hub::HubRepo) -> Option { + // Try common Python tokenizer filenames + let py_path = repo.get("tokenization_kimi.py").or_else(|_| repo.get("tokenizer.py")).ok()?; + + let source = std::fs::read_to_string(&py_path).ok()?; + + // Look for pat_str assignment. Common patterns: + // pat_str = "..." or pat_str = '...' or pat_str = "|".join([...]) + extract_pat_str_from_source(&source) +} + +/// Parse pat_str from Python source code. +fn extract_pat_str_from_source(source: &str) -> Option { + // Strategy: find `pat_str = "|".join([` and collect the raw string fragments + // This handles the common Kimi/Qwen pattern of joining a list of regex strings. + + // First try: look for pat_str = "|".join([...]) pattern + if let Some(join_start) = source.find("pat_str") { + let after = &source[join_start..]; + + // Check for "|".join([ pattern + if let Some(join_pos) = after.find(".join(") { + let after_join = &after[join_pos + 6..]; // skip ".join(" + if let Some(bracket_start) = after_join.find('[') { + let inside = &after_join[bracket_start + 1..]; + // Collect all string literals inside the list + let mut fragments = Vec::new(); + let mut remaining = inside; + + while let Some(frag) = extract_next_python_string(remaining) { + fragments.push(frag.0); + remaining = frag.1; + // Check if we hit the closing bracket + let trimmed = remaining.trim_start(); + if trimmed.starts_with(']') { + break; + } + } + + if !fragments.is_empty() { + let pattern = fragments.join("|"); + println!( + "Extracted pat_str from Python source: {} fragments", + fragments.len() + ); + return Some(pattern); + } + } + } + + // Fallback: simple pat_str = r"..." or pat_str = "..." + if let Some(eq_pos) = after.find('=') { + let after_eq = after[eq_pos + 1..].trim_start(); + if let Some(frag) = extract_next_python_string(after_eq) { + return Some(frag.0); + } + } + } + + None +} + +/// Extract the next Python string literal (r"...", "...", r'''...''', etc.) +/// Returns (string_content, remaining_text) +fn extract_next_python_string(s: &str) -> Option<(String, &str)> { + let s = s.trim_start_matches(|c: char| c == ',' || c.is_whitespace()); + + // Skip comments + if s.starts_with('#') { + let next_line = s.find('\n').map(|i| i + 1).unwrap_or(s.len()); + return extract_next_python_string(&s[next_line..]); + } + + // Check for r""" (triple-quoted raw string) + for prefix in &["r\"\"\"", "r'''"] { + if let Some(inner) = s.strip_prefix(prefix) { + let delim = &prefix[1..]; // """ or ''' + if let Some(end) = inner.find(delim) { + let content = &inner[..end]; + let rest = &inner[end + delim.len()..]; + return Some((content.to_string(), rest)); + } + } + } + + // Check for r"..." (raw string) + if let Some(inner) = s.strip_prefix("r\"") + && let Some(end) = inner.find('"') + { + let content = &inner[..end]; + let rest = &inner[end + 1..]; + return Some((content.to_string(), rest)); + } + if let Some(inner) = s.strip_prefix("r'") + && let Some(end) = inner.find('\'') + { + let content = &inner[..end]; + let rest = &inner[end + 1..]; + return Some((content.to_string(), rest)); + } + + // Check for "..." or '...' (regular string — same as raw for regex patterns) + if let Some(inner) = s.strip_prefix('"') + && let Some(end) = find_unescaped(inner, '"') + { + let content = &inner[..end]; + let rest = &inner[end + 1..]; + return Some((content.to_string(), rest)); + } + if let Some(inner) = s.strip_prefix('\'') + && let Some(end) = find_unescaped(inner, '\'') + { + let content = &inner[..end]; + let rest = &inner[end + 1..]; + return Some((content.to_string(), rest)); + } + + // Check for closing bracket — stop + if s.starts_with(']') { + return None; + } + + None +} + +/// Find position of `ch` that is not preceded by a backslash. +fn find_unescaped(s: &str, ch: char) -> Option { + let mut escaped = false; + for (i, c) in s.char_indices() { + if escaped { + escaped = false; + continue; + } + if c == '\\' { + escaped = true; + continue; + } + if c == ch { + return Some(i); + } + } + None +} diff --git a/rust/src/bench/src/tokenizer.rs b/rust/src/bench/src/tokenizer.rs new file mode 100644 index 000000000000..3d0d4c9c3155 --- /dev/null +++ b/rust/src/bench/src/tokenizer.rs @@ -0,0 +1,300 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::HashSet; +use std::path::Path; + +use tokenizers::Tokenizer; + +use crate::error::{BenchError, Result}; +use crate::tiktoken::TiktokenTokenizer; + +/// Abstraction over local HuggingFace tokenizer, tiktoken, or server-side tokenization. +pub enum TokenizerKind { + Local(Box), + Tiktoken(TiktokenTokenizer), + Server(ServerTokenizer), +} + +/// Server-side tokenizer using vLLM's /tokenize and /detokenize endpoints. +pub struct ServerTokenizer { + client: reqwest::blocking::Client, + tokenize_url: String, + detokenize_url: String, + model: String, + cached_vocab_size: u32, +} + +impl ServerTokenizer { + /// Create a new server tokenizer and verify connectivity. + pub fn new(base_url: &str, model: &str) -> Result { + let client = reqwest::blocking::Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .map_err(|e| BenchError::Tokenizer(format!("Failed to build HTTP client: {e}")))?; + + let tokenize_url = format!("{base_url}/tokenize"); + let detokenize_url = format!("{base_url}/detokenize"); + + let st = Self { + client, + tokenize_url, + detokenize_url, + model: model.to_string(), + cached_vocab_size: 0, + }; + + // Probe the endpoint to verify it works and discover vocab size + let test_tokens = st.encode_inner("test")?; + let max_id = test_tokens.iter().copied().max().unwrap_or(0); + let estimated_vocab = (max_id * 2).max(131072); + + Ok(Self { + cached_vocab_size: estimated_vocab, + ..st + }) + } + + fn encode_inner(&self, text: &str) -> Result> { + let payload = serde_json::json!({ + "model": self.model, + "prompt": text, + }); + + let resp = self + .client + .post(&self.tokenize_url) + .json(&payload) + .send() + .map_err(|e| BenchError::Tokenizer(format!("Server tokenize failed: {e}")))?; + + if !resp.status().is_success() { + return Err(BenchError::Tokenizer(format!( + "Server tokenize returned HTTP {}", + resp.status() + ))); + } + + let data: serde_json::Value = resp.json().map_err(|e| { + BenchError::Tokenizer(format!("Failed to parse tokenize response: {e}")) + })?; + + let tokens = data + .get("tokens") + .and_then(|t| t.as_array()) + .ok_or_else(|| BenchError::Tokenizer("Missing 'tokens' in tokenize response".into()))?; + + tokens + .iter() + .map(|v| { + v.as_u64() + .map(|id| id as u32) + .ok_or_else(|| BenchError::Tokenizer("Invalid token ID in response".into())) + }) + .collect() + } + + fn decode_inner(&self, ids: &[u32]) -> Result { + let payload = serde_json::json!({ + "model": self.model, + "tokens": ids, + }); + + let resp = self + .client + .post(&self.detokenize_url) + .json(&payload) + .send() + .map_err(|e| BenchError::Tokenizer(format!("Server detokenize failed: {e}")))?; + + if !resp.status().is_success() { + return Err(BenchError::Tokenizer(format!( + "Server detokenize returned HTTP {}", + resp.status() + ))); + } + + let data: serde_json::Value = resp.json().map_err(|e| { + BenchError::Tokenizer(format!("Failed to parse detokenize response: {e}")) + })?; + + data.get("prompt") + .and_then(|p| p.as_str()) + .map(|s| s.to_string()) + .ok_or_else(|| BenchError::Tokenizer("Missing 'prompt' in detokenize response".into())) + } +} + +// --- TokenizerKind methods --- + +impl TokenizerKind { + pub fn encode(&self, text: &str, add_special_tokens: bool) -> Result> { + match self { + TokenizerKind::Local(tok) => tok + .encode(text, add_special_tokens) + .map(|enc| enc.get_ids().to_vec()) + .map_err(|e| BenchError::Tokenizer(format!("Encode failed: {e}"))), + TokenizerKind::Tiktoken(tok) => Ok(tok.encode(text)), + TokenizerKind::Server(srv) => srv.encode_inner(text), + } + } + + pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> Result { + match self { + TokenizerKind::Local(tok) => tok + .decode(ids, skip_special_tokens) + .map_err(|e| BenchError::Tokenizer(format!("Decode failed: {e}"))), + TokenizerKind::Tiktoken(tok) => tok.decode(ids), + TokenizerKind::Server(srv) => srv.decode_inner(ids), + } + } + + pub fn vocab_size(&self) -> u32 { + match self { + TokenizerKind::Local(tok) => tok.get_vocab_size(true) as u32, + TokenizerKind::Tiktoken(tok) => tok.vocab_size(), + TokenizerKind::Server(srv) => srv.cached_vocab_size, + } + } + + pub fn num_special_tokens_to_add(&self) -> usize { + match self { + TokenizerKind::Local(tok) => match tok.encode("", true) { + Ok(enc) => enc.get_ids().len(), + Err(_) => 0, + }, + TokenizerKind::Tiktoken(_) | TokenizerKind::Server(_) => 0, + } + } + + pub fn get_allowed_tokens(&self) -> Vec { + match self { + TokenizerKind::Local(tok) => { + let vs = tok.get_vocab_size(true) as u32; + let mut special_ids = HashSet::new(); + for (id, token) in tok.get_added_tokens_decoder() { + if token.special { + special_ids.insert(id); + } + } + (0..vs).filter(|id| !special_ids.contains(id)).collect() + } + TokenizerKind::Tiktoken(tok) => tok.get_allowed_tokens(), + TokenizerKind::Server(srv) => (0..srv.cached_vocab_size).collect(), + } + } +} + +/// Load a tokenizer with fallback chain: +/// 0. Built-in tiktoken encoding (o200k_base, cl100k_base, etc.) — no download needed +/// 1. Local tokenizer.json (HuggingFace fast tokenizer) +/// 2. Tiktoken model file (for Kimi, Qwen, etc.) +/// 3. Server-side /tokenize + /detokenize endpoints +/// +/// `server_info` is `Some((base_url, model))` to enable server-side fallback. +pub fn load_tokenizer( + model_id: &str, + _trust_remote_code: bool, + server_info: Option<(&str, &str)>, +) -> Result { + // 0. Check for built-in tiktoken encoding names (no HF download needed). These are useful for + // consistent cross-model token counting (e.g. Artificial Analysis). + const BUILTIN_TIKTOKEN: &[&str] = &[ + "o200k_base", + "cl100k_base", + "p50k_base", + "p50k_edit", + "r50k_base", + "gpt2", + ]; + if BUILTIN_TIKTOKEN.contains(&model_id) { + return crate::tiktoken::load_builtin_tiktoken(model_id).map(TokenizerKind::Tiktoken); + } + + // 1. Try local HuggingFace tokenizer (tokenizer.json) + match try_load_local(model_id) { + Ok(tok) => { + println!("Tokenizer: Local (vocab_size={})", tok.get_vocab_size(true)); + Ok(TokenizerKind::Local(Box::new(tok))) + } + Err(local_err) => { + // 2. Try tiktoken format + println!("No tokenizer.json for '{model_id}', trying tiktoken format..."); + match crate::tiktoken::try_load_tiktoken(model_id) { + Ok(tok) => { + println!("Tokenizer: Tiktoken (vocab_size={})", tok.vocab_size()); + Ok(TokenizerKind::Tiktoken(tok)) + } + Err(tiktoken_err) => { + // 3. Try server-side fallback + if let Some((base_url, model)) = server_info { + println!( + "Tiktoken also not available ({tiktoken_err}), \ + trying server-side tokenization..." + ); + match ServerTokenizer::new(base_url, model) { + Ok(srv) => { + println!( + "Tokenizer: Server (vocab_size≈{})", + srv.cached_vocab_size + ); + return Ok(TokenizerKind::Server(srv)); + } + Err(srv_err) => { + return Err(BenchError::Tokenizer(format!( + "All tokenizer loading methods failed:\n \ + Local: {local_err}\n \ + Tiktoken: {tiktoken_err}\n \ + Server: {srv_err}\n \ + Try --tokenizer with a model that has tokenizer.json." + ))); + } + } + } + Err(BenchError::Tokenizer(format!( + "Failed to load tokenizer:\n \ + Local: {local_err}\n \ + Tiktoken: {tiktoken_err}\n \ + Try --tokenizer or provide --base-url for server fallback." + ))) + } + } + } + } +} + +/// Try loading tokenizer.json from local path or HuggingFace Hub. +fn try_load_local(model_id: &str) -> Result { + // 1. Try local directory with tokenizer.json + let local_path = Path::new(model_id).join("tokenizer.json"); + if local_path.exists() { + return Tokenizer::from_file(&local_path).map_err(|e| { + BenchError::Tokenizer(format!( + "Failed to load tokenizer from {}: {e}", + local_path.display() + )) + }); + } + + // 2. Try direct path to tokenizer.json + if Path::new(model_id).exists() && model_id.ends_with("tokenizer.json") { + return Tokenizer::from_file(model_id) + .map_err(|e| BenchError::Tokenizer(format!("Failed to load tokenizer: {e}"))); + } + + // 3. If model_id is a local directory, don't try HF Hub — let tiktoken fallback handle it + if Path::new(model_id).is_dir() { + return Err(BenchError::Tokenizer(format!( + "No tokenizer.json in local directory '{model_id}'" + ))); + } + + // 4. Download from HuggingFace Hub (hf-hub handles auth via HF_TOKEN / cached token) + let repo = crate::hub::HubRepo::model(model_id.to_string()); + let tokenizer_path = repo + .get("tokenizer.json") + .map_err(|e| BenchError::Tokenizer(format!("No tokenizer.json for '{model_id}': {e}")))?; + + Tokenizer::from_file(&tokenizer_path) + .map_err(|e| BenchError::Tokenizer(format!("Failed to load downloaded tokenizer: {e}"))) +} diff --git a/rust/src/chat/src/backend/hf.rs b/rust/src/chat/src/backend/hf.rs index f938e1905cb3..d4999d07cd05 100644 --- a/rust/src/chat/src/backend/hf.rs +++ b/rust/src/chat/src/backend/hf.rs @@ -20,6 +20,7 @@ use crate::output::{ use crate::renderer::hf::{HfChatRenderer, MultimodalRenderInfo}; use crate::renderer::{ DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, HarmonyChatRenderer, + InklingChatRenderer, }; use crate::request::ChatRequest; use crate::{DynChatOutputProcessor, RendererSelection}; @@ -71,6 +72,7 @@ impl HfChatBackend { RendererSelection::DeepSeekV32 => Arc::new(DeepSeekV32ChatRenderer::new()), RendererSelection::DeepSeekV4 => Arc::new(DeepSeekV4ChatRenderer::new()), RendererSelection::Harmony => Arc::new(HarmonyChatRenderer::new()?), + RendererSelection::Inkling => Arc::new(InklingChatRenderer::new(tokenizer.clone())?), }; info!( diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 625a2d7818b5..3188657b9096 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -33,7 +33,7 @@ pub use parser::tool::{ToolParser, ToolParserError, ToolParserFactory}; pub use renderer::hf::ChatTemplateContentFormatOption; pub use renderer::{ ChatRenderer, DeepSeekV4ChatRenderer, DeepSeekV32ChatRenderer, DynChatRenderer, - HarmonyChatRenderer, RenderedPrompt, RendererSelection, + HarmonyChatRenderer, InklingChatRenderer, RenderedPrompt, RendererSelection, }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, @@ -294,7 +294,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, inkling, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] @@ -305,6 +305,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, inkling, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/multimodal/audio.rs b/rust/src/chat/src/multimodal/audio.rs index 275963775450..4910f235eea7 100644 --- a/rust/src/chat/src/multimodal/audio.rs +++ b/rust/src/chat/src/multimodal/audio.rs @@ -79,6 +79,38 @@ mod tests { use super::*; const AUDIO_PAD_ID: u32 = 151_676; + const INKLING_AUDIO_MARKER_ID: u32 = 200_020; + const INKLING_AUDIO_EMBED_ID: i32 = 200_053; + + fn inkling_info(decoder_dmodel: serde_json::Value) -> MultimodalModelInfo { + let config = serde_json::json!({ + "model_type": "inkling_mm_model", + "audio_config": { + "decoder_dmodel": decoder_dmodel, + "n_mel_bins": 80, + "mel_vocab_size": 16, + "dmel_min_value": -7.0, + "dmel_max_value": 2.0 + } + }); + let tokenizer = TestTokenizer::new() + .with_regular_token("<|content_image|>", 200_005) + .with_regular_token("<|content_audio_input|>", INKLING_AUDIO_MARKER_ID); + let context = MultimodalModelContext { + model_id: "inkling-test".to_string(), + model_type: Some("inkling_mm_model".to_string()), + config, + tokenizer: TokenizerResolver(Arc::new(tokenizer)), + }; + + MultimodalModelInfo::from_loaded( + context, + PreProcessorConfig::default(), + PreProcessorConfig::default(), + ) + .unwrap() + .expect("Inkling multimodal support") + } fn qwen3_asr_info() -> MultimodalModelInfo { let context = MultimodalModelContext { @@ -120,6 +152,40 @@ mod tests { bytes } + #[test] + fn resolves_inkling_audio_from_model_spec() { + let info = inkling_info(serde_json::json!(1024)); + let support = info.audio.as_ref().expect("audio support"); + + assert_eq!( + info.placeholder_token(Modality::Audio), + Some("<|content_audio_input|>") + ); + assert_eq!(support.placeholder.marker_token_id, INKLING_AUDIO_MARKER_ID); + assert_eq!( + support.placeholder.embed_token_id, + INKLING_AUDIO_EMBED_ID as u32 + ); + assert_eq!(support.spec.primary_key(), AUDIO_PRIMARY_KEY); + assert!(matches!( + &support.spec.field_layouts.encoder_input, + llm_multimodal::FieldLayout::Flat { sizes_key } + if sizes_key == "num_audio_tokens" + )); + assert!(matches!( + support.spec.field_layouts.model_specific.get("num_audio_tokens"), + Some(llm_multimodal::FieldLayout::Batched) + )); + } + + #[test] + fn inkling_decoder_config_gates_audio_capability() { + let info = inkling_info(serde_json::Value::Null); + + assert!(info.audio.is_none()); + assert_eq!(info.placeholder_token(Modality::Audio), None); + } + #[test] fn resolves_qwen_audio_from_model_spec() { let info = qwen3_asr_info(); @@ -210,4 +276,50 @@ mod tests { if tensor.dtype == "int64" && tensor.shape.is_empty() )); } + + #[tokio::test] + async fn inkling_tracker_and_processor_use_standard_audio_key() { + let info = inkling_info(serde_json::json!(1024)); + let wav = wav_i16_mono(16_000, &[0; 1_600]); + let expected_hash = llm_multimodal::hasher::hash_audio(&wav); + let fetched = info + .fetch_media(vec![MediaContentPart::AudioData { + data: wav, + mime_type: Some("audio/wav".to_string()), + uuid: Some("audio-1".to_string()), + }]) + .await + .unwrap(); + + let prepared = info.prepare_audios(fetched.audios, fetched.audio_uuids).await.unwrap(); + + assert_eq!(prepared.replacements.len(), 1); + assert_eq!( + prepared.replacements[0].tokens[0], + INKLING_AUDIO_MARKER_ID as i32 + ); + assert!( + prepared.replacements[0].tokens[1..] + .iter() + .all(|token| *token == INKLING_AUDIO_EMBED_ID) + ); + let item = &prepared.items[0]; + assert_eq!(item.hash, expected_hash); + assert_eq!(item.uuid.as_deref(), Some("audio-1")); + + let features = &item.data[AUDIO_PRIMARY_KEY]; + assert!(matches!(&features.field, MmField::Flat(_))); + assert!(matches!( + features.data.as_ref(), + Some(MmKwargValue::Tensor(tensor)) + if tensor.dtype == "float32" && tensor.shape.get(1) == Some(&80) + )); + let count = &item.data["num_audio_tokens"]; + assert!(matches!(&count.field, MmField::Batched(_))); + assert!(matches!( + count.data.as_ref(), + Some(MmKwargValue::Tensor(tensor)) + if tensor.dtype == "int64" && tensor.shape.is_empty() + )); + } } diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index 10d48b2b395b..ca025268b042 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -23,6 +23,7 @@ pub mod names { pub const DEEPSEEK_V3: &str = "deepseek_v3"; pub const DEEPSEEK_V4: &str = "deepseek_v4"; pub const GEMMA4: &str = "gemma4"; + pub const INKLING: &str = "inkling"; pub const GLM45: &str = "glm45"; pub const KIMI: &str = "kimi"; pub const KIMI_K2: &str = "kimi_k2"; @@ -63,6 +64,7 @@ impl ReasoningParserFactory { .register_parser::(names::DEEPSEEK_V3) .register_parser::(names::DEEPSEEK_V4) .register_unified_dummy(names::GEMMA4) + .register_unified_dummy(names::INKLING) .register_parser::(names::GLM45) .register_parser::(names::KIMI) .register_parser::(names::KIMI_K2) diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 3e394c417bcc..78ccde8a4b34 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -10,7 +10,7 @@ pub use vllm_parser::tool::{ Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser, - Qwen3XmlToolParser, ToolParser, ToolParserError, + Qwen3XmlToolParser, SeedOssToolParser, ToolParser, ToolParserError, }; use crate::parser::ParserFactory; @@ -25,6 +25,7 @@ pub mod names { pub const GLM45: &str = "glm45"; pub const GLM47: &str = "glm47"; pub const GEMMA4: &str = "gemma4"; + pub const INKLING: &str = "inkling"; pub const GRANITE4: &str = "granite4"; pub const HERMES: &str = "hermes"; pub const HY_V3: &str = "hy_v3"; @@ -40,6 +41,7 @@ pub mod names { pub const PHI4_MINI_JSON: &str = "phi4_mini_json"; pub const QWEN3_CODER: &str = "qwen3_coder"; pub const QWEN3_XML: &str = "qwen3_xml"; + pub const SEED_OSS: &str = "seed_oss"; } /// Constructor signature for one registered tool parser implementation. @@ -70,6 +72,7 @@ impl ToolParserFactory { .register_parser::(names::GLM45) .register_parser::(names::GLM47) .register_unified_dummy(names::GEMMA4) + .register_unified_dummy(names::INKLING) .register_parser::(names::GRANITE4) .register_parser::(names::HERMES) .register_parser::(names::HY_V3) @@ -82,7 +85,8 @@ impl ToolParserFactory { .register_parser::(names::MISTRAL) .register_parser::(names::PHI4_MINI_JSON) .register_parser::(names::QWEN3_XML) - .register_parser::(names::QWEN3_CODER); + .register_parser::(names::QWEN3_CODER) + .register_parser::(names::SEED_OSS); factory .register_pattern("mistral-", names::MISTRAL) @@ -120,7 +124,9 @@ impl ToolParserFactory { .register_pattern("minimax-m3", names::MINIMAX_M3) .register_pattern("mm-m3", names::MINIMAX_M3) .register_pattern("minimax", names::MINIMAX_M2) - .register_pattern("mm-m2", names::MINIMAX_M2); + .register_pattern("mm-m2", names::MINIMAX_M2) + .register_pattern("seed-oss", names::SEED_OSS) + .register_pattern("seedoss", names::SEED_OSS); factory } diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs index bb89ea38feb8..c89a50b22455 100644 --- a/rust/src/chat/src/parser/tool/tests.rs +++ b/rust/src/chat/src/parser/tool/tests.rs @@ -176,6 +176,10 @@ fn factory_new_resolves_default_patterns() { factory.resolve_name_for_model("org/mm-m2-base"), Some(names::MINIMAX_M2) ); + assert_eq!( + factory.resolve_name_for_model("ByteDance-Seed/Seed-OSS-36B-Instruct"), + Some(names::SEED_OSS) + ); // InternLM2 positive: both dashed and underscored versioned names route. assert_eq!( diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index ff97bb38db13..2536bf78c252 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -5,7 +5,7 @@ use std::sync::LazyLock; -pub use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParser}; +pub use vllm_parser::unified::{Gemma4UnifiedParser, InklingUnifiedParser, UnifiedParser}; use vllm_tokenizer::DynTokenizer; use crate::parser::ParserFactory; @@ -14,6 +14,7 @@ use crate::request::ChatTool; /// Canonical public names for registered unified parsers. pub mod names { pub const GEMMA4: &str = "gemma4"; + pub const INKLING: &str = "inkling"; } /// Constructor signature for one registered unified parser implementation. @@ -37,10 +38,12 @@ impl UnifiedParserFactory { let mut factory = Self::default(); factory.register_parser::(names::GEMMA4); + factory.register_parser::(names::INKLING); factory .register_pattern("gemma-4", names::GEMMA4) - .register_pattern("gemma4", names::GEMMA4); + .register_pattern("gemma4", names::GEMMA4) + .register_pattern("inkling", names::INKLING); factory } @@ -88,6 +91,13 @@ mod tests { .with_regular_token("", 257) } + fn inkling_tokenizer() -> TestTokenizer { + TestTokenizer::new() + .with_regular_token("<|message_model|>", 200001) + .with_regular_token("<|content_text|>", 200004) + .with_regular_token("<|content_thinking|>", 200008) + } + #[test] fn factory_registers_gemma4() { let factory = UnifiedParserFactory::new(); @@ -99,4 +109,16 @@ mod tests { ); factory.create(names::GEMMA4, &[], Arc::new(tokenizer())).unwrap(); } + + #[test] + fn factory_registers_inkling() { + let factory = UnifiedParserFactory::new(); + + assert!(factory.contains(names::INKLING)); + assert_eq!( + factory.resolve_name_for_model("thinkingmachines/Inkling"), + Some(names::INKLING) + ); + factory.create(names::INKLING, &[], Arc::new(inkling_tokenizer())).unwrap(); + } } diff --git a/rust/src/chat/src/renderer/inkling/fixtures/text_audio_input.json b/rust/src/chat/src/renderer/inkling/fixtures/text_audio_input.json new file mode 100644 index 000000000000..be704bdf1fce --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/fixtures/text_audio_input.json @@ -0,0 +1,21 @@ +{ + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "transcribe" + }, + { + "type": "input_audio", + "data": "" + }, + { + "type": "audio_url", + "audio_url": "data:audio/wav;base64," + } + ] + } + ] +} diff --git a/rust/src/chat/src/renderer/inkling/fixtures/text_audio_output.txt b/rust/src/chat/src/renderer/inkling/fixtures/text_audio_output.txt new file mode 100644 index 000000000000..659e2614e2f3 --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/fixtures/text_audio_output.txt @@ -0,0 +1 @@ +<|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_user|><|content_text|>transcribe<|end_message|><|message_user|><|content_audio_input|><|audio_end|><|end_message|><|message_user|><|content_audio_input|><|audio_end|><|end_message|><|message_model|> diff --git a/rust/src/chat/src/renderer/inkling/fixtures/text_image_input.json b/rust/src/chat/src/renderer/inkling/fixtures/text_image_input.json new file mode 100644 index 000000000000..d0afac640138 --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/fixtures/text_image_input.json @@ -0,0 +1,17 @@ +{ + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "look" + }, + { + "type": "image_url", + "image_url": "data:image/png;base64," + } + ] + } + ] +} diff --git a/rust/src/chat/src/renderer/inkling/fixtures/text_image_output.txt b/rust/src/chat/src/renderer/inkling/fixtures/text_image_output.txt new file mode 100644 index 000000000000..28a83a0ee846 --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/fixtures/text_image_output.txt @@ -0,0 +1 @@ +<|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_user|><|content_text|>look<|end_message|><|message_user|><|content_image|><|end_message|><|message_model|> diff --git a/rust/src/chat/src/renderer/inkling/fixtures/tool_declare_input.json b/rust/src/chat/src/renderer/inkling/fixtures/tool_declare_input.json new file mode 100644 index 000000000000..a1ca1ca3ccad --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/fixtures/tool_declare_input.json @@ -0,0 +1,46 @@ +{ + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather information", + "parameters": { + "type": "object", + "required": [ + "city" + ], + "properties": { + "city": { + "type": "string" + } + } + } + } + } + ], + "messages": [ + { + "role": "developer", + "content": "rules", + "tools": [ + { + "type": "function", + "function": { + "name": "local_tool", + "parameters": { + "z": 1, + "a": { + "b": 2 + } + } + } + } + ] + }, + { + "role": "user", + "content": "hi" + } + ] +} diff --git a/rust/src/chat/src/renderer/inkling/fixtures/tool_declare_output.txt b/rust/src/chat/src/renderer/inkling/fixtures/tool_declare_output.txt new file mode 100644 index 000000000000..5a3c890325d1 --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/fixtures/tool_declare_output.txt @@ -0,0 +1 @@ +<|message_system|>tool_declare<|content_xml|>[{"description":"Get weather information","name":"get_weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"},"type":"function"},{"description":"","name":"local_tool","parameters":{"a":{"b":2},"z":1},"type":"function"}]<|end_message|><|message_system|><|content_text|>rules<|end_message|><|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_user|><|content_text|>hi<|end_message|><|message_model|> diff --git a/rust/src/chat/src/renderer/inkling/fixtures/tool_round_trip_input.json b/rust/src/chat/src/renderer/inkling/fixtures/tool_round_trip_input.json new file mode 100644 index 000000000000..239e1348f18e --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/fixtures/tool_round_trip_input.json @@ -0,0 +1,24 @@ +{ + "add_generation_prompt": false, + "messages": [ + { + "role": "assistant", + "reasoning_content": "think", + "content": "answer", + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "get_weather", + "arguments": "{\"city\":\"SF\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "sunny" + } + ] +} diff --git a/rust/src/chat/src/renderer/inkling/fixtures/tool_round_trip_output.txt b/rust/src/chat/src/renderer/inkling/fixtures/tool_round_trip_output.txt new file mode 100644 index 000000000000..f5069f13bedd --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/fixtures/tool_round_trip_output.txt @@ -0,0 +1 @@ +<|message_system|><|content_text|>Thinking effort level: 0.9<|end_message|><|message_model|><|content_thinking|>think<|end_message|><|message_model|><|content_text|>answer<|end_message|><|message_model|>get_weather<|content_invoke_tool_json|>{"name":"get_weather","args":{"city":"SF"}}<|end_message|><|content_model_end_sampling|><|message_tool|>get_weather<|content_text|>sunny<|end_message|> diff --git a/rust/src/chat/src/renderer/inkling/mod.rs b/rust/src/chat/src/renderer/inkling/mod.rs new file mode 100644 index 000000000000..8bf1a204fbf2 --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/mod.rs @@ -0,0 +1,442 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::collections::HashMap; + +use serde_json::{Map, Value, json}; +use thiserror_ext::AsReport as _; +use vllm_text::Prompt; +use vllm_text::tokenizer::{DynTokenizer, Tokenizer}; + +use super::{ChatRenderer, RenderedPrompt, request_template_kwargs}; +use crate::error::{Error, Result}; +use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool}; +use crate::{AssistantContentBlock, AssistantToolCall}; + +const MESSAGE_USER: &str = "<|message_user|>"; +const MESSAGE_MODEL: &str = "<|message_model|>"; +const MESSAGE_SYSTEM: &str = "<|message_system|>"; +const MESSAGE_TOOL: &str = "<|message_tool|>"; +const CONTENT_TEXT: &str = "<|content_text|>"; +const CONTENT_IMAGE: &str = "<|content_image|>"; +const CONTENT_MODEL_END_SAMPLING: &str = "<|content_model_end_sampling|>"; +const CONTENT_AUDIO_INPUT: &str = "<|content_audio_input|>"; +const CONTENT_THINKING: &str = "<|content_thinking|>"; +// Inkling renderer semantics use this slot for structured payloads such as tool +// declarations. +const CONTENT_XML: &str = "<|content_xml|>"; +const CONTENT_INVOKE_TOOL_JSON: &str = "<|content_invoke_tool_json|>"; +const END_MESSAGE: &str = "<|end_message|>"; +const AUDIO_END: &str = "<|audio_end|>"; +const MAX_REASONING_EFFORT: f64 = 0.99; + +/// Native Inkling renderer that emits token IDs directly. +#[derive(Clone)] +pub struct InklingChatRenderer { + tokenizer: DynTokenizer, + special: InklingSpecialTokenIds, +} + +#[derive(Debug, Clone, Copy)] +struct InklingSpecialTokenIds { + message_user: u32, + message_model: u32, + message_system: u32, + message_tool: u32, + content_text: u32, + content_image: u32, + content_model_end_sampling: u32, + content_audio_input: u32, + content_thinking: u32, + content_xml: u32, + content_invoke_tool_json: u32, + end_message: u32, + audio_end: u32, +} + +impl InklingSpecialTokenIds { + fn resolve(tokenizer: &dyn Tokenizer) -> Result { + Ok(Self { + message_user: resolve_special_token(tokenizer, MESSAGE_USER)?, + message_model: resolve_special_token(tokenizer, MESSAGE_MODEL)?, + message_system: resolve_special_token(tokenizer, MESSAGE_SYSTEM)?, + message_tool: resolve_special_token(tokenizer, MESSAGE_TOOL)?, + content_text: resolve_special_token(tokenizer, CONTENT_TEXT)?, + content_image: resolve_special_token(tokenizer, CONTENT_IMAGE)?, + content_model_end_sampling: resolve_special_token( + tokenizer, + CONTENT_MODEL_END_SAMPLING, + )?, + content_audio_input: resolve_special_token(tokenizer, CONTENT_AUDIO_INPUT)?, + content_thinking: resolve_special_token(tokenizer, CONTENT_THINKING)?, + content_xml: resolve_special_token(tokenizer, CONTENT_XML)?, + content_invoke_tool_json: resolve_special_token(tokenizer, CONTENT_INVOKE_TOOL_JSON)?, + end_message: resolve_special_token(tokenizer, END_MESSAGE)?, + audio_end: resolve_special_token(tokenizer, AUDIO_END)?, + }) + } +} + +impl InklingChatRenderer { + pub fn new(tokenizer: DynTokenizer) -> Result { + let special = InklingSpecialTokenIds::resolve(tokenizer.as_ref())?; + Ok(Self { tokenizer, special }) + } + + fn write_text_tokens(&self, out: &mut Vec, text: &str) -> Result<()> { + out.extend(self.tokenizer.encode(text, false)?); + Ok(()) + } + + fn write_message_start( + &self, + out: &mut Vec, + role_token_id: u32, + author_name: Option<&str>, + ) -> Result<()> { + out.push(role_token_id); + if let Some(author_name) = author_name + && !author_name.is_empty() + { + self.write_text_tokens(out, author_name)?; + } + Ok(()) + } + + fn write_text_block( + &self, + out: &mut Vec, + role_token_id: u32, + author_name: Option<&str>, + text: &str, + ) -> Result<()> { + self.write_message_start(out, role_token_id, author_name)?; + out.push(self.special.content_text); + self.write_text_tokens(out, text)?; + out.push(self.special.end_message); + Ok(()) + } + + /// Write one image block holding only the `<|content_image|>` marker. + /// Multimodal preprocessing later expands the marker into per-patch + /// image placeholder tokens once the patch count is known, mirroring the + /// Python `InklingMultiModalProcessor` marker-anchored prompt updates. + fn write_image_block(&self, out: &mut Vec, role_token_id: u32) { + out.push(role_token_id); + out.push(self.special.content_image); + out.push(self.special.end_message); + } + + /// Write one audio block holding the `<|content_audio_input|>` marker + /// followed by the `<|audio_end|>` terminator. Multimodal preprocessing + /// later expands the marker into per-frame audio placeholder tokens + /// (landing before `<|audio_end|>`) once the clip length is known. + fn write_audio_block(&self, out: &mut Vec, role_token_id: u32) { + out.push(role_token_id); + out.push(self.special.content_audio_input); + out.push(self.special.audio_end); + out.push(self.special.end_message); + } + + fn write_reasoning_block(&self, out: &mut Vec, text: &str) -> Result<()> { + if text.is_empty() { + return Ok(()); + } + out.push(self.special.message_model); + out.push(self.special.content_thinking); + self.write_text_tokens(out, text)?; + out.push(self.special.end_message); + Ok(()) + } + + fn write_reasoning_effort(&self, out: &mut Vec, effort: f64) -> Result<()> { + if !(0.0..=MAX_REASONING_EFFORT).contains(&effort) { + return Err(Error::ChatTemplate(format!( + "Inkling reasoning_effort must be in [0.0, 0.99], got {effort}" + ))); + } + let formatted = format!("{effort:.2}"); + let effort = formatted.trim_end_matches('0').trim_end_matches('.'); + let effort = if matches!(effort, "0" | "-0") { + "0.0" + } else { + effort + }; + self.write_text_block( + out, + self.special.message_system, + None, + &format!("Thinking effort level: {effort}"), + ) + } + + fn write_tool_declarations(&self, out: &mut Vec, tools: &[&ChatTool]) -> Result<()> { + if tools.is_empty() { + return Ok(()); + } + + let mut specs = Vec::with_capacity(tools.len()); + for tool in tools { + specs.push(json!({ + "description": tool.description.as_deref().unwrap_or(""), + "name": tool.name, + "parameters": sort_json(&tool.parameters), + "type": "function", + })); + } + let payload = compact_json(&sort_json(&Value::Array(specs)))?; + + self.write_message_start(out, self.special.message_system, Some("tool_declare"))?; + out.push(self.special.content_xml); + self.write_text_tokens(out, &payload)?; + out.push(self.special.end_message); + Ok(()) + } + + fn write_chat_content( + &self, + out: &mut Vec, + role_token_id: u32, + content: &ChatContent, + ) -> Result<()> { + match content { + ChatContent::Text(text) => { + if !text.is_empty() { + self.write_text_block(out, role_token_id, None, text)?; + } + } + ChatContent::Parts(parts) => { + for part in parts { + match part { + ChatContentPart::Text { text } => { + self.write_text_block(out, role_token_id, None, text)?; + } + ChatContentPart::ImageUrl { .. } => { + self.write_image_block(out, role_token_id); + } + ChatContentPart::InputAudio { .. } | ChatContentPart::AudioUrl { .. } => { + self.write_audio_block(out, role_token_id); + } + // Inkling has no video modality. + ChatContentPart::VideoUrl { .. } => { + return Err(Error::UnsupportedMultimodalContent("video_url")); + } + } + } + } + } + Ok(()) + } + + fn write_assistant_tool_call( + &self, + out: &mut Vec, + tool_call: &AssistantToolCall, + ) -> Result<()> { + let payload = tool_call_json(tool_call)?; + + self.write_message_start(out, self.special.message_model, Some(&tool_call.name))?; + out.push(self.special.content_invoke_tool_json); + self.write_text_tokens(out, &payload)?; + out.push(self.special.end_message); + Ok(()) + } + + fn write_assistant_content( + &self, + out: &mut Vec, + content: &[AssistantContentBlock], + tool_call_id_to_name: &mut HashMap, + ) -> Result<()> { + for block in content { + match block { + AssistantContentBlock::Reasoning { text } => { + self.write_reasoning_block(out, text)?; + } + AssistantContentBlock::Text { text } => { + if !text.is_empty() { + self.write_text_block(out, self.special.message_model, None, text)?; + } + } + AssistantContentBlock::ToolCall(tool_call) => { + if !tool_call.id.is_empty() { + tool_call_id_to_name.insert(tool_call.id.clone(), tool_call.name.clone()); + } + self.write_assistant_tool_call(out, tool_call)?; + } + } + } + out.push(self.special.content_model_end_sampling); + Ok(()) + } + + fn write_tool_response( + &self, + out: &mut Vec, + content: &ChatContent, + tool_call_id: &str, + tool_call_id_to_name: &HashMap, + ) -> Result<()> { + let text = content.try_flatten_to_text()?; + let tool_name = tool_call_id_to_name.get(tool_call_id).map(String::as_str).unwrap_or(""); + self.write_text_block(out, self.special.message_tool, Some(tool_name), &text) + } +} + +impl ChatRenderer for InklingChatRenderer { + fn render(&self, request: &ChatRequest) -> Result { + request.validate()?; + if request.chat_options.continue_final_message() { + return Err(Error::ChatTemplate( + "Inkling renderer does not support continue_final_message".to_string(), + )); + } + + let mut out = Vec::new(); + let mut tool_call_id_to_name = HashMap::new(); + let effective_template_kwargs = request_template_kwargs(request); + let tools = rendered_tools(request); + self.write_tool_declarations(&mut out, &tools)?; + let mut reasoning_effort = + resolve_reasoning_effort(effective_template_kwargs.get("reasoning_effort")); + + for message in &request.messages { + if !matches!( + message, + ChatMessage::System { .. } | ChatMessage::Developer { .. } + ) && let Some(effort) = reasoning_effort.take() + { + self.write_reasoning_effort(&mut out, effort)?; + } + + match message { + ChatMessage::System { content } => { + self.write_chat_content(&mut out, self.special.message_system, content)?; + } + ChatMessage::Developer { content, .. } => { + self.write_chat_content(&mut out, self.special.message_system, content)?; + } + ChatMessage::User { content } => { + self.write_chat_content(&mut out, self.special.message_user, content)?; + } + ChatMessage::Assistant { content } => { + self.write_assistant_content(&mut out, content, &mut tool_call_id_to_name)?; + } + ChatMessage::ToolResponse { + content, + tool_call_id, + } => { + self.write_tool_response( + &mut out, + content, + tool_call_id, + &tool_call_id_to_name, + )?; + } + } + } + + if let Some(effort) = reasoning_effort { + self.write_reasoning_effort(&mut out, effort)?; + } + + if request.chat_options.add_generation_prompt() { + out.push(self.special.message_model); + } + + Ok(RenderedPrompt { + prompt: Prompt::TokenIds(out), + effective_template_kwargs, + }) + } +} + +fn resolve_reasoning_effort(value: Option<&Value>) -> Option { + let Some(value) = value else { + return Some(0.9); + }; + match value { + Value::String(name) => match name.as_str() { + "none" => Some(0.0), + "minimal" => Some(0.1), + "low" => Some(0.2), + "medium" => Some(0.7), + "high" => Some(0.9), + "xhigh" | "max" => Some(0.99), + _ => None, + }, + Value::Number(number) => number.as_f64(), + _ => None, + } +} + +fn resolve_special_token(tokenizer: &dyn Tokenizer, token: &str) -> Result { + tokenizer.token_to_id(token).ok_or_else(|| { + Error::ChatTemplate(format!( + "Inkling tokenizer is missing special token `{token}`" + )) + }) +} + +fn rendered_tools(request: &ChatRequest) -> Vec<&ChatTool> { + if !request.tool_parsing_enabled() { + return Vec::new(); + } + + let mut tools = Vec::with_capacity(request.tools.len()); + tools.extend(request.tools.iter()); + for message in &request.messages { + if let ChatMessage::Developer { + tools: Some(local_tools), + .. + } = message + { + tools.extend(local_tools.iter()); + } + } + tools +} + +fn tool_call_json(tool_call: &AssistantToolCall) -> Result { + let name_json = serde_json::to_string(&tool_call.name) + .map_err(|error| Error::ChatTemplate(error.as_report().to_string()))?; + let arguments = if tool_call.arguments.trim().is_empty() { + Value::Object(Map::new()) + } else { + serde_json::from_str(&tool_call.arguments).map_err(|error| { + Error::ChatTemplate(format!( + "Inkling tool call arguments must decode to a JSON object: {error}" + )) + })? + }; + let Value::Object(_) = arguments else { + return Err(Error::ChatTemplate( + "Inkling tool call arguments must decode to a JSON object".to_string(), + )); + }; + let args_json = compact_json(&sort_json(&arguments))?; + Ok(format!("{{\"name\":{name_json},\"args\":{args_json}}}")) +} + +fn compact_json(value: &Value) -> Result { + serde_json::to_string(value).map_err(|error| Error::ChatTemplate(error.as_report().to_string())) +} + +fn sort_json(value: &Value) -> Value { + match value { + Value::Array(items) => Value::Array(items.iter().map(sort_json).collect()), + Value::Object(map) => { + let mut sorted = Map::new(); + let mut keys = map.keys().collect::>(); + keys.sort(); + for key in keys { + sorted.insert(key.clone(), sort_json(&map[key])); + } + Value::Object(sorted) + } + _ => value.clone(), + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/chat/src/renderer/inkling/tests.rs b/rust/src/chat/src/renderer/inkling/tests.rs new file mode 100644 index 000000000000..d0d55e1be2d5 --- /dev/null +++ b/rust/src/chat/src/renderer/inkling/tests.rs @@ -0,0 +1,375 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::path::PathBuf; +use std::sync::Arc; + +use expect_test::{ExpectFile, expect_file}; +use serde_json::json; +use thiserror_ext::AsReport; +use vllm_text::tokenizer::Tokenizer; + +use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; + +use super::{ + AUDIO_END, CONTENT_AUDIO_INPUT, CONTENT_IMAGE, CONTENT_INVOKE_TOOL_JSON, CONTENT_TEXT, + CONTENT_THINKING, CONTENT_XML, END_MESSAGE, InklingChatRenderer, MESSAGE_MODEL, MESSAGE_SYSTEM, + MESSAGE_TOOL, MESSAGE_USER, +}; +use crate::event::{AssistantContentBlock, AssistantToolCall}; +use crate::request::{ChatMessage, ChatRequest, GenerationPromptMode, ReasoningEffort}; +use crate::{ChatRenderer, Error}; + +struct FixtureTokenizer; + +impl Tokenizer for FixtureTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_text::tokenizer::Result> { + Ok(text.bytes().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_text::tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + "<|message_user|>" => Some(200000), + "<|message_model|>" => Some(200001), + "<|message_system|>" => Some(200002), + "<|message_tool|>" => Some(200003), + "<|content_text|>" => Some(200004), + "<|content_image|>" => Some(200005), + "<|content_model_end_sampling|>" => Some(200006), + "<|content_thinking|>" => Some(200008), + "<|end_message|>" => Some(200010), + "<|content_audio_input|>" => Some(200020), + CONTENT_XML => Some(200024), + "<|audio_end|>" => Some(200043), + "<|content_invoke_tool_json|>" => Some(200049), + _ => None, + } + } + + fn id_to_token(&self, id: u32) -> Option { + let token = match id { + 200000 => "<|message_user|>", + 200001 => "<|message_model|>", + 200002 => "<|message_system|>", + 200003 => "<|message_tool|>", + 200004 => "<|content_text|>", + 200005 => "<|content_image|>", + 200006 => "<|content_model_end_sampling|>", + 200008 => "<|content_thinking|>", + 200010 => "<|end_message|>", + 200020 => "<|content_audio_input|>", + 200024 => CONTENT_XML, + 200043 => "<|audio_end|>", + 200049 => "<|content_invoke_tool_json|>", + _ => return None, + }; + Some(token.to_string()) + } +} + +fn renderer() -> InklingChatRenderer { + InklingChatRenderer::new(Arc::new(FixtureTokenizer)).unwrap() +} + +fn render_token_ids(request: &ChatRequest) -> Vec { + renderer() + .render(request) + .unwrap() + .prompt + .into_token_ids() + .expect("Inkling renderer returns token IDs") +} + +fn fixture_request(name: &str) -> ChatRequest { + fixture_chat_request(&fixture_path(name), inkling_fixture_options()) +} + +fn inkling_fixture_options() -> FixtureRequestOptions { + FixtureRequestOptions { + enable_thinking: false, + no_generation_prompt_when_last_assistant: false, + } +} + +fn fixture_path(name: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src/renderer/inkling") + .join("fixtures") + .join(name) +} + +fn assert_fixture(input_name: &str, expected: ExpectFile) { + let request = fixture_request(input_name); + let rendered = format!("{}\n", render_symbolic_tokens(&render_token_ids(&request))); + + expected.assert_eq(&rendered); +} + +fn render_symbolic_tokens(token_ids: &[u32]) -> String { + let mut out = String::new(); + let mut text_bytes = Vec::new(); + for token_id in token_ids { + if let Some(marker) = symbolic_tokens() + .iter() + .find_map(|(marker, marker_id)| (token_id == marker_id).then_some(*marker)) + { + flush_text_bytes(&mut out, &mut text_bytes); + out.push_str(marker); + continue; + } + + let byte = u8::try_from(*token_id) + .unwrap_or_else(|_| panic!("unexpected non-byte fixture text token {token_id}")); + text_bytes.push(byte); + } + flush_text_bytes(&mut out, &mut text_bytes); + out +} + +fn flush_text_bytes(out: &mut String, text_bytes: &mut Vec) { + if text_bytes.is_empty() { + return; + } + out.push_str(std::str::from_utf8(text_bytes).unwrap()); + text_bytes.clear(); +} + +fn symbolic_tokens() -> [(&'static str, u32); 13] { + [ + (MESSAGE_USER, 200000), + (MESSAGE_MODEL, 200001), + (MESSAGE_SYSTEM, 200002), + (MESSAGE_TOOL, 200003), + (CONTENT_TEXT, 200004), + (CONTENT_IMAGE, 200005), + (super::CONTENT_MODEL_END_SAMPLING, 200006), + (CONTENT_THINKING, 200008), + (END_MESSAGE, 200010), + (CONTENT_AUDIO_INPUT, 200020), + (CONTENT_XML, 200024), + (AUDIO_END, 200043), + (CONTENT_INVOKE_TOOL_JSON, 200049), + ] +} + +#[test] +fn renders_text_and_image_as_inkling_tokens() { + assert_fixture( + "text_image_input.json", + expect_file!["fixtures/text_image_output.txt"], + ); +} + +#[test] +fn renders_audio_marker_only_blocks() { + assert_fixture( + "text_audio_input.json", + expect_file!["fixtures/text_audio_output.txt"], + ); +} + +#[test] +fn renders_reasoning_text_tool_call_and_tool_response() { + assert_fixture( + "tool_round_trip_input.json", + expect_file!["fixtures/tool_round_trip_output.txt"], + ); +} + +#[test] +fn renders_request_and_developer_tools_as_tool_declare() { + assert_fixture( + "tool_declare_input.json", + expect_file!["fixtures/tool_declare_output.txt"], + ); +} + +#[test] +fn renders_named_reasoning_effort_after_tool_declarations() { + for (effort, expected) in [ + (ReasoningEffort::None, "0.0"), + (ReasoningEffort::Minimal, "0.1"), + (ReasoningEffort::Low, "0.2"), + (ReasoningEffort::Medium, "0.7"), + (ReasoningEffort::High, "0.9"), + (ReasoningEffort::XHigh, "0.99"), + (ReasoningEffort::Max, "0.99"), + ] { + let mut request = fixture_request("tool_declare_input.json"); + request.chat_options.reasoning_effort = Some(effort); + + let rendered = render_symbolic_tokens(&render_token_ids(&request)); + let tool_end = rendered.find("<|end_message|>").unwrap(); + let effort_block = format!( + "<|message_system|><|content_text|>Thinking effort level: \ + {expected}<|end_message|>" + ); + let effort_start = rendered.find(&effort_block).unwrap(); + let system_end = rendered.find("rules<|end_message|>").unwrap(); + + assert!(effort_start > tool_end); + assert!(effort_start > system_end); + } +} + +#[test] +fn emits_one_reasoning_effort_for_multi_turn_conversation() { + let request = ChatRequest { + messages: vec![ + ChatMessage::system("rules"), + ChatMessage::user("user1"), + ChatMessage::assistant_text("assistant1"), + ChatMessage::user("user2"), + ], + ..ChatRequest::for_test() + }; + + let rendered = render_symbolic_tokens(&render_token_ids(&request)); + let effort = "<|message_system|><|content_text|>Thinking effort level: 0.9\ + <|end_message|>"; + + assert_eq!(rendered.matches(effort).count(), 1); + assert!(rendered.find("rules").unwrap() < rendered.find(effort).unwrap()); + assert!(rendered.find(effort).unwrap() < rendered.find("user1").unwrap()); + assert!(rendered.find("user1").unwrap() < rendered.find("assistant1").unwrap()); + assert!(rendered.find("assistant1").unwrap() < rendered.find("user2").unwrap()); +} + +#[test] +fn canonicalizes_numeric_zero_reasoning_effort() { + for value in [0.0, -0.0] { + let mut request = ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!(value)); + + assert!( + render_symbolic_tokens(&render_token_ids(&request)).starts_with( + "<|message_system|><|content_text|>Thinking effort level: 0.0\ + <|end_message|>" + ) + ); + } +} + +#[test] +fn defaults_reasoning_effort_to_high() { + let request = ChatRequest::for_test(); + + assert!( + render_symbolic_tokens(&render_token_ids(&request)).starts_with( + "<|message_system|><|content_text|>Thinking effort level: 0.9\ + <|end_message|>" + ) + ); +} + +#[test] +fn renders_numeric_reasoning_effort_template_kwarg() { + let mut request = ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!(0.8)); + + assert_eq!( + render_symbolic_tokens(&render_token_ids(&request)), + "<|message_system|><|content_text|>Thinking effort level: 0.8\ + <|end_message|><|message_user|><|content_text|>test\ + <|end_message|><|message_model|>" + ); +} + +#[test] +fn ignores_unsupported_reasoning_effort_values() { + for value in [json!(true), json!("invalid"), json!(null)] { + let mut request = ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), value); + + assert_eq!( + render_symbolic_tokens(&render_token_ids(&request)), + "<|message_user|><|content_text|>test<|end_message|><|message_model|>" + ); + } +} + +#[test] +fn rejects_out_of_range_reasoning_effort() { + for value in [0.990_000_1, 1.0, 1.5, -0.1] { + let mut request = ChatRequest::for_test(); + request + .chat_options + .template_kwargs + .insert("reasoning_effort".to_string(), json!(value)); + + let error = renderer().render(&request).unwrap_err(); + assert!(error.as_report().to_string().contains("must be in [0.0, 0.99]")); + } +} + +#[test] +fn rejects_continue_final_message() { + let request = ChatRequest { + messages: vec![ChatMessage::assistant_text("partial")], + chat_options: crate::ChatOptions { + generation_prompt_mode: GenerationPromptMode::ContinueFinalAssistant, + ..Default::default() + }, + ..ChatRequest::for_test() + }; + + let error = renderer().render(&request).unwrap_err(); + assert!(matches!(error, Error::ChatTemplate(_))); + assert!(error.as_report().to_string().contains("continue_final_message")); +} + +#[test] +fn renders_developer_messages_as_system() { + let request = ChatRequest { + messages: vec![ChatMessage::developer("rules", None)], + ..ChatRequest::for_test() + }; + + assert_eq!( + render_symbolic_tokens(&render_token_ids(&request)), + "<|message_system|><|content_text|>rules<|end_message|>\ + <|message_system|><|content_text|>Thinking effort level: 0.9\ + <|end_message|><|message_model|>" + ); +} + +#[test] +fn rejects_non_object_tool_call_arguments() { + let request = ChatRequest { + messages: vec![ChatMessage::assistant_blocks(vec![ + AssistantContentBlock::ToolCall(AssistantToolCall { + id: "call_1".to_string(), + name: "get_weather".to_string(), + arguments: "[]".to_string(), + }), + ])], + ..ChatRequest::for_test() + }; + + let error = renderer().render(&request).unwrap_err(); + assert!(error.as_report().to_string().contains("JSON object")); +} diff --git a/rust/src/chat/src/renderer/mod.rs b/rust/src/chat/src/renderer/mod.rs index 639d6df30e59..4d9c1581a01e 100644 --- a/rust/src/chat/src/renderer/mod.rs +++ b/rust/src/chat/src/renderer/mod.rs @@ -14,6 +14,7 @@ pub mod deepseek_v32; pub mod deepseek_v4; pub mod harmony; pub mod hf; +mod inkling; mod selection; #[cfg(test)] mod test_utils; @@ -21,6 +22,7 @@ mod test_utils; pub use deepseek_v4::DeepSeekV4ChatRenderer; pub use deepseek_v32::DeepSeekV32ChatRenderer; pub use harmony::HarmonyChatRenderer; +pub use inkling::InklingChatRenderer; pub use selection::RendererSelection; /// Rendered chat prompt submitted to the text backend. diff --git a/rust/src/chat/src/renderer/selection.rs b/rust/src/chat/src/renderer/selection.rs index 508bf8d17c18..711c0d87ddfb 100644 --- a/rust/src/chat/src/renderer/selection.rs +++ b/rust/src/chat/src/renderer/selection.rs @@ -24,6 +24,8 @@ pub enum RendererSelection { DeepSeekV4, /// Force the GPT-OSS Harmony renderer. Harmony, + /// Force the Inkling native token renderer. + Inkling, } impl RendererSelection { @@ -33,6 +35,8 @@ impl RendererSelection { pub const GPT_OSS_MODEL_TYPE: &str = "gpt_oss"; pub const HARMONY_LITERAL: &str = "harmony"; pub const HF_LITERAL: &str = "hf"; + pub const INKLING_LITERAL: &str = "inkling"; + pub const INKLING_MODEL_TYPE: &str = "inkling_mm_model"; /// Resolve the renderer selection using the given model type string, if /// it's `Auto`. @@ -42,6 +46,7 @@ impl RendererSelection { Self::DEEPSEEK_V32_LITERAL => Self::DeepSeekV32, Self::DEEPSEEK_V4_LITERAL => Self::DeepSeekV4, Self::GPT_OSS_MODEL_TYPE => Self::Harmony, + Self::INKLING_MODEL_TYPE => Self::Inkling, _ => Self::Hf, }, selection => selection, @@ -63,6 +68,8 @@ impl FromStr for RendererSelection { Ok(Self::DeepSeekV4) } else if value.eq_ignore_ascii_case(Self::HARMONY_LITERAL) { Ok(Self::Harmony) + } else if value.eq_ignore_ascii_case(Self::INKLING_LITERAL) { + Ok(Self::Inkling) } else { Err(format!( "unknown renderer `{value}` (expected one of: {})", @@ -80,6 +87,7 @@ impl fmt::Display for RendererSelection { Self::DeepSeekV32 => f.write_str(Self::DEEPSEEK_V32_LITERAL), Self::DeepSeekV4 => f.write_str(Self::DEEPSEEK_V4_LITERAL), Self::Harmony => f.write_str(Self::HARMONY_LITERAL), + Self::Inkling => f.write_str(Self::INKLING_LITERAL), } } } @@ -106,7 +114,7 @@ mod tests { fn renderer_selection_expected_error_message() { let err = RendererSelection::from_str("unknown").unwrap_err(); expect_test::expect![ - "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony)" + "unknown renderer `unknown` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling)" ] .assert_eq(&err); } diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index a12531150c81..08eeca8176d4 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -103,6 +103,8 @@ pub(crate) enum FixtureContent { pub(crate) enum FixtureContentPart { Text { text: String }, ImageUrl { image_url: String }, + InputAudio { data: String }, + AudioUrl { audio_url: String }, } #[derive(Debug, Deserialize)] @@ -226,6 +228,15 @@ fn to_chat_content(content: FixtureContent) -> ChatContent { FixtureContentPart::ImageUrl { image_url } => { ChatContentPart::image_url(image_url) } + FixtureContentPart::InputAudio { data } => ChatContentPart::InputAudio { + data, + format: None, + uuid: None, + }, + FixtureContentPart::AudioUrl { audio_url } => ChatContentPart::AudioUrl { + audio_url, + uuid: None, + }, }) .collect(), ), diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 58a3adfc4eec..ff15aa929ef1 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -263,19 +263,32 @@ impl RoundtripCase { sort_json_keys: false, } } + + /// Inkling typed content blocks with native token-id rendering. + fn inkling() -> Self { + Self { + model_id: "thinkingmachines/Inkling", + assistant_stop_suffix: "", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Always { value: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: true, + } + } } macro_rules! roundtrip_tests { - ($($case:ident => [$($(#[$fixture_attr:meta])* $fixture:ident),* $(,)?]),+ $(,)?) => { + ($($case:ident => $(#[$case_attr:meta])* [$($fixture:ident),* $(,)?]),+ $(,)?) => { paste::paste! { $( #[tokio::test] #[file_serial([])] + $(#[$case_attr])* async fn []() -> Result<()> { let case = RoundtripCase::$case(); let backends = load_roundtrip_backends(&case).await?; $( - $(#[$fixture_attr])* [](&case, &backends).await?; )* Ok(()) @@ -294,12 +307,13 @@ roundtrip_tests! { deepseek_v32 => [tool_call_mix], glm45 => [reasoning_and_content, tool_call_mix], glm47 => [reasoning_and_content, tool_call_mix], - seed_oss => [reasoning_and_content], + seed_oss => [reasoning_and_content, tool_call_mix], step3p5 => [reasoning_and_content], nemotron_v3 => [reasoning_and_content], gemma4 => [tool_call_mix], // Gemma4 strips reasoning in history if there's no tool call kimi_k25 => [tool_call_mix], // Kimi K2.5 strips reasoning in history gpt_oss => [tool_call_mix], // Harmony strips reasoning in history if there's no tool call + inkling => [reasoning_and_content, tool_call_mix], } /// Run the fixed reasoning+content fixture for one model/parser case. diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index f3af2a18fd2f..ee0d3ebf0486 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -646,7 +646,7 @@ fn serve_args_reject_unknown_renderer_value() { .unwrap_err(); expect![[r#" - error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony) + error: invalid value 'definitely_missing' for '--tokenizer-mode ': unknown renderer `definitely_missing` (expected one of: auto, hf, deepseek_v32, deepseek_v4, harmony, inkling) For more information, try '--help'. "#]] diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 628104e9c337..e8f277b128d0 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -1,10 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use std::path::Path; use std::time::Duration; -use tokio::time::timeout; +use tokio::time::{sleep, timeout}; use zeromq::prelude::{Socket, SocketRecv, SocketSend}; use zeromq::util::PeerIdentity; use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage}; @@ -32,7 +31,7 @@ pub struct MockEngineConfig { /// Engine-ready payload reported after INIT, including max model length, /// KV block count, and dtype. pub ready_response: EngineCoreReadyResponse, - /// Maximum time to wait for IPC endpoints to appear before connecting. + /// Maximum time to wait for endpoints to accept connections. pub connect_timeout: Duration, } @@ -121,22 +120,33 @@ fn peer_identity(engine_id: impl Into) -> Result { }) } -/// Wait for an IPC endpoint path to appear before attempting to connect. -async fn wait_for_ipc_endpoint(endpoint: &str, connect_timeout: Duration) -> Result<()> { - let Some(socket_path) = endpoint.strip_prefix("ipc://") else { - return Ok(()); - }; - - timeout(connect_timeout, async { - while !Path::new(socket_path).exists() { - tokio::time::sleep(Duration::from_millis(20)).await; - } - }) - .await - .map_err(|_| Error::HandshakeTimeout { - stage: "mock engine IPC endpoint", - timeout: connect_timeout, - }) +/// Wait for an endpoint to accept connections before attempting the ZMQ connect. +async fn wait_for_endpoint(endpoint: &str, connect_timeout: Duration) -> Result<()> { + if let Some(socket_path) = endpoint.strip_prefix("ipc://") { + timeout(connect_timeout, async { + while tokio::net::UnixStream::connect(socket_path).await.is_err() { + sleep(Duration::from_millis(20)).await; + } + }) + .await + .map_err(|_| Error::HandshakeTimeout { + stage: "mock engine IPC endpoint", + timeout: connect_timeout, + }) + } else if let Some(address) = endpoint.strip_prefix("tcp://") { + timeout(connect_timeout, async { + while tokio::net::TcpStream::connect(address).await.is_err() { + sleep(Duration::from_millis(20)).await; + } + }) + .await + .map_err(|_| Error::HandshakeTimeout { + stage: "mock engine TCP endpoint", + timeout: connect_timeout, + }) + } else { + Ok(()) + } } /// Encode the engine-ready response sent on input socket registration. @@ -151,7 +161,7 @@ pub async fn connect_to_frontend( config: MockEngineConfig, ) -> Result { let engine_handshake = engine_handshake.as_ref(); - wait_for_ipc_endpoint(engine_handshake, config.connect_timeout).await?; + wait_for_endpoint(engine_handshake, config.connect_timeout).await?; let peer_identity = peer_identity(engine_id)?; let mut options = SocketOptions::default(); @@ -192,8 +202,8 @@ pub async fn connect_to_frontend( for (input_address, output_address) in init.addresses.inputs.iter().zip(init.addresses.outputs.iter()) { - wait_for_ipc_endpoint(input_address, config.connect_timeout).await?; - wait_for_ipc_endpoint(output_address, config.connect_timeout).await?; + wait_for_endpoint(input_address, config.connect_timeout).await?; + wait_for_endpoint(output_address, config.connect_timeout).await?; let mut input_options = SocketOptions::default(); input_options.peer_identity(peer_identity.clone()); @@ -260,8 +270,8 @@ pub async fn connect_to_bootstrapped_frontend( ) -> Result<(DealerSocket, PushSocket)> { let input_address = input_address.as_ref(); let output_address = output_address.as_ref(); - wait_for_ipc_endpoint(input_address, config.connect_timeout).await?; - wait_for_ipc_endpoint(output_address, config.connect_timeout).await?; + wait_for_endpoint(input_address, config.connect_timeout).await?; + wait_for_endpoint(output_address, config.connect_timeout).await?; let peer_identity = peer_identity(engine_id)?; let mut input_options = SocketOptions::default(); diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index 6765faeab75f..979da139e869 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -90,8 +90,8 @@ async fn shutdown_mock( shutdown: CancellationToken, task: tokio::task::JoinHandle>, ) { - client.shutdown().await.expect("client shutdown"); shutdown.cancel(); + client.shutdown().await.expect("client shutdown"); task.await.expect("mock join").expect("mock run"); } diff --git a/rust/src/parser/src/tool/json/mod.rs b/rust/src/parser/src/tool/json/mod.rs index b707bff6fd19..d545ece23e43 100644 --- a/rust/src/parser/src/tool/json/mod.rs +++ b/rust/src/parser/src/tool/json/mod.rs @@ -31,24 +31,24 @@ use super::utils::{ }; use super::{Result, ToolCallDelta, ToolParserOutput}; -type JsonToolInput<'i> = Partial<&'i str>; +pub(crate) type JsonToolInput<'i> = Partial<&'i str>; #[derive(Debug, Clone, Copy)] -struct JsonToolCallConfig { - parser_name: &'static str, - start_marker: &'static str, - end_marker: &'static str, - marker_whitespace: JsonToolCallWhitespace, - delimiter: Option<&'static str>, - name_key: &'static str, +pub(crate) struct JsonToolCallConfig { + pub parser_name: &'static str, + pub start_marker: &'static str, + pub end_marker: &'static str, + pub marker_whitespace: JsonToolCallWhitespace, + pub delimiter: Option<&'static str>, + pub name_key: &'static str, /// Candidate JSON keys naming the arguments payload, tried in order. /// Most parsers use a single key like `["arguments"]`, but some accept /// multiple (e.g. InternLM2 accepts `parameters` or `arguments`). - arguments_key: &'static [&'static str], + pub arguments_key: &'static [&'static str], } #[derive(Debug, Clone, Copy)] -enum JsonToolCallWhitespace { +pub(crate) enum JsonToolCallWhitespace { Optional, Exact(&'static str), } @@ -61,7 +61,7 @@ enum JsonToolCallMode { } #[derive(Debug, Clone, PartialEq, Eq)] -enum JsonToolCallEvent { +pub(crate) enum JsonToolCallEvent { Text { len: usize }, ToolCallStart, ToolCallHeader { function_name: String }, @@ -220,7 +220,7 @@ fn tool_call_start_event( /// Parse a marker-wrapped JSON tool-call header before the raw arguments /// payload. -fn tool_call_header_event( +pub(crate) fn tool_call_header_event( input: &mut JsonToolInput<'_>, config: JsonToolCallConfig, ) -> ModalResult { diff --git a/rust/src/parser/src/tool/mod.rs b/rust/src/parser/src/tool/mod.rs index d0e88e19362d..216d3dff4cf0 100644 --- a/rust/src/parser/src/tool/mod.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -9,12 +9,13 @@ mod deepseek_dsml; pub(crate) mod deepseek_json; mod glm_xml; mod hy_v3; -mod json; +pub(crate) mod json; mod kimi_k2; mod minimax_m2; mod minimax_m3; mod parameters; mod qwen_coder; +mod seed_oss; #[cfg(any(test, feature = "test-util"))] pub mod test_utils; use std::collections::{BTreeMap, btree_map}; @@ -32,6 +33,7 @@ pub use kimi_k2::KimiK2ToolParser; pub use minimax_m2::MinimaxM2ToolParser; pub use minimax_m3::MinimaxM3ToolParser; pub use qwen_coder::Qwen3CoderToolParser; +pub use seed_oss::SeedOssToolParser; use serde::{Deserialize, Serialize}; use serde_json::Value; pub use xgrammar_structural_tag::Model as StructuralTagModel; diff --git a/rust/src/parser/src/tool/qwen_coder.rs b/rust/src/parser/src/tool/qwen_coder.rs index f864984192fc..361de524e8ba 100644 --- a/rust/src/parser/src/tool/qwen_coder.rs +++ b/rust/src/parser/src/tool/qwen_coder.rs @@ -19,6 +19,28 @@ const FUNCTION_END: &str = ""; const PARAMETER_START: &str = "` / `` tags are always +/// byte-identical. Seed-OSS, for example, wraps the same body in +/// `` / ``. +#[derive(Debug, Clone, Copy)] +pub(crate) struct QwenCoderConfig { + /// Human-readable parser name used in error messages. + pub(crate) parser_name: &'static str, + /// Marker that opens a tool-call block. + pub(crate) tool_call_start: &'static str, + /// Marker that closes a tool-call block. + pub(crate) tool_call_end: &'static str, +} + +const QWEN_CODER_CONFIG: QwenCoderConfig = QwenCoderConfig { + parser_name: "Qwen Coder", + tool_call_start: TOOL_CALL_START, + tool_call_end: TOOL_CALL_END, +}; + type QwenCoderInput<'i> = Partial<&'i str>; #[derive(Debug, Clone, PartialEq, Eq)] @@ -60,16 +82,24 @@ pub struct Qwen3CoderToolParser { mode: QwenCoderMode, emitted_tool_count: usize, tool_parameters: ToolSchemas, + config: QwenCoderConfig, } impl Qwen3CoderToolParser { /// Create a Qwen Coder tool parser. fn new(tools: &[Tool]) -> Self { + Self::with_config(tools, QWEN_CODER_CONFIG) + } + + /// Create a parser for a model that reuses the Qwen Coder grammar with a + /// different tool-call wrapper (e.g. Seed-OSS). + pub(crate) fn with_config(tools: &[Tool], config: QwenCoderConfig) -> Self { Self { buffer: String::new(), mode: QwenCoderMode::Text, emitted_tool_count: 0, tool_parameters: ToolSchemas::from_tools(tools), + config, } } @@ -122,9 +152,10 @@ impl ToolParser for Qwen3CoderToolParser { fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { self.buffer.push_str(chunk); + let config = self.config; while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { - parse_next_qwen_coder_event(input, &mut self.mode) + parse_next_qwen_coder_event(input, &mut self.mode, config) })? { self.apply_event(event, output)?; self.buffer.drain(..consumed_len); @@ -137,9 +168,12 @@ impl ToolParser for Qwen3CoderToolParser { let mut output = ToolParserOutput::default(); if !self.buffer.is_empty() { if matches!(self.mode, QwenCoderMode::ToolCall { .. }) - || self.buffer.starts_with(TOOL_CALL_START) + || self.buffer.starts_with(self.config.tool_call_start) { - return Err(parsing_failed!("incomplete Qwen Coder tool call")); + return Err(parsing_failed!( + "incomplete {} tool call", + self.config.parser_name + )); } output.push_text(&self.buffer); } @@ -156,37 +190,54 @@ impl ToolParser for Qwen3CoderToolParser { fn parse_next_qwen_coder_event( input: &mut QwenCoderInput<'_>, mode: &mut QwenCoderMode, + config: QwenCoderConfig, ) -> ModalResult { match mode { - QwenCoderMode::Text => parse_text_event(input), - QwenCoderMode::ToolCall { end_marker_scan } => tool_call_event(input, end_marker_scan), + QwenCoderMode::Text => parse_text_event(input, config), + QwenCoderMode::ToolCall { end_marker_scan } => { + tool_call_event(input, end_marker_scan, config.tool_call_end) + } } } /// Parse a text-mode Qwen Coder event. -fn parse_text_event(input: &mut QwenCoderInput<'_>) -> ModalResult { - alt((tool_call_start_event, safe_text_event)).parse_next(input) +fn parse_text_event( + input: &mut QwenCoderInput<'_>, + config: QwenCoderConfig, +) -> ModalResult { + alt(( + |input: &mut QwenCoderInput<'_>| tool_call_start_event(input, config.tool_call_start), + |input: &mut QwenCoderInput<'_>| safe_text_event(input, config.tool_call_start), + )) + .parse_next(input) } /// Parse a Qwen Coder tool-call start marker. -fn tool_call_start_event(input: &mut QwenCoderInput<'_>) -> ModalResult { - literal(TOOL_CALL_START).value(QwenCoderEvent::ToolCallStart).parse_next(input) +fn tool_call_start_event( + input: &mut QwenCoderInput<'_>, + tool_call_start: &'static str, +) -> ModalResult { + literal(tool_call_start).value(QwenCoderEvent::ToolCallStart).parse_next(input) } /// Parse a safe text run before the next Qwen Coder marker. -fn safe_text_event(input: &mut QwenCoderInput<'_>) -> ModalResult { - safe_text_len(input, TOOL_CALL_START).map(|len| QwenCoderEvent::Text { len }) +fn safe_text_event( + input: &mut QwenCoderInput<'_>, + tool_call_start: &'static str, +) -> ModalResult { + safe_text_len(input, tool_call_start).map(|len| QwenCoderEvent::Text { len }) } /// Parse a complete Qwen Coder tool call. fn tool_call_event( input: &mut QwenCoderInput<'_>, end_marker_scan: &mut MarkerScanState, + tool_call_end: &'static str, ) -> ModalResult { let (body,) = seq!( _: ws0, - take_until_marker(TOOL_CALL_END, end_marker_scan), - _: literal(TOOL_CALL_END), + take_until_marker(tool_call_end, end_marker_scan), + _: literal(tool_call_end), ) .parse_next(input)?; diff --git a/rust/src/parser/src/tool/seed_oss.rs b/rust/src/parser/src/tool/seed_oss.rs new file mode 100644 index 000000000000..dae7c5b2c530 --- /dev/null +++ b/rust/src/parser/src/tool/seed_oss.rs @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use super::qwen_coder::{Qwen3CoderToolParser, QwenCoderConfig}; +use crate::tool::{Result, Tool, ToolParser, ToolParserOutput}; + +const SEED_OSS_CONFIG: QwenCoderConfig = QwenCoderConfig { + parser_name: "Seed-OSS", + tool_call_start: "", + tool_call_end: "", +}; + +/// Tool parser for Seed-OSS XML-style tool calls. +/// +/// Example tool call content: +/// +/// ```text +/// +/// +/// 杭州 +/// +/// +/// ``` +/// +/// Seed-OSS shares the Qwen3 Coder grammar exactly; only the two tool-call +/// wrapper tokens differ (`` / ``). The inner +/// `` / `` grammar and schema-driven argument +/// conversion are byte-identical, so this delegates to a +/// [`Qwen3CoderToolParser`] configured with the Seed-OSS markers. This mirrors +/// Python `SeedOssParser(Qwen3Parser)` in `vllm/parser/seed_oss.py`. +/// +/// Structured-output tags are intentionally unsupported: the trait default +/// `structural_tag_model() -> None` matches Python `SeedOssEngineToolParser`, +/// which sets `structural_tag_model = None` (xgrammar has no Seed-OSS model). +/// +/// The `` wrapper tokens are added-vocabulary tokens +/// (`special = false`), not tokenizer special tokens, so they survive decoding +/// under the production default `skip_special_tokens = true`; the trait default +/// `preserve_special_tokens() == false` is therefore correct, matching the +/// sibling `SeedOssReasoningParser`, which relies on the same for ``. +pub struct SeedOssToolParser { + inner: Qwen3CoderToolParser, +} + +impl SeedOssToolParser { + /// Create a Seed-OSS tool parser. + fn new(tools: &[Tool]) -> Self { + Self { + inner: Qwen3CoderToolParser::with_config(tools, SEED_OSS_CONFIG), + } + } +} + +impl ToolParser for SeedOssToolParser { + fn create(tools: &[Tool]) -> Result> + where + Self: Sized + 'static, + { + Ok(Box::new(Self::new(tools))) + } + + fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> { + self.inner.parse_into(chunk, output) + } + + fn finish(&mut self) -> Result { + self.inner.finish() + } + + fn reset(&mut self) -> String { + self.inner.reset() + } +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + use thiserror_ext::AsReport; + + use super::SeedOssToolParser; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + + fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { + let params = params + .iter() + .map(|(name, value)| format!("{value}")) + .collect::>() + .join("\n"); + format!( + "\n\n{params}\n\n" + ) + } + + #[test] + fn seed_oss_parse_complete_without_tool_call_keeps_text() { + let mut parser = SeedOssToolParser::new(&test_tools()); + let output = parser.parse_complete("Hello, world!").unwrap(); + + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); + } + + #[test] + fn seed_oss_does_not_treat_plain_tool_call_marker_as_start() { + // The Qwen Coder wrapper `` must NOT trigger the Seed-OSS + // parser; only `` does. + let mut parser = SeedOssToolParser::new(&test_tools()); + let output = parser + .parse_complete("\n\n\n") + .unwrap(); + + assert_eq!( + output.normal_text(), + "\n\n\n" + ); + assert!(output.calls().is_empty()); + } + + #[test] + fn seed_oss_parse_complete_extracts_single_tool_call() { + let mut parser = SeedOssToolParser::new(&test_tools()); + let output = parser + .parse_complete(&build_tool_call( + "get_weather", + &[("location", "SF"), ("date", "2026-04-29")], + )) + .unwrap(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "location": "SF", "date": "2026-04-29" }) + ); + } + + #[test] + fn seed_oss_parse_complete_preserves_prefix_text() { + let mut parser = SeedOssToolParser::new(&test_tools()); + let input = format!( + "Thinking... {}", + build_tool_call("get_weather", &[("location", "NYC")]) + ); + let output = parser.parse_complete(&input).unwrap(); + + assert_eq!(output.normal_text(), "Thinking... "); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn seed_oss_parse_complete_converts_schema_types() { + let mut parser = SeedOssToolParser::new(&test_tools()); + let output = parser + .parse_complete(&build_tool_call( + "convert", + &[ + ("whole", "5.0"), + ("flag", "true"), + ("payload", r#"{"nested":true}"#), + ("items", "[1,2]"), + ("empty", "42"), + ], + )) + .unwrap(); + + assert_eq!(output.calls().len(), 1); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ + "whole": 5.0, + "flag": true, + "payload": { "nested": true }, + "items": [1, 2], + "empty": "42", + }) + ); + } + + #[test] + fn seed_oss_streaming_extracts_multiple_tool_calls_in_order() { + let text = format!( + "{}\n{}", + build_tool_call("get_weather", &[("location", "SF")]), + build_tool_call("get_weather", &[("location", "NYC")]) + ); + let chunks = split_by_chars(&text, 7); + let mut parser = SeedOssToolParser::new(&test_tools()); + + let output = collect_stream(&mut parser, &chunks); + + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[1].tool_index, 1); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "location": "SF" }) + ); + assert_eq!( + serde_json::from_str::(&output.calls()[1].arguments).unwrap(), + json!({ "location": "NYC" }) + ); + } + + #[test] + fn seed_oss_streaming_handles_end_marker_split_across_chunks() { + let mut parser = SeedOssToolParser::new(&test_tools()); + let mut output = ToolParserOutput::default(); + output.append( + parser + .parse_chunk( + "\n\ + \n\ + SF\n\ + \n\ + ").unwrap()); + output.append(parser.finish().unwrap()); + let output = output.coalesce(); + + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), + json!({ "location": "SF" }) + ); + } + + #[test] + fn seed_oss_finish_fails_incomplete_tool_call() { + let mut parser = SeedOssToolParser::new(&test_tools()); + parser + .parse_chunk( + "\n\nSF", + ) + .unwrap(); + + let error = parser.finish().unwrap_err(); + + assert_eq!( + error.to_report_string(), + "tool parser parsing failed: incomplete Seed-OSS tool call" + ); + } +} diff --git a/rust/src/parser/src/unified/gemma4.rs b/rust/src/parser/src/unified/gemma4.rs index bc5925189e77..8577a3fc286d 100644 --- a/rust/src/parser/src/unified/gemma4.rs +++ b/rust/src/parser/src/unified/gemma4.rs @@ -10,7 +10,7 @@ use winnow::prelude::*; use winnow::stream::{Partial, Stream}; use winnow::token::{literal, take_till, take_until}; -use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; use crate::reasoning::last_reasoning_boundary; use crate::tool::{Tool, ToolCallDelta}; use crate::unified::parsing_failed; @@ -79,17 +79,8 @@ pub struct Gemma4UnifiedParser { impl Gemma4UnifiedParser { /// Create a Gemma4 parser. pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { - let channel_start_token_id = tokenizer.token_to_id(CHANNEL_START).ok_or_else(|| { - UnifiedParserError::MissingToken { - token: CHANNEL_START.to_string(), - } - })?; - let channel_end_token_id = - tokenizer - .token_to_id(CHANNEL_END) - .ok_or_else(|| UnifiedParserError::MissingToken { - token: CHANNEL_END.to_string(), - })?; + let channel_start_token_id = token_id(tokenizer.as_ref(), CHANNEL_START)?; + let channel_end_token_id = token_id(tokenizer.as_ref(), CHANNEL_END)?; Ok(Self { buffer: String::new(), @@ -524,10 +515,10 @@ mod tests { use super::{ CHANNEL_END, CHANNEL_START, Gemma4UnifiedParser, ToolCallDelta, UnifiedParser, - UnifiedParserError, UnifiedParserOutput, gemma4_array_content, parse_gemma4_args, + UnifiedParserOutput, gemma4_array_content, parse_gemma4_args, }; use crate::tool::Tool; - use crate::unified::{UnifiedParserEvent, parsing_failed}; + use crate::unified::{UnifiedParserError, UnifiedParserEvent, parsing_failed}; const CHANNEL_START_ID: u32 = 256; const CHANNEL_END_ID: u32 = 257; diff --git a/rust/src/parser/src/unified/inkling.rs b/rust/src/parser/src/unified/inkling.rs new file mode 100644 index 000000000000..eb78321d1269 --- /dev/null +++ b/rust/src/parser/src/unified/inkling.rs @@ -0,0 +1,768 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use winnow::ascii::multispace0 as ws0; +use winnow::combinator::{alt, seq}; +use winnow::error::ModalResult; +use winnow::prelude::*; +use winnow::stream::Partial; +use winnow::token::literal; + +use vllm_tokenizer::DynTokenizer; + +use super::{Result, UnifiedParser, UnifiedParserOutput, token_id}; +use crate::tool::json::{ + JsonToolCallConfig, JsonToolCallEvent, JsonToolCallWhitespace, JsonToolInput, + tool_call_header_event, +}; +use crate::tool::{Tool, ToolCallDelta}; +use crate::unified::parsing_failed; +use crate::utils::{ + JsonObjectScanState, parse_buffered_event, safe_text_len_mul, take_json_object, +}; + +const CONTENT_TEXT: &str = "<|content_text|>"; +const CONTENT_THINKING: &str = "<|content_thinking|>"; +const CONTENT_INVOKE_TOOL_JSON: &str = "<|content_invoke_tool_json|>"; +const CONTENT_INVOKE_TOOL_TEXT: &str = "<|content_invoke_tool_text|>"; +const CONTENT_MODEL_END_SAMPLING: &str = "<|content_model_end_sampling|>"; +const CONTENT_TOOL_ERROR: &str = "<|content_tool_error|>"; +const MESSAGE_MODEL: &str = "<|message_model|>"; +const END_MESSAGE: &str = "<|end_message|>"; + +const IDLE_MARKERS: &[&str] = &[ + MESSAGE_MODEL, + CONTENT_TEXT, + CONTENT_THINKING, + CONTENT_INVOKE_TOOL_JSON, + CONTENT_INVOKE_TOOL_TEXT, + CONTENT_MODEL_END_SAMPLING, + CONTENT_TOOL_ERROR, + END_MESSAGE, +]; +const BLOCK_END_MARKERS: &[&str] = &[END_MESSAGE, CONTENT_MODEL_END_SAMPLING]; + +const INKLING_TOOL_CONFIG: JsonToolCallConfig = JsonToolCallConfig { + parser_name: "Inkling", + start_marker: CONTENT_INVOKE_TOOL_JSON, + end_marker: END_MESSAGE, + marker_whitespace: JsonToolCallWhitespace::Optional, + delimiter: None, + name_key: "name", + arguments_key: &["args"], +}; + +type InklingInput<'i> = Partial<&'i str>; + +#[derive(Debug, Clone, PartialEq, Eq)] +enum InklingEvent { + Text { len: usize }, + Reasoning { len: usize }, + TextStart, + ReasoningStart, + MessageStart, + Header, + ToolJsonStart, + ToolJsonHeader { name: String }, + ToolJsonArgs { len: usize, complete: bool }, + BlockEnd, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +enum InklingMode { + #[default] + Idle, + MessageHeader, + Text, + Reasoning, + ToolJsonHeader, + ToolJsonArgs { + json_scan: JsonObjectScanState, + }, + ToolJsonClose, +} + +/// Unified parser for Inkling typed content blocks. +pub struct InklingUnifiedParser { + buffer: String, + mode: InklingMode, + emitted_tool_count: usize, + active_tool_index: Option, + tokenizer: DynTokenizer, + message_model_token_id: u32, + content_text_token_id: u32, + content_thinking_token_id: u32, +} + +impl InklingUnifiedParser { + /// Create a Inkling parser. + pub fn new(_tools: &[Tool], tokenizer: DynTokenizer) -> Result { + let message_model_token_id = token_id(tokenizer.as_ref(), MESSAGE_MODEL)?; + let content_text_token_id = token_id(tokenizer.as_ref(), CONTENT_TEXT)?; + let content_thinking_token_id = token_id(tokenizer.as_ref(), CONTENT_THINKING)?; + + Ok(Self { + buffer: String::new(), + mode: InklingMode::Idle, + emitted_tool_count: 0, + active_tool_index: None, + tokenizer, + message_model_token_id, + content_text_token_id, + content_thinking_token_id, + }) + } + + fn initialize_mode(&mut self, prompt_token_ids: &[u32]) { + self.mode = InklingMode::Idle; + for token_id in prompt_token_ids.iter().rev().copied() { + if token_id == self.message_model_token_id { + self.mode = InklingMode::MessageHeader; + return; + } + if token_id == self.content_thinking_token_id { + self.mode = InklingMode::Reasoning; + return; + } + if token_id == self.content_text_token_id { + self.mode = InklingMode::Text; + return; + } + if self.tokenizer.is_special_id(token_id) { + return; + } + } + } + + fn apply_event(&mut self, event: InklingEvent, output: &mut UnifiedParserOutput) -> Result<()> { + match event { + InklingEvent::Text { len } => output.push_text(self.buffer[..len].to_string()), + InklingEvent::Reasoning { len } => { + output.push_reasoning(self.buffer[..len].to_string()); + } + InklingEvent::MessageStart => self.mode = InklingMode::MessageHeader, + InklingEvent::Header => {} + InklingEvent::TextStart => self.mode = InklingMode::Text, + InklingEvent::ReasoningStart => self.mode = InklingMode::Reasoning, + InklingEvent::ToolJsonStart => self.mode = InklingMode::ToolJsonHeader, + InklingEvent::ToolJsonHeader { name } => { + let tool_index = self.emitted_tool_count; + self.emitted_tool_count += 1; + self.active_tool_index = Some(tool_index); + self.mode = InklingMode::ToolJsonArgs { + json_scan: JsonObjectScanState::default(), + }; + output.push_call(ToolCallDelta { + tool_index, + name: Some(name), + arguments: String::new(), + }); + } + InklingEvent::ToolJsonArgs { len, complete } => { + let Some(tool_index) = self.active_tool_index else { + return Err(parsing_failed!( + "Inkling arguments without an active tool call" + )); + }; + output.push_call(ToolCallDelta { + tool_index, + name: None, + arguments: self.buffer[..len].to_string(), + }); + if complete { + self.mode = InklingMode::ToolJsonClose; + } + } + InklingEvent::BlockEnd => { + self.mode = InklingMode::Idle; + self.active_tool_index = None; + } + } + Ok(()) + } + + fn reset(&mut self) -> String { + self.mode = InklingMode::Idle; + self.active_tool_index = None; + self.emitted_tool_count = 0; + std::mem::take(&mut self.buffer) + } +} + +impl UnifiedParser for InklingUnifiedParser { + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Self::new(tools, tokenizer).map(|parser| Box::new(parser) as Box) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.buffer.clear(); + self.emitted_tool_count = 0; + self.active_tool_index = None; + self.initialize_mode(prompt_token_ids); + Ok(()) + } + + fn preserve_special_tokens(&self) -> bool { + true + } + + fn parse_into(&mut self, chunk: &str, output: &mut UnifiedParserOutput) -> Result<()> { + self.buffer.push_str(chunk); + + while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| { + parse_next_inkling_event(input, &mut self.mode) + })? { + self.apply_event(event, output)?; + self.buffer.drain(..consumed_len); + } + + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); + + match &self.mode { + InklingMode::Idle | InklingMode::Text => { + output.push_text(std::mem::take(&mut self.buffer)) + } + InklingMode::MessageHeader => self.buffer.clear(), + InklingMode::Reasoning => output.push_reasoning(std::mem::take(&mut self.buffer)), + InklingMode::ToolJsonHeader + | InklingMode::ToolJsonArgs { .. } + | InklingMode::ToolJsonClose => { + return Err(parsing_failed!("incomplete Inkling tool call")); + } + } + + let _ = self.reset(); + Ok(output) + } + + fn reset(&mut self) -> String { + InklingUnifiedParser::reset(self) + } +} + +/// Parse one Inkling event from buffered streaming input. +fn parse_next_inkling_event( + input: &mut InklingInput<'_>, + mode: &mut InklingMode, +) -> ModalResult { + match mode { + InklingMode::Idle => parse_idle_event(input), + InklingMode::MessageHeader => parse_message_header_event(input), + InklingMode::Text => parse_text_event(input), + InklingMode::Reasoning => parse_reasoning_event(input), + InklingMode::ToolJsonHeader => parse_tool_json_header_event(input), + InklingMode::ToolJsonArgs { json_scan } => parse_tool_json_args_event(input, json_scan), + InklingMode::ToolJsonClose => parse_tool_json_close_event(input), + } +} + +/// Parse an event while waiting for a Inkling content kind. +fn parse_idle_event(input: &mut InklingInput<'_>) -> ModalResult { + alt(( + message_start_event, + reasoning_start_event, + text_start_event, + tool_json_start_event, + raw_text_start_event, + block_end_event, + safe_idle_text_event, + )) + .parse_next(input) +} + +/// Parse a Inkling model-authored message start marker. +fn message_start_event(input: &mut InklingInput<'_>) -> ModalResult { + literal(MESSAGE_MODEL).value(InklingEvent::MessageStart).parse_next(input) +} + +/// Parse an event while waiting for an Inkling message content kind. +fn parse_message_header_event(input: &mut InklingInput<'_>) -> ModalResult { + alt(( + reasoning_start_event, + text_start_event, + tool_json_start_event, + raw_text_start_event, + block_end_event, + safe_header_event, + )) + .parse_next(input) +} + +/// Parse an event inside a Inkling text block. +fn parse_text_event(input: &mut InklingInput<'_>) -> ModalResult { + alt((block_end_event, safe_text_event)).parse_next(input) +} + +/// Parse an event inside a Inkling reasoning block. +fn parse_reasoning_event(input: &mut InklingInput<'_>) -> ModalResult { + alt((block_end_event, safe_reasoning_event)).parse_next(input) +} + +/// Parse a Inkling text start marker. +fn text_start_event(input: &mut InklingInput<'_>) -> ModalResult { + literal(CONTENT_TEXT).value(InklingEvent::TextStart).parse_next(input) +} + +/// Parse a Inkling reasoning start marker. +fn reasoning_start_event(input: &mut InklingInput<'_>) -> ModalResult { + literal(CONTENT_THINKING).value(InklingEvent::ReasoningStart).parse_next(input) +} + +/// Parse a Inkling JSON tool-call start marker. +fn tool_json_start_event(input: &mut InklingInput<'_>) -> ModalResult { + literal(CONTENT_INVOKE_TOOL_JSON) + .value(InklingEvent::ToolJsonStart) + .parse_next(input) +} + +/// Parse a Inkling content kind treated as visible text. +fn raw_text_start_event(input: &mut InklingInput<'_>) -> ModalResult { + alt(( + literal(CONTENT_INVOKE_TOOL_TEXT), + literal(CONTENT_TOOL_ERROR), + )) + .value(InklingEvent::TextStart) + .parse_next(input) +} + +/// Parse a Inkling block end marker. +fn block_end_event(input: &mut InklingInput<'_>) -> ModalResult { + alt((literal(END_MESSAGE), literal(CONTENT_MODEL_END_SAMPLING))) + .value(InklingEvent::BlockEnd) + .parse_next(input) +} + +/// Parse safe text while waiting for the next Inkling marker. +fn safe_idle_text_event(input: &mut InklingInput<'_>) -> ModalResult { + safe_text_len_mul(input, IDLE_MARKERS).map(|len| InklingEvent::Text { len }) +} + +/// Parse safe header text before the next Inkling marker. +fn safe_header_event(input: &mut InklingInput<'_>) -> ModalResult { + safe_text_len_mul(input, IDLE_MARKERS).map(|_| InklingEvent::Header) +} + +/// Parse safe text before the end of a Inkling text block. +fn safe_text_event(input: &mut InklingInput<'_>) -> ModalResult { + safe_text_len_mul(input, BLOCK_END_MARKERS).map(|len| InklingEvent::Text { len }) +} + +/// Parse safe reasoning before the end of a Inkling reasoning block. +fn safe_reasoning_event(input: &mut InklingInput<'_>) -> ModalResult { + safe_text_len_mul(input, BLOCK_END_MARKERS).map(|len| InklingEvent::Reasoning { len }) +} + +/// Parse a Inkling JSON tool-call header. +fn parse_tool_json_header_event(input: &mut InklingInput<'_>) -> ModalResult { + match tool_call_header_event(input, INKLING_TOOL_CONFIG)? { + JsonToolCallEvent::ToolCallHeader { function_name } => Ok(InklingEvent::ToolJsonHeader { + name: function_name, + }), + _ => unreachable!("tool_call_header_event only emits ToolCallHeader"), + } +} + +/// Parse raw Inkling JSON tool-call argument bytes. +fn parse_tool_json_args_event( + input: &mut JsonToolInput<'_>, + json_scan: &mut JsonObjectScanState, +) -> ModalResult { + let len = take_json_object(input, json_scan)?; + Ok(InklingEvent::ToolJsonArgs { + len, + complete: json_scan.complete(), + }) +} + +/// Parse the close of a Inkling JSON tool-call block. +fn parse_tool_json_close_event(input: &mut InklingInput<'_>) -> ModalResult { + seq!( + _: ws0, + _: literal("}"), + _: ws0, + _: literal(END_MESSAGE), + ) + .value(InklingEvent::BlockEnd) + .parse_next(input) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::{CONTENT_TEXT, CONTENT_THINKING, InklingUnifiedParser, MESSAGE_MODEL}; + use crate::tool::Tool; + use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; + use thiserror_ext::AsReport; + use vllm_tokenizer::Tokenizer; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + MESSAGE_MODEL => Some(200001), + "<|content_text|>" => Some(200004), + "<|content_model_end_sampling|>" => Some(200006), + "<|content_thinking|>" => Some(200008), + "<|end_message|>" => Some(200010), + "<|content_tool_error|>" => Some(200022), + "<|content_invoke_tool_json|>" => Some(200049), + "<|content_invoke_tool_text|>" => Some(200057), + _ => None, + } + } + + fn id_to_token(&self, id: u32) -> Option { + let token = match id { + 200001 => MESSAGE_MODEL, + 200004 => "<|content_text|>", + 200006 => "<|content_model_end_sampling|>", + 200008 => "<|content_thinking|>", + 200010 => "<|end_message|>", + 200022 => "<|content_tool_error|>", + 200049 => "<|content_invoke_tool_json|>", + 200057 => "<|content_invoke_tool_text|>", + _ => return None, + }; + Some(token.to_string()) + } + + fn is_special_id(&self, token_id: u32) -> bool { + (199999..=200057).contains(&token_id) + } + } + + trait UnifiedParserTestExt { + fn parse_chunk(&mut self, chunk: &str) -> super::Result; + fn parse_complete(&mut self, text: &str) -> super::Result; + } + + impl UnifiedParserTestExt for T { + fn parse_chunk(&mut self, chunk: &str) -> super::Result { + let mut output = UnifiedParserOutput::default(); + self.parse_into(chunk, &mut output)?; + Ok(output) + } + + fn parse_complete(&mut self, text: &str) -> super::Result { + let mut output = self.parse_chunk(text)?; + output.append(self.finish()?); + Ok(output) + } + } + + trait UnifiedParserOutputTestExt { + fn normal_text(&self) -> String; + fn reasoning_text(&self) -> String; + fn calls(&self) -> Vec; + } + + impl UnifiedParserOutputTestExt for UnifiedParserOutput { + fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn reasoning_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::Reasoning(text) => Some(text.as_str()), + _ => None, + }) + .collect() + } + + fn calls(&self) -> Vec { + self.events + .iter() + .filter_map(|event| match event { + UnifiedParserEvent::ToolCall(call) => Some(call.clone()), + _ => None, + }) + .collect() + } + } + + fn test_tools() -> Vec { + vec![Tool { + name: "get_weather".to_string(), + description: Some("Get weather".to_string()), + parameters: serde_json::json!({ + "type": "object", + "properties": { + "city": {"type": "string"} + }, + }), + strict: None, + }] + } + + fn test_parser() -> InklingUnifiedParser { + InklingUnifiedParser::new(&test_tools(), Arc::new(FakeTokenizer)).unwrap() + } + + fn collect_stream(chunks: &[&str]) -> UnifiedParserOutput { + let mut parser = test_parser(); + let mut output = UnifiedParserOutput::default(); + for chunk in chunks { + output.append(parser.parse_chunk(chunk).unwrap()); + } + output.append(parser.finish().unwrap()); + output + } + + #[test] + fn inkling_streaming_emits_reasoning_then_text() { + let output = collect_stream(&[concat!( + "<|content_thinking|>reason<|end_message|>", + "<|message_model|><|content_text|>answer<|end_message|>", + "<|content_model_end_sampling|>" + )]); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + assert!(output.calls().is_empty()); + } + + #[test] + fn inkling_streaming_holds_split_markers() { + let output = collect_stream(&[ + "<|content_thin", + "king|>rea", + "son<|end_mes", + "sage|><|message_", + "model|><|content_text|>answer", + "<|end_message|>", + ]); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn inkling_streaming_accepts_content_blocks_without_message_start() { + let output = collect_stream(&[concat!( + "<|content_thinking|>reason<|end_message|>", + "<|content_text|>answer<|end_message|>", + )]); + + assert_eq!(output.reasoning_text(), "reason"); + assert_eq!(output.normal_text(), "answer"); + } + + #[test] + fn inkling_tool_json_streams_name_then_argument_deltas() { + let mut parser = test_parser(); + let chunks = [ + "<|message_model|><|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":", + "{\"city\":", + "\"SF\"", + "}}<|end_message|>", + ]; + + let mut output = UnifiedParserOutput::default(); + let mut observed_args = Vec::new(); + for chunk in chunks { + let next = parser.parse_chunk(chunk).unwrap(); + observed_args.extend( + next.calls() + .into_iter() + .filter(|call| call.name.is_none()) + .map(|call| call.arguments), + ); + output.append(next); + } + output.append(parser.finish().unwrap()); + + let calls = output.calls(); + assert_eq!(calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(observed_args, ["{\"city\":", "\"SF\"", "}"]); + assert_eq!( + calls.iter().map(|call| call.arguments.as_str()).collect::(), + "{\"city\":\"SF\"}" + ); + assert!(output.normal_text().is_empty()); + } + + #[test] + fn inkling_discards_tool_name_from_message_header() { + let output = collect_stream(&[concat!( + "<|message_model|>get_weather<|content_invoke_tool_json|>", + "{\"name\":\"get_weather\",\"args\":{}}<|end_message|>" + )]); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + } + + #[test] + fn inkling_streaming_handles_multiple_tool_blocks() { + let output = collect_stream(&[concat!( + "<|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":{\"city\":\"SF\"}}<|end_message|>", + "<|message_model|><|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":{\"city\":\"NYC\"}}<|end_message|>" + )]); + + let calls = output.calls(); + assert_eq!(calls.iter().filter(|call| call.name.is_some()).count(), 2); + assert_eq!(calls[0].tool_index, 0); + assert_eq!(calls[2].tool_index, 1); + } + + #[test] + fn inkling_initialize_open_text_prompt_starts_in_text() { + let mut parser = test_parser(); + parser.initialize(&[200004]).unwrap(); + + let output = parser.parse_complete("answer<|end_message|>").unwrap(); + + assert_eq!(output.normal_text(), "answer"); + assert!(output.reasoning_text().is_empty()); + } + + #[test] + fn inkling_initialize_open_reasoning_prompt_starts_in_reasoning() { + let mut parser = test_parser(); + parser.initialize(&[200008]).unwrap(); + + let output = parser.parse_complete("reason<|end_message|>").unwrap(); + + assert_eq!(output.reasoning_text(), "reason"); + assert!(output.normal_text().is_empty()); + } + + #[test] + fn inkling_initialize_model_opener_starts_in_message_header() { + let mut parser = test_parser(); + parser.initialize(&[200001]).unwrap(); + + let output = parser + .parse_complete(concat!( + "get_weather<|content_invoke_tool_json|>", + "{\"name\":\"get_weather\",\"args\":{}}<|end_message|>" + )) + .unwrap(); + + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + } + + #[test] + fn inkling_plain_text_falls_through_as_text() { + let output = collect_stream(&["plain ", "answer"]); + + assert_eq!(output.normal_text(), "plain answer"); + assert!(output.reasoning_text().is_empty()); + } + + #[test] + fn inkling_raw_tool_text_and_tool_error_are_visible_text() { + let output = collect_stream(&[concat!( + "<|content_invoke_tool_text|>search SF<|end_message|>", + "<|content_tool_error|>failed<|end_message|>" + )]); + + assert_eq!(output.normal_text(), "search SFfailed"); + assert!(output.calls().is_empty()); + } + + #[test] + fn inkling_finish_fails_incomplete_tool_call() { + let mut parser = test_parser(); + parser + .parse_chunk("<|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":{\"city\"") + .unwrap(); + + let error = parser.finish().unwrap_err(); + + assert!(error.as_report().to_string().contains("incomplete Inkling tool call")); + } + + #[test] + fn inkling_rejects_non_object_args() { + let mut parser = test_parser(); + let error = parser + .parse_chunk("<|content_invoke_tool_json|>{\"name\":\"get_weather\",\"args\":42") + .unwrap_err(); + + assert!(error.as_report().to_string().contains("JSON object argument")); + } + + #[test] + fn inkling_missing_token_fails_create() { + struct MissingTokenizer; + + impl Tokenizer for MissingTokenizer { + fn encode( + &self, + _text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(vec![]) + } + + fn decode( + &self, + _token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(String::new()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + MESSAGE_MODEL => Some(200001), + CONTENT_TEXT => Some(200004), + _ => None, + } + } + + fn id_to_token(&self, id: u32) -> Option { + match id { + 200001 => Some(MESSAGE_MODEL.to_string()), + 200004 => Some(CONTENT_TEXT.to_string()), + _ => None, + } + } + } + + let error = match InklingUnifiedParser::new(&[], Arc::new(MissingTokenizer)) { + Ok(_) => panic!("expected parser creation to fail"), + Err(error) => error, + }; + + assert!(error.as_report().to_string().contains(CONTENT_THINKING)); + } +} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 1fec8ac09ba9..d5ecb6492904 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -5,9 +5,11 @@ mod combined; mod gemma4; +mod inkling; pub use combined::CombinedParser; pub use gemma4::Gemma4UnifiedParser; +pub use inkling::InklingUnifiedParser; use thiserror::Error; use thiserror_ext::Macro; use vllm_tokenizer::DynTokenizer; @@ -208,3 +210,10 @@ pub enum UnifiedParserError { #[error(transparent)] Tool(#[from] ToolParserError), } + +/// Returns the ID for the given token, or an error if it's not found. +fn token_id(tokenizer: &dyn vllm_tokenizer::Tokenizer, token: &str) -> Result { + tokenizer.token_to_id(token).ok_or_else(|| UnifiedParserError::MissingToken { + token: token.to_string(), + }) +} diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index 41ef192b7363..b08c533a51bd 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -125,6 +125,7 @@ async fn collect_chat_completion( // Ignored: non-streaming responses are collected before usage is attached. include_continuous_usage: _, requested_logprobs, + output_top_logprobs, include_prompt_logprobs, include_reasoning, echo, @@ -173,6 +174,7 @@ async fn collect_chat_completion( logprobs.as_ref().ok_or_else(|| { server_error!("chat response requested logprobs but generation returned none") })?, + output_top_logprobs, return_tokens_as_token_ids, )?) } else { @@ -248,6 +250,7 @@ async fn chat_completion_chunk_stream( include_usage, include_continuous_usage, requested_logprobs, + output_top_logprobs, // Ignored: chat streaming prompt logprobs are rejected for Python parity. include_prompt_logprobs: _, include_reasoning, @@ -337,7 +340,13 @@ async fn chat_completion_chunk_stream( let openai_logprobs = if include_metadata { logprobs .as_ref() - .map(|lp| decoded_logprobs_to_openai_chat(lp, return_tokens_as_token_ids)) + .map(|lp| { + decoded_logprobs_to_openai_chat( + lp, + output_top_logprobs, + return_tokens_as_token_ids, + ) + }) .transpose()? } else { None diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index cfb4b1a7674a..5b2fa3c19ed7 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -42,6 +42,8 @@ pub(super) struct ResponseOptions { pub include_continuous_usage: bool, /// Whether the caller requested output logprobs on chat choices. pub requested_logprobs: bool, + /// Number of top logprobs to include for each output token. + pub output_top_logprobs: i32, /// Whether the caller requested top-level prompt logprobs. pub include_prompt_logprobs: bool, /// Whether to include reasoning content in OpenAI responses. @@ -172,6 +174,7 @@ pub(super) fn prepare_chat_request( include_usage, include_continuous_usage, requested_logprobs, + output_top_logprobs: top_logprobs, include_prompt_logprobs, include_reasoning, echo, diff --git a/rust/src/server/src/routes/openai/utils/logprobs.rs b/rust/src/server/src/routes/openai/utils/logprobs.rs index 19215755f57c..078c42b507fd 100644 --- a/rust/src/server/src/routes/openai/utils/logprobs.rs +++ b/rust/src/server/src/routes/openai/utils/logprobs.rs @@ -120,12 +120,13 @@ pub fn decoded_prompt_logprobs_to_maps( /// shape. pub fn decoded_logprobs_to_openai_chat( logprobs: &DecodedLogprobs, + top_logprobs: i32, return_tokens_as_token_ids: bool, ) -> Result { let content = logprobs .positions .iter() - .map(|pos| position_to_chat_logprobs_content(pos, return_tokens_as_token_ids)) + .map(|pos| position_to_chat_logprobs_content(pos, top_logprobs, return_tokens_as_token_ids)) .try_collect()?; Ok(ChatLogProbs { @@ -224,6 +225,7 @@ fn position_top_logprobs_map( fn position_to_chat_logprobs_content( position: &DecodedPositionLogprobs, + top_logprobs: i32, return_tokens_as_token_ids: bool, ) -> Result { let chosen = position.entries.first().ok_or_else(|| { @@ -235,9 +237,7 @@ fn position_to_chat_logprobs_content( token: token_str.clone(), logprob: clamp_logprob(chosen.logprob), bytes: Some(token_bytes(&token_str)), - top_logprobs: position - .entries - .iter() + top_logprobs: chat_top_logprob_entries(position, top_logprobs) .map(|entry| { let t = format_token(entry, return_tokens_as_token_ids); TopLogProb { @@ -250,6 +250,19 @@ fn position_to_chat_logprobs_content( }) } +fn chat_top_logprob_entries( + position: &DecodedPositionLogprobs, + top_logprobs: i32, +) -> impl Iterator { + let limit = if top_logprobs == -1 { + position.entries.len() + } else { + usize::try_from(top_logprobs).unwrap_or(0) + }; + + position.entries.iter().take(limit) +} + fn token_bytes(token: &str) -> Vec { token.as_bytes().to_vec() } @@ -257,3 +270,51 @@ fn token_bytes(token: &str) -> Vec { pub fn clamp_logprob(logprob: f32) -> f32 { logprob.max(-9999.0) } + +#[cfg(test)] +mod tests { + use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob}; + + use super::decoded_logprobs_to_openai_chat; + + fn sample_logprobs() -> DecodedLogprobs { + DecodedLogprobs { + positions: vec![DecodedPositionLogprobs { + entries: vec![ + DecodedTokenLogprob { + token_id: 1, + token: "A".to_string(), + logprob: -0.1, + rank: 1, + }, + DecodedTokenLogprob { + token_id: 2, + token: "B".to_string(), + logprob: -1.0, + rank: 2, + }, + DecodedTokenLogprob { + token_id: 3, + token: "C".to_string(), + logprob: -2.0, + rank: 3, + }, + ], + }], + } + } + + fn chat_top_logprobs_len(top_logprobs: i32) -> usize { + let chat_logprobs = + decoded_logprobs_to_openai_chat(&sample_logprobs(), top_logprobs, false) + .expect("chat logprobs"); + chat_logprobs.content.expect("content")[0].top_logprobs.len() + } + + #[test] + fn chat_logprobs_respects_requested_top_logprobs_count() { + assert_eq!(chat_top_logprobs_len(0), 0); + assert_eq!(chat_top_logprobs_len(1), 1); + assert_eq!(chat_top_logprobs_len(-1), 3); + } +} diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index d72226c2fc65..64e904b562b6 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -2380,6 +2380,14 @@ async fn non_stream_chat_includes_logprobs_and_prompt_logprobs() { json["choices"][0]["logprobs"]["content"][1]["token"], json!("i") ); + assert_eq!( + json["choices"][0]["logprobs"]["content"][0]["top_logprobs"], + json!([]) + ); + assert_eq!( + json["choices"][0]["logprobs"]["content"][1]["top_logprobs"], + json!([]) + ); assert_eq!(json["prompt_logprobs"][0], serde_json::Value::Null); assert!(json["prompt_logprobs"][1].is_object()); } diff --git a/rust/src/text/src/backend/hf/config.rs b/rust/src/text/src/backend/hf/config.rs index 140bfdc30489..2a69632c872b 100644 --- a/rust/src/text/src/backend/hf/config.rs +++ b/rust/src/text/src/backend/hf/config.rs @@ -95,6 +95,7 @@ impl HfSpecialTokens { pub struct ModelConfig { model_type: Option, vocab_size: Option, + eos_token_id: Option, num_experts: Option, moe_num_experts: Option, n_routed_experts: Option, @@ -125,12 +126,16 @@ pub(super) enum OneOrManyTokenIds { } impl OneOrManyTokenIds { - pub(super) fn into_set(self) -> BTreeSet { + pub(super) fn as_slice(&self) -> &[u32] { match self { - Self::One(id) => BTreeSet::from([id]), - Self::Many(ids) => ids.into_iter().collect(), + Self::One(id) => std::slice::from_ref(id), + Self::Many(ids) => ids.as_slice(), } } + + pub(super) fn into_set(self) -> BTreeSet { + self.as_slice().iter().copied().collect() + } } /// Hugging Face configs may expose the expert count either as one integer or @@ -195,6 +200,18 @@ impl ModelConfig { } } + /// Return the effective model-side EOS token ids, following the same + /// simplified text-config selection as `vocab_size`. + pub(super) fn eos_token_ids(&self) -> &[u32] { + if let Some(eos_token_id) = self.eos_token_id.as_ref() { + eos_token_id.as_slice() + } else if let Some(text_config) = self.text_config.as_deref() { + text_config.eos_token_ids() + } else { + &[] + } + } + /// Match Python's current expert-count priority on the selected text /// config. /// @@ -354,6 +371,32 @@ mod tests { assert_eq!(config.vocab_size().unwrap(), 151936); } + #[test] + fn model_config_reads_top_level_eos_token_ids() { + let single: ModelConfig = serde_json::from_str(r#"{"eos_token_id":151645}"#).unwrap(); + let many: ModelConfig = + serde_json::from_str(r#"{"eos_token_id":[128001,128008,128009]}"#).unwrap(); + let null: ModelConfig = serde_json::from_str(r#"{"eos_token_id":null}"#).unwrap(); + + assert_eq!(single.eos_token_ids(), &[151645]); + assert_eq!(many.eos_token_ids(), &[128001, 128008, 128009]); + assert!(null.eos_token_ids().is_empty()); + } + + #[test] + fn model_config_uses_nested_eos_token_ids_when_top_level_is_absent() { + let config: ModelConfig = serde_json::from_str( + r#"{ + "text_config": { + "eos_token_id": [59246, 59253, 59255] + } + }"#, + ) + .unwrap(); + + assert_eq!(config.eos_token_ids(), &[59246, 59253, 59255]); + } + #[test] fn model_config_rejects_missing_vocab_size() { let config: ModelConfig = serde_json::from_str(r#"{}"#).unwrap(); diff --git a/rust/src/text/src/backend/hf/mod.rs b/rust/src/text/src/backend/hf/mod.rs index a145ecc39b1d..49ae5dbd6b95 100644 --- a/rust/src/text/src/backend/hf/mod.rs +++ b/rust/src/text/src/backend/hf/mod.rs @@ -8,7 +8,9 @@ use std::collections::BTreeSet; use std::sync::Arc; use tracing::info; -use vllm_tokenizer::{DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, TiktokenTokenizer}; +use vllm_tokenizer::{ + DynTokenizer, HuggingFaceTokenizer, TekkenTokenizer, TiktokenTokenizer, Tokenizer, +}; use self::config::{GenerationConfig, load_generation_config}; pub use self::config::{ @@ -56,23 +58,15 @@ impl HfTextBackend { pub fn from_resolved_model_files(files: ResolvedModelFiles, model_id: String) -> Result { let tokenizer_config = load_tokenizer_config(files.tokenizer_config_path.as_deref())?; let tokenizer = load_tokenizer(&files.tokenizer)?; - let primary_eos_token_id = tokenizer_config - .special_tokens - .eos_token - .as_ref() - .and_then(|token| tokenizer.token_to_id(token.as_str())); - let model_config = load_model_config(files.config_path.as_deref())?; let model_vocab_size = model_config.vocab_size()? as usize; let generation_config = load_generation_config(files.generation_config_path.as_deref())?; - let mut extra_eos_token_ids = generation_config - .eos_token_id - .clone() - .map(|value| value.into_set()) - .unwrap_or_default(); - if let Some(primary_eos_token_id) = primary_eos_token_id { - extra_eos_token_ids.remove(&primary_eos_token_id); - } + let (primary_eos_token_id, extra_eos_token_ids) = resolve_eos_token_ids( + &tokenizer_config, + &model_config, + &generation_config, + tokenizer.as_ref(), + ); info!( model_id, @@ -98,6 +92,42 @@ impl HfTextBackend { } } +/// Resolve EOS hints from tokenizer, model, and generation configs. +/// +/// Resolution rules: +/// 1. Use the tokenizer-side EOS token as the primary EOS when it resolves to a +/// token id. +/// 2. Fall back to the first model-config EOS id when the tokenizer config does +/// not provide a primary EOS. +/// 3. Keep any remaining model/generation EOS ids as extra stop-token ids so +/// they still participate in stopping and min-token handling. +fn resolve_eos_token_ids( + tokenizer_config: &HfTokenizerConfig, + model_config: &ModelConfig, + generation_config: &GenerationConfig, + tokenizer: &dyn Tokenizer, +) -> (Option, BTreeSet) { + let model_config_eos_token_ids = model_config.eos_token_ids(); + let primary_eos_token_id = tokenizer_config + .special_tokens + .eos_token + .as_ref() + .and_then(|token| tokenizer.token_to_id(token.as_str())) + .or_else(|| model_config_eos_token_ids.first().copied()); + + let mut extra_eos_token_ids = generation_config + .eos_token_id + .clone() + .map(|value| value.into_set()) + .unwrap_or_default(); + extra_eos_token_ids.extend(model_config_eos_token_ids.iter().copied()); + if let Some(primary_eos_token_id) = primary_eos_token_id { + extra_eos_token_ids.remove(&primary_eos_token_id); + } + + (primary_eos_token_id, extra_eos_token_ids) +} + impl TextBackend for HfTextBackend { fn tokenizer(&self) -> DynTokenizer { self.tokenizer.clone() @@ -128,3 +158,77 @@ impl TextBackend for HfTextBackend { }) } } + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use super::{GenerationConfig, HfTokenizerConfig, ModelConfig, resolve_eos_token_ids}; + use vllm_tokenizer::Tokenizer; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + _text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(vec![]) + } + + fn decode( + &self, + _token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(String::new()) + } + + fn token_to_id(&self, token: &str) -> Option { + (token == "").then_some(2) + } + + fn id_to_token(&self, id: u32) -> Option { + (id == 2).then(|| "".to_string()) + } + } + + #[test] + fn eos_resolution_uses_model_config_when_tokenizer_has_no_eos() { + let tokenizer_config: HfTokenizerConfig = serde_json::from_str("{}").unwrap(); + let model_config: ModelConfig = + serde_json::from_str(r#"{"eos_token_id":[200006,200010]}"#).unwrap(); + let generation_config: GenerationConfig = serde_json::from_str("{}").unwrap(); + + let (primary, extra) = resolve_eos_token_ids( + &tokenizer_config, + &model_config, + &generation_config, + &FakeTokenizer, + ); + + assert_eq!(primary, Some(200006)); + assert_eq!(extra, BTreeSet::from([200010])); + } + + #[test] + fn eos_resolution_keeps_tokenizer_eos_primary() { + let tokenizer_config: HfTokenizerConfig = + serde_json::from_str(r#"{"eos_token":""}"#).unwrap(); + let model_config: ModelConfig = + serde_json::from_str(r#"{"eos_token_id":[2,200006]}"#).unwrap(); + let generation_config: GenerationConfig = + serde_json::from_str(r#"{"eos_token_id":[2,200010]}"#).unwrap(); + + let (primary, extra) = resolve_eos_token_ids( + &tokenizer_config, + &model_config, + &generation_config, + &FakeTokenizer, + ); + + assert_eq!(primary, Some(2)); + assert_eq!(extra, BTreeSet::from([200006, 200010])); + } +} diff --git a/rust/src/tokenizer/src/incremental.rs b/rust/src/tokenizer/src/incremental.rs index cf2263d35353..5e470ae4b005 100644 --- a/rust/src/tokenizer/src/incremental.rs +++ b/rust/src/tokenizer/src/incremental.rs @@ -158,17 +158,19 @@ impl IncrementalDecoder for DecodeStream<'_, T> { } fn flush(&mut self, truncate_output_to: Option) -> Result<(Option, String)> { - if !self.ids.is_empty() { + // If the prefix was never seeded (no push_token was called), `ids` + // holds only prompt context — decoding it would re-emit prompt text. + if self.prefix_seeded && !self.ids.is_empty() { let string = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?; let prefix_len = self.prefix.len(); - self.ids.clear(); - self.prefix.clear(); - self.prefix_index = 0; - self.prefix_seeded = true; // Ensure we split at a utf-8 char boundary. self.cumulative_output .push_str(&string[string.floor_char_boundary(prefix_len)..]); } + self.ids.clear(); + self.prefix.clear(); + self.prefix_index = 0; + self.prefix_seeded = true; if let Some(truncate_output_to) = truncate_output_to { self.cumulative_output.truncate(truncate_output_to); } @@ -489,4 +491,31 @@ mod tests { assert_eq!(full_text, "你好A"); assert_eq!(out, "你好A"); } + + #[test] + fn flush_without_push_token_does_not_leak_prompt() { + let backend = Utf8Backend; + let prompt: Vec = b"The quick brown fox jumps over the lazy dog. " + .iter() + .cycle() + .take(7001) + .map(|&b| b as u32) + .collect(); + let mut decoder = backend.create_decode_stream(&prompt, false, 0); + + let (last_chunk, full_text) = decoder.flush(None).unwrap(); + assert_eq!(last_chunk, None); + assert_eq!(full_text, ""); + } + + #[test] + fn flush_without_push_token_does_not_leak_undecodable_prompt_tail() { + let backend = Utf8Backend; + let prompt = vec![0xe4, 0xbd]; + let mut decoder = backend.create_decode_stream(&prompt, false, 0); + + let (last_chunk, full_text) = decoder.flush(None).unwrap(); + assert_eq!(last_chunk, None); + assert_eq!(full_text, ""); + } } diff --git a/setup.py b/setup.py index e8f529701845..a8685157933c 100644 --- a/setup.py +++ b/setup.py @@ -445,6 +445,17 @@ def run(self): dirs_exist_ok=True, ) + tml_fa4_build = os.path.join( + self.build_lib, "vllm", "third_party", "tml_fa4" + ) + if os.path.exists(tml_fa4_build): + print(f"Copying {tml_fa4_build} to vllm/third_party/tml_fa4") + shutil.copytree( + tml_fa4_build, + "vllm/third_party/tml_fa4", + dirs_exist_ok=True, + ) + class precompiled_build_ext(build_ext): """Disables extension building when using precompiled binaries.""" @@ -803,6 +814,7 @@ def extract_precompiled_and_patch_package( # DeepGEMM: extract all files (.py, .so, .cuh, .h, .hpp, etc.) deep_gemm_regex = re.compile(r"vllm/third_party/deep_gemm/.*") fmha_sm100_regex = re.compile(r"vllm/third_party/fmha_sm100/.*") + tml_fa4_regex = re.compile(r"vllm/third_party/tml_fa4/.*") file_members = [] for member in wheel.filelist: if member.filename in exact_members: @@ -828,6 +840,7 @@ def extract_precompiled_and_patch_package( or triton_kernels_regex.match(member.filename) or flashmla_regex.match(member.filename) or deep_gemm_regex.match(member.filename) + or tml_fa4_regex.match(member.filename) or fmha_sm100_regex.match(member.filename) ): file_members.append(member) @@ -1075,6 +1088,11 @@ def _read_requirements(filename: str) -> list[str]: # vllm-flash-attn is built only for CUDA 12.x. # Skip for other versions. continue + if "flashinfer-cubin" in req: + # Not on PyPI since 0.6.14 (only https://flashinfer.ai/whl), so + # it cannot be a wheel dependency; flashinfer falls back to + # fetching cubins at runtime when the package is absent. + continue if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12": # [cu13] extra is the default; strip it on CUDA 12 builds. req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl") @@ -1141,6 +1159,8 @@ def _read_requirements(filename: str) -> list[str]: ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True)) # fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party. ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True)) + # tml-fa4 is copied into an isolated vllm.third_party package. + ext_modules.append(CMakeExtension(name="vllm.tml_fa4", optional=True)) if _is_cpu(): import platform @@ -1168,6 +1188,7 @@ def _read_requirements(filename: str) -> list[str]: "entrypoints/serve/instrumentator/static/*.js", "entrypoints/serve/instrumentator/static/*.css", "distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp", + "third_party/flash_linear_attention/LICENSE", # DeepGEMM JIT include headers (vendored via cmake) "third_party/deep_gemm/include/**/*.cuh", "third_party/deep_gemm/include/**/*.h", diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py new file mode 100644 index 000000000000..b5e354ddfd82 --- /dev/null +++ b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py @@ -0,0 +1,478 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os + +import pytest +import torch + +import vllm.config +from tests.compile.backend import TestBackend +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops +from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP +from vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion import ( + QkNormRopeKvCacheFusionPass, +) +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass +from vllm.compilation.passes.utility.scatter_split_replace import ( + ScatterSplitReplacementPass, +) +from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass +from vllm.config import ( + CacheConfig, + CompilationConfig, + CompilationMode, + ModelConfig, + PassConfig, + VllmConfig, +) +from vllm.forward_context import get_forward_context, set_forward_context +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding +from vllm.platforms import current_platform +from vllm.v1.attention.backend import ( + AttentionBackend, + CommonAttentionMetadata, +) +from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.kv_cache_interface import AttentionSpec + +INDEX_SELECT_OP = torch.ops.aten.index.Tensor +FP8_DTYPE = current_platform.fp8_dtype() + + +class QKNormRoPEKVCacheTestModel(torch.nn.Module): + """Minimal model that reproduces the QK-norm + RoPE + KV cache update + pattern matched by QkNormRopeKvCacheFusionPass: + + q, k, v = split(qkv) + q = rms_norm(q.view(heads, dim), q_weight).view(flat) + k = rms_norm(k.view(heads, dim), k_weight).view(flat) + q, k = rotary_emb(positions, q, k) + q = q.view(num_heads, head_dim) + k = k.view(num_kv_heads, head_dim) + v = v.view(num_kv_heads, head_dim) + dummy = unified_kv_cache_update(k, v, layer_name) + """ + + def __init__( + self, + vllm_config: VllmConfig, + attn_backend: AttentionBackendEnum, + num_heads: int, + num_kv_heads: int, + head_size: int, + is_neox: bool, + rms_norm_eps: float, + dtype: torch.dtype, + device: torch.device, + rotary_dim: int | None = None, + prefix: str = "model.layers.0.self_attn.attn", + ): + super().__init__() + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_size = head_size + self.rotary_dim = rotary_dim if rotary_dim is not None else head_size + self.block_size = vllm_config.cache_config.block_size + self.q_size = num_heads * head_size + self.kv_size = num_kv_heads * head_size + self.is_neox = is_neox + self.dtype = dtype + self.device = device + self.layer_name = prefix + + self.q_norm = RMSNorm(head_size, eps=rms_norm_eps) + self.k_norm = RMSNorm(head_size, eps=rms_norm_eps) + + self.rotary_emb = RotaryEmbedding( + head_size, + rotary_dim=self.rotary_dim, + max_position_embeddings=4096, + base=10000, + is_neox_style=is_neox, + dtype=self.dtype, + ) + + self.enable_rope_custom_op = self.rotary_emb.enabled() + + self.attn = Attention( + num_heads=num_heads, + head_size=head_size, + scale=1.0 / head_size**0.5, + num_kv_heads=num_kv_heads, + cache_config=vllm_config.cache_config, + quant_config=vllm_config.quant_config, + prefix=prefix, + attn_backend=attn_backend.get_class(), + ) + self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend() + assert not self.attn_backend.forward_includes_kv_cache_update, ( + f"Attention backend {self.attn_backend} does not support " + "fuse_qk_norm_rope_kvcache." + ) + kv_cache_dtype_str = vllm_config.cache_config.cache_dtype + self.kv_cache_dtype = ( + FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype + ) + + if self.kv_cache_dtype != self.dtype: + self.attn._k_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + self.attn._v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) + self.attn._k_scale_float = 1.0 + self.attn._v_scale_float = 1.0 + else: + self.attn._k_scale = self.attn._k_scale.to(device) + self.attn._v_scale = self.attn._v_scale.to(device) + + self.builder = self.attn.attn_backend.get_builder_cls()( + kv_cache_spec=AttentionSpec( + block_size=self.block_size, + num_kv_heads=self.num_kv_heads, + head_size=head_size, + dtype=self.kv_cache_dtype, + ), + layer_names=[self.attn.layer_name], + vllm_config=vllm_config, + device=device, + ) + + def build_attn_metadata( + self, batch_size: int, kv_stride_order: tuple[int, ...] | None = None + ) -> CommonAttentionMetadata: + batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size) + common_attn_metadata = create_common_attn_metadata( + batch_spec, self.block_size, self.device, arange_block_indices=True + ) + + max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size + num_blocks = batch_size * max_blocks + + attn_backend = self.attn.attn_backend + kv_cache_shape = attn_backend.get_kv_cache_shape( + num_blocks, self.block_size, self.num_kv_heads, self.head_size + ) + # Caller can force a physical layout; else use the backend's. + if kv_stride_order is None: + try: + kv_stride_order = attn_backend.get_kv_cache_stride_order() + except (AttributeError, NotImplementedError): + kv_stride_order = tuple(range(len(kv_cache_shape))) + + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_stride_order) + inv_order = [kv_stride_order.index(i) for i in range(len(kv_stride_order))] + + raw_tensor = torch.zeros( + 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size, + dtype=self.kv_cache_dtype, + device=self.device, + ) + raw_tensor = raw_tensor.view(kv_cache_shape) + kv_cache = raw_tensor.permute(*inv_order) + + # Store as a bare tensor (not wrapped in a list) to match production + # `bind_kv_cache` behavior. `get_attention_context` returns this + # attribute directly to the fused/unfused `do_kv_cache_update` impls, + # which call `kv_cache.unbind(0)` and therefore require a tensor. + self.attn.kv_cache = kv_cache + + attn_metadata = self.builder.build( + common_prefix_len=0, common_attn_metadata=common_attn_metadata + ) + + return attn_metadata + + def forward( + self, qkv: torch.Tensor, positions: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + qkv = qkv.clone() + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + # QK-norm: RMSNorm on per-head Q and K + q = q.view(-1, self.num_heads, self.head_size) + q = self.q_norm(q) + q = q.view(-1, self.q_size) + + k = k.view(-1, self.num_kv_heads, self.head_size) + k = self.k_norm(k) + k = k.view(-1, self.kv_size) + + # RoPE + q, k = self.rotary_emb(positions, q, k) + + # Mirror Attention.forward: quant-query impls consume an fp8 q. + if ( + self.kv_cache_dtype != self.dtype + and self.attn.impl.supports_quant_query_input + ): + q_fp8 = torch.empty_like(q, dtype=FP8_DTYPE) + torch.ops.vllm.rocm_aiter_per_tensor_quant( + q_fp8, q, self.attn._q_scale, False + ) + q = q_fp8 + + # Final views + KV cache update + q = q.view(-1, self.num_heads, self.head_size) + k = k.view(-1, self.num_kv_heads, self.head_size) + v = v.view(-1, self.num_kv_heads, self.head_size) + kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update( + k, v, self.layer_name + ) + return q, k, v, kv_cache_dummy_dep + + def ops_in_model_before(self) -> list[torch._ops.OpOverload]: + ops: list[torch._ops.OpOverload] = [] + # RoPE is not yet IR-migrated, so its custom op still surfaces + # directly in the graph based on `enable_rope_custom_op`. + if self.enable_rope_custom_op: + if rocm_aiter_ops.is_triton_rotary_embed_enabled(): + ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default) + else: + ops.append(ROTARY_OP) + else: + ops.append(INDEX_SELECT_OP) + ops.append(torch.ops.vllm.unified_kv_cache_update.default) + return ops + + def ops_in_model_after(self) -> list[torch._ops.OpOverload]: + return [torch.ops.vllm.fused_qk_norm_rope_and_unified_kv_cache_update.default] + + +def _run_qk_norm_rope_kvcache_fusion_test( + *, + attn_backend: AttentionBackendEnum, + enable_aiter_triton_rope: bool, + num_tokens: int, + num_heads: int, + num_kv_heads: int, + head_size: int, + rotary_dim: int, + block_size: int, + is_neox: bool, + use_shuffle_kv_layout: str, + kv_stride_order: tuple[int, ...], + dtype: torch.dtype, + kv_cache_dtype: str, + rms_norm_eps: float, + custom_op: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + device = os.environ.get("VLLM_TEST_CUDA_DEVICE", "cuda") + torch.set_default_device(device) + torch.set_default_dtype(dtype) + torch.manual_seed(0) + + vllm_config = VllmConfig( + model_config=ModelConfig(dtype=dtype), + cache_config=CacheConfig( + block_size=block_size, + cache_dtype=kv_cache_dtype, + ), + compilation_config=CompilationConfig( + mode=CompilationMode.VLLM_COMPILE, + custom_ops=[custom_op], + pass_config=PassConfig( + fuse_qk_norm_rope_kvcache=True, + eliminate_noops=True, + ), + ), + ) + + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.setenv( + "VLLM_ROCM_USE_AITER_TRITON_ROPE", + "1" if enable_aiter_triton_rope else "0", + ) + m.setenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", use_shuffle_kv_layout) + rocm_aiter_ops.refresh_env_variables() + + model = QKNormRoPEKVCacheTestModel( + vllm_config=vllm_config, + attn_backend=attn_backend, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_size, + rotary_dim=rotary_dim, + is_neox=is_neox, + rms_norm_eps=rms_norm_eps, + dtype=dtype, + device=torch.get_default_device(), + ) + + fusion_pass = QkNormRopeKvCacheFusionPass(vllm_config) + passes = [ + NoOpEliminationPass(vllm_config), + SplitCoalescingPass(vllm_config), + ScatterSplitReplacementPass(vllm_config), + fusion_pass, + PostCleanupPass(vllm_config), + ] + backend = TestBackend(*passes) + + qkv = torch.randn( + num_tokens, + num_heads * head_size + 2 * num_kv_heads * head_size, + dtype=dtype, + ) + pos = torch.arange(num_tokens, dtype=torch.long) + + qkv_unfused = qkv.clone() + pos_unfused = pos.clone() + + # Run unfused (eager) forward + with set_forward_context(None, vllm_config): + forward_context = get_forward_context() + attn_metadata = model.build_attn_metadata(num_tokens, kv_stride_order) + forward_context.slot_mapping = { + model.layer_name: attn_metadata.slot_mapping + } + q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused) + attn_layer = forward_context.no_compile_layers[model.layer_name] + kv_cache_unfused = attn_layer.kv_cache + del dummy + + # Run fused (compiled) forward + torch._dynamo.mark_dynamic(qkv, 0) + torch._dynamo.mark_dynamic(pos, 0) + with set_forward_context(None, vllm_config): + model_fused = torch.compile(model, backend=backend) + forward_context = get_forward_context() + attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_stride_order) + forward_context.slot_mapping = { + model.layer_name: attn_metadata.slot_mapping + } + q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos) + attn_layer = forward_context.no_compile_layers[model.layer_name] + kv_cache_fused = attn_layer.kv_cache + del dummy + + assert fusion_pass.matched_count == 1 + + backend.check_before_ops(model.ops_in_model_before()) + backend.check_after_ops(model.ops_in_model_after()) + + # Sweep-backed (18.2k pts, PR #42749): native-rope ref worst 7.7e-3 -> 1e-2; + # AITER-triton-rope ref is itself approximate (plateau 1.28e-2) -> 2e-2. + ATOL, RTOL = (2e-2, 2e-2) if enable_aiter_triton_rope else (1e-2, 1e-2) + is_fp8_cache = model.kv_cache_dtype != dtype + + if q_fused.dtype == FP8_DTYPE: + # Quant-query path: both q are fp8; compare dequant within 1 fp8 ULP. + torch.testing.assert_close( + q_unfused.float(), q_fused.float(), atol=1.25e-1, rtol=1.25e-1 + ) + else: + torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL) + + if not is_fp8_cache: + # The AITER PTS kernel populates k_out only for non-FP8 caches. + # With FP8, the kernel writes quantized K directly to the cache + # and may leave k_out uninitialised. In production this is fine + # because downstream attention reads K from the cache. + torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL) + + # Should be bit exact since no processing had been done on v for both paths + torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0) + + # fp8 vs triton-rope ref requires loosening tolerance to 1.25e-1. + if is_fp8_cache and enable_aiter_triton_rope: + cache_atol = cache_rtol = 1.25e-1 + else: + cache_atol, cache_rtol = ATOL, RTOL + + torch.testing.assert_close( + kv_cache_unfused[0].float(), + kv_cache_fused[0].float(), + atol=cache_atol, + rtol=cache_rtol, + ) + + +_FUSION_CONFIGS = [ + # Full rotary, both neox styles (the original coverage). + pytest.param(64, 8, 64, 64, True, id="full-neox"), + pytest.param(64, 8, 64, 64, False, id="full-non_neox"), + # GLM-4.5/4.6/4.7 (glm4_moe.py:275 partial_rotary_factor=0.5, neox-style) + pytest.param(32, 8, 128, 64, True, id="glm4_moe"), + # GLM-4 dense (glm4.py:97,124 partial_rotary_factor=0.5, non-neox) + pytest.param(32, 2, 128, 64, False, id="glm4_dense"), + # Moondream3-style small head (head_size=64, rotary_dim=32) + pytest.param(16, 2, 64, 32, True, id="partial_small_head"), +] + + +@pytest.mark.parametrize( + "num_heads, num_kv_heads, head_size, rotary_dim, is_neox", + _FUSION_CONFIGS, +) +@pytest.mark.parametrize( + "attn_backend", + [ + AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN, + AttentionBackendEnum.ROCM_AITER_FA, + ], +) +@pytest.mark.parametrize("num_tokens", [5, 16, 2048]) +@pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"]) +@pytest.mark.parametrize( + "kv_stride_order", + [ + pytest.param((0, 1, 2, 3, 4), id="block_first"), + pytest.param((1, 0, 2, 3, 4), id="kv_first"), + ], +) +@pytest.mark.parametrize("enable_aiter_triton_rope", [True, False]) +@pytest.mark.parametrize("block_size", [16, 32, 64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("rms_norm_eps", [1e-5, 1e-6]) +@pytest.mark.parametrize("custom_op", ["+rotary_embedding", "+rms_norm"]) +@pytest.mark.skipif( + not is_aiter_found_and_supported(), + reason="Only test on ROCm with AITER installed and supported", +) +def test_qk_norm_rope_kvcache_fusion( + num_tokens: int, + num_heads: int, + num_kv_heads: int, + head_size: int, + rotary_dim: int, + is_neox: bool, + attn_backend: AttentionBackendEnum, + enable_aiter_triton_rope: bool, + use_shuffle_kv_layout: str, + kv_stride_order: tuple[int, ...], + block_size: int, + dtype: torch.dtype, + kv_cache_dtype: str, + rms_norm_eps: float, + custom_op: str, + monkeypatch: pytest.MonkeyPatch, +): + if ( + attn_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN + and use_shuffle_kv_layout == "1" + ): + pytest.skip("ROCM_AITER_UNIFIED_ATTN is NHD-only; shuffle env is ignored") + _run_qk_norm_rope_kvcache_fusion_test( + attn_backend=attn_backend, + enable_aiter_triton_rope=enable_aiter_triton_rope, + num_tokens=num_tokens, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_size=head_size, + rotary_dim=rotary_dim, + block_size=block_size, + is_neox=is_neox, + use_shuffle_kv_layout=use_shuffle_kv_layout, + kv_stride_order=kv_stride_order, + dtype=dtype, + kv_cache_dtype=kv_cache_dtype, + rms_norm_eps=rms_norm_eps, + custom_op=custom_op, + monkeypatch=monkeypatch, + ) diff --git a/tests/compile/test_graph_partition.py b/tests/compile/test_graph_partition.py index 8e20b704facc..bb2a6f2aee55 100644 --- a/tests/compile/test_graph_partition.py +++ b/tests/compile/test_graph_partition.py @@ -701,3 +701,67 @@ def test_decompose_size_with_getitem_user(): f"getitem node '{node.name}' has {len(node.args)} args " f"(expected 2): {node.args}" ) + + +def test_decompose_size_leaves_scalar_size_with_dim(): + """ + Regression test: _decompose_size_nodes must leave x.size(dim) alone. + + x.size() returns a torch.Size tuple that can't cross split boundaries and + must be decomposed. x.size(dim), however, already returns a scalar + SymInt/int that crosses fine, so the pass must not touch it. + + The punica LoRA path traces token_lora_mapping[:x.size(0)] under a dynamic + batch dim, so the size(0) node ends up nested inside a slice object: + + %size = call_method[target="size"](args = (%x, 0)) + %slice = call_function[target=getitem]( + args = (%mapping, slice(None, %size, None))) + + The old pass tried to decompose this scalar node too and then erase it, but + the slice still referenced it, raising "Tried to erase Node size but it + still had N users". The fix skips size calls that carry a dim argument. + """ + from torch._dynamo.source import LocalSource + from torch._subclasses.fake_tensor import FakeTensorMode + from torch.fx.experimental.symbolic_shapes import ShapeEnv + + # Build graph: + # %x = placeholder + # %mapping = placeholder + # %size = x.size(0) # scalar, with a dim arg + # %sliced = mapping[slice(None, %size, None)] # size node inside a slice + graph = fx.Graph() + x = graph.placeholder("x") + mapping = graph.placeholder("token_lora_mapping") + size_node = graph.call_method("size", args=(x, 0)) + sliced_node = graph.call_function( + operator.getitem, + args=(mapping, slice(None, size_node, None)), + ) + graph.output((sliced_node,)) + + # dim 0 dynamic (SymInt) — the realistic Unsloth + LoRA case. Without the + # skip, the pass would build per-dim replacements and then crash trying to + # erase the still-referenced size node. + shape_env = ShapeEnv() + src = LocalSource("tokens") + sym_tokens = shape_env.create_symintnode(shape_env.create_symbol(4, src), hint=4) + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + fake_x = torch.empty_strided((sym_tokens, 8), (8, 1)) + x.meta["example_value"] = fake_x + + gm = fx.GraphModule(torch.nn.Module(), graph) + + # Must not raise "Tried to erase Node ... still had N users". + _decompose_size_nodes(gm) + + # The scalar x.size(0) node is left in place, untouched. + remaining = list(gm.graph.find_nodes(op="call_method", target="size")) + assert len(remaining) == 1, ( + f"x.size(0) should be left untouched, found {len(remaining)} size nodes" + ) + assert remaining[0].args == (x, 0), ( + f"size node args changed: {remaining[0].args} (expected (x, 0))" + ) diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index ddb8752a80d6..7e425d68eecb 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -85,6 +85,34 @@ def record(hf_config: PretrainedConfig) -> PretrainedConfig: assert seen_architectures == ["MiMoMTPModel"] +@pytest.mark.cpu_test +def test_inkling_override_exposes_only_first_mtp_depth(): + text_config = _make_hf_config( + architectures=["InklingForCausalLM"], + model_type="inkling_model", + local_layer_ids=[1, 3], + ) + config = _make_hf_config( + architectures=["InklingForConditionalGeneration"], + model_type="inkling_mm_model", + text_config=text_config, + mtp_config={ + "num_nextn_predict_layers": 8, + "local_layer_ids": [0, 2, 4], + }, + ) + + out = SpeculativeConfig.hf_config_override(config) + + assert out is text_config + assert out.model_type == "inkling_mtp" + assert out.architectures == ["InklingMTPModel"] + assert out.n_predict == 1 + assert out.num_nextn_predict_layers == 8 + assert out.chain_hidden_post_norm is False + assert out.local_layer_ids == [0, 2, 4] + + def _module_level_shrink(hf_config: PretrainedConfig) -> PretrainedConfig: hf_config.num_hidden_layers = 1 return hf_config diff --git a/tests/distributed/test_kv_cache_events.py b/tests/distributed/test_kv_cache_events.py index aa39ab17b30a..72e0fec54f7e 100644 --- a/tests/distributed/test_kv_cache_events.py +++ b/tests/distributed/test_kv_cache_events.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + +import msgspec import pytest from vllm.distributed.kv_events import BlockRemoved, BlockStored @@ -9,9 +12,44 @@ _FAKE_HASH: bytes = b"\xab" * 32 +class _LegacyBlockStored( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] + tag="BlockStored", # type: ignore[call-arg] +): + """BlockStored wire schema before locality was added.""" + + block_hashes: list[bytes] + parent_block_hash: bytes | None + token_ids: list[int] + block_size: int + lora_id: int | None + medium: str | None + lora_name: str | None + extra_keys: list[tuple[Any, ...] | None] | None = None + group_idx: int | None = None + kv_cache_spec_kind: str | None = None + kv_cache_spec_sliding_window: int | None = None + + +class _LegacyBlockRemoved( + msgspec.Struct, + omit_defaults=True, # type: ignore[call-arg] + gc=False, # type: ignore[call-arg] + tag="BlockRemoved", # type: ignore[call-arg] +): + """BlockRemoved wire schema before locality was added.""" + + block_hashes: list[bytes] + medium: str | None + group_idx: int | None = None + + def _make_block_stored( group_idx: int | None = None, kv_cache_spec_sliding_window: int | None = None, + locality: str | None = None, ) -> BlockStored: return BlockStored( block_hashes=[_FAKE_HASH], @@ -23,16 +61,19 @@ def _make_block_stored( lora_name=None, group_idx=group_idx, kv_cache_spec_sliding_window=kv_cache_spec_sliding_window, + locality=locality, ) def _make_block_removed( group_idx: int | None = None, + locality: str | None = None, ) -> BlockRemoved: return BlockRemoved( block_hashes=[_FAKE_HASH], medium="GPU", group_idx=group_idx, + locality=locality, ) @@ -84,3 +125,61 @@ def test_block_stored_hash_differs_by_sliding_window(): event_a = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=128) event_b = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=256) assert hash(event_a) != hash(event_b) + + +@pytest.mark.parametrize( + ("event_a", "event_b"), + [ + ( + _make_block_stored(locality="LOCAL"), + _make_block_stored(locality="REMOTE"), + ), + ( + _make_block_removed(locality="LOCAL"), + _make_block_removed(locality="REMOTE"), + ), + ], +) +def test_event_hash_differs_by_locality( + event_a: BlockStored | BlockRemoved, + event_b: BlockStored | BlockRemoved, +): + assert hash(event_a) != hash(event_b) + + +def test_block_stored_locality_is_wire_compatible(): + legacy = _LegacyBlockStored( + block_hashes=[_FAKE_HASH], + parent_block_hash=None, + token_ids=[1, 2, 3, 4], + block_size=4, + lora_id=None, + medium="GPU", + lora_name=None, + group_idx=2, + kv_cache_spec_sliding_window=128, + ) + legacy_payload = msgspec.msgpack.encode(legacy) + assert ( + msgspec.msgpack.encode( + _make_block_stored( + group_idx=2, + kv_cache_spec_sliding_window=128, + ) + ) + == legacy_payload + ) + assert msgspec.msgpack.decode(legacy_payload, type=BlockStored).locality is None + new_payload = msgspec.msgpack.encode(_make_block_stored(locality="LOCAL")) + assert msgspec.msgpack.decode(new_payload)["locality"] == "LOCAL" + assert msgspec.msgpack.decode(new_payload, type=_LegacyBlockStored).medium == "GPU" + + +def test_block_removed_locality_is_wire_compatible(): + legacy = _LegacyBlockRemoved(block_hashes=[_FAKE_HASH], medium="GPU") + legacy_payload = msgspec.msgpack.encode(legacy) + assert msgspec.msgpack.encode(_make_block_removed()) == legacy_payload + assert msgspec.msgpack.decode(legacy_payload, type=BlockRemoved).locality is None + new_payload = msgspec.msgpack.encode(_make_block_removed(locality="REMOTE")) + assert msgspec.msgpack.decode(new_payload)["locality"] == "REMOTE" + assert msgspec.msgpack.decode(new_payload, type=_LegacyBlockRemoved).medium == "GPU" diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index 95c905fc0803..fb2d9bb832e7 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -206,6 +206,36 @@ class _AttnMeta: # should run even when FlashInfer NVLink backends are not installed. +@pytest.mark.parametrize("supports_output", [False, True]) +def test_one_sided_combine_into_compatibility(supports_output): + from vllm.distributed.device_communicators.all2all import ( + FlashInferNVLinkOneSidedManager, + ) + + class FakeMoeAlltoAll: + def combine( + self, + payload, + runtime_max_tokens_per_rank, + output=None, + ): + result = payload + runtime_max_tokens_per_rank + if output is None: + return result + output.copy_(result) + return output + + manager = FlashInferNVLinkOneSidedManager.__new__(FlashInferNVLinkOneSidedManager) + manager.moe_alltoall = FakeMoeAlltoAll() + manager._combine_supports_output = supports_output + payload = torch.arange(4, dtype=torch.float32) + output = torch.empty_like(payload) + + manager.combine_into(payload, runtime_max_tokens_per_rank=2, output=output) + + torch.testing.assert_close(output, payload + 2) + + # --------------------------------------------------------------------------- # Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init) # --------------------------------------------------------------------------- diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index f3423745ca59..b79aa1974d13 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -17,7 +17,19 @@ from vllm.config.parallel import ParallelConfig from vllm.config.weight_transfer import WeightTransferConfig -from vllm.distributed.weight_transfer import WeightTransferEngineFactory +from vllm.distributed.weight_transfer import ( + HTTPVLLMWeightSyncClient, + ModuleSource, + RayVLLMWeightSyncClient, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightTransferEngineFactory, + WeightTransferTrainerFactory, +) +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) from vllm.distributed.weight_transfer.ipc_engine import ( IPCWeightTransferEngine, IPCWeightTransferInitInfo, @@ -1214,3 +1226,194 @@ def test_ipc_receive_weights_missing_gpu_uuid_raises(): with pytest.raises(ValueError, match="IPC handle not found"): engine.receive_weights(update_info) + + +class RecordingClient: + """A fake VLLMWeightSyncClient that records the order of calls.""" + + def __init__(self): + self.order: list[str] = [] + self.last_init_info: dict | None = None + self.last_update_info: dict | None = None + + def init_weight_transfer_engine(self, init_info: dict) -> None: + self.order.append("init") + self.last_init_info = init_info + + def start_weight_update(self) -> None: + self.order.append("start") + + def update_weights(self, update_info: dict) -> None: + self.order.append("update") + self.last_update_info = update_info + + def finish_weight_update(self) -> None: + self.order.append("finish") + + +def _module_with(*pairs): + """A tiny nn.Module exposing the given (name, tensor) pairs as parameters, + so trainer tests can build a ModuleSource without a real model.""" + module = torch.nn.Module() + for name, tensor in pairs: + module.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False)) + return module + + +class _DummyTrainerEngine(TrainerWeightTransferEngine): + """Minimal concrete trainer engine to exercise base-class + factory.""" + + @classmethod + def trainer_init(cls, config, init_info, *, client, source): + return cls(config, client=client, source=source) + + def send_weights(self): + pass + + +class TestTrainerClients: + """Structural protocol conformance for the built-in clients.""" + + def test_recording_client_is_protocol(self): + assert isinstance(RecordingClient(), VLLMWeightSyncClient) + + def test_http_client_is_protocol(self): + assert isinstance( + HTTPVLLMWeightSyncClient("http://localhost:8000"), VLLMWeightSyncClient + ) + + def test_ray_client_is_protocol(self): + assert isinstance(RayVLLMWeightSyncClient(MagicMock()), VLLMWeightSyncClient) + + def test_ray_client_sends_typed_requests(self, monkeypatch): + """Ray client must hand the actor typed Request objects, not raw dicts.""" + import ray + + monkeypatch.setattr(ray, "get", lambda refs: None) + handle = MagicMock() + client = RayVLLMWeightSyncClient(handle) + + client.init_weight_transfer_engine({"master_addr": "x"}) + (init_req,), _ = handle.init_weight_transfer_engine.remote.call_args + assert isinstance(init_req, WeightTransferInitRequest) + assert init_req.init_info == {"master_addr": "x"} + + client.update_weights({"names": ["w"]}) + (update_req,), _ = handle.update_weights.remote.call_args + assert isinstance(update_req, WeightTransferUpdateRequest) + assert update_req.update_info == {"names": ["w"]} + + def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch): + """HTTP update_weights must encode raw ipc_handles as a base64 pickle.""" + captured = {} + + def fake_post(self, path, json=None): + captured["path"] = path + captured["json"] = json + + monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post) + client = HTTPVLLMWeightSyncClient("http://localhost:8000") + client.update_weights({"names": ["w"], "ipc_handles": [{"gpu": ("args",)}]}) + sent = captured["json"]["update_info"] + assert "ipc_handles" not in sent + assert "ipc_handles_pickled" in sent + assert pickle.loads(base64.b64decode(sent["ipc_handles_pickled"])) == [ + {"gpu": ("args",)} + ] + + def test_http_client_passes_through_nccl_update_info(self, monkeypatch): + """NCCL update_info has only JSON-native fields and passes unchanged.""" + captured = {} + + def fake_post(self, path, json=None): + captured["json"] = json + + monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post) + client = HTTPVLLMWeightSyncClient("http://localhost:8000") + update_info = {"names": ["w"], "dtype_names": ["float32"], "shapes": [[4]]} + client.update_weights(update_info) + assert captured["json"]["update_info"] == update_info + + +class TestModuleSource: + """`ModuleSource` metadata vs. materialized iteration (dense, no GPU).""" + + def test_metadata_reads_shape_and_dtype(self): + source = ModuleSource( + _module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3))) + ) + meta = source.metadata() + assert [m.name for m in meta] == ["w", "b"] + assert [m.shape for m in meta] == [(2, 3), (3,)] + assert all(m.dtype == torch.float32 for m in meta) + + def test_iteration_yields_materialized_tensors(self): + w = torch.arange(6, dtype=torch.float32).reshape(2, 3) + source = ModuleSource(_module_with(("w", w))) + pairs = list(source) + assert [name for name, _ in pairs] == ["w"] + assert torch.equal(pairs[0][1], w) + + def test_source_is_reiterable(self): + source = ModuleSource(_module_with(("w", torch.zeros(2)))) + assert [n for n, _ in source] == [n for n, _ in source] == ["w"] + + +class TestTrainerFactory: + """WeightTransferTrainerFactory registry mechanics.""" + + def test_builtin_registry_has_no_trainer_backends_yet(self): + # Concrete backends register in the per-backend migration PRs. + assert WeightTransferTrainerFactory._registry == {} + + def test_register_and_dispatch(self): + saved = dict(WeightTransferTrainerFactory._registry) + try: + WeightTransferTrainerFactory.register_engine("dummy", _DummyTrainerEngine) + engine = WeightTransferTrainerFactory.trainer_init( + "dummy", + WeightTransferConfig(backend="dummy"), + MagicMock(), + client=RecordingClient(), + source=ModuleSource(_module_with(("w", torch.zeros(2)))), + ) + assert isinstance(engine, _DummyTrainerEngine) + with pytest.raises(ValueError, match="already registered"): + WeightTransferTrainerFactory.register_engine( + "dummy", _DummyTrainerEngine + ) + finally: + WeightTransferTrainerFactory._registry = saved + + def test_unknown_backend_raises(self): + with pytest.raises(ValueError, match="Invalid weight transfer backend"): + WeightTransferTrainerFactory.trainer_init( + "nope", + WeightTransferConfig(backend="nope"), + MagicMock(), + client=RecordingClient(), + source=ModuleSource(_module_with(("w", torch.zeros(2)))), + ) + + +class TestTrainerEngineBase: + """Base-class construction (no GPU).""" + + def test_source_stored_and_sender_by_default(self): + engine = _DummyTrainerEngine( + WeightTransferConfig(backend="nccl"), + client=RecordingClient(), + source=ModuleSource(_module_with(("w", torch.zeros(2)))), + ) + assert engine.is_sender is True + assert [name for name, _ in engine.source] == ["w"] + + def test_shutdown_default_is_noop(self): + engine = _DummyTrainerEngine( + WeightTransferConfig(backend="nccl"), + client=RecordingClient(), + source=ModuleSource(_module_with(("w", torch.zeros(2)))), + is_sender=False, + ) + assert engine.is_sender is False + engine.shutdown() # must not raise diff --git a/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py new file mode 100644 index 000000000000..8c03fc1daa01 --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_extra_content_fields.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for the TranslateGemma extra-fields-on-content-parts pathway. + +TranslateGemma's bundled chat template reads ``source_lang_code`` and +``target_lang_code`` off each content part to assemble the translation prompt. +This PR makes sure those extra fields survive request parsing and reach the +template; these tests verify the full pathway with the real model on both +the text and image content branches of the template. +""" + +import json + +import openai +import pytest +import pytest_asyncio + +from tests.utils import RemoteOpenAIServer +from vllm.assets.image import ImageAsset +from vllm.multimodal.utils import encode_image_url + +MODEL_NAME = "google/translategemma-4b-it" + + +@pytest.fixture(scope="module") +def server(): + args = [ + "--max-model-len", + "2048", + "--max-num-seqs", + "16", + "--enforce-eager", + "--chat-template-content-format", + "openai", + "--limit-mm-per-prompt", + json.dumps({"image": 1}), + ] + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client(server): + async with server.get_async_client() as async_client: + yield async_client + + +@pytest.fixture(scope="module") +def stop_sign_image_url(): + return encode_image_url(ImageAsset("stop_sign").pil_image) + + +@pytest.mark.asyncio +async def test_translategemma_extra_lang_code_fields(client: openai.AsyncOpenAI): + """en -> es translation through TranslateGemma's bundled chat template's + ``text`` content branch, which depends on ``source_lang_code`` / + ``target_lang_code`` extra fields being preserved on the content part.""" + completion = await client.chat.completions.create( + model=MODEL_NAME, + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "The quick brown fox jumps over the lazy dog.", + "source_lang_code": "en", + "target_lang_code": "es", + } + ], + } + ], + max_tokens=64, + temperature=0.0, + ) + + content = completion.choices[0].message.content + assert content is not None and content.strip(), ( + "expected a non-empty Spanish translation" + ) + # If extra fields had been stripped, the bundled template would not + # produce a translation prompt and the output would not be Spanish. + spanish_markers = (" el ", " la ", " los ", " las ", " perro", " zorro") + lowered = f" {content.lower()} " + assert any(m in lowered for m in spanish_markers), ( + f"output does not look like Spanish: {content!r}" + ) + + +@pytest.mark.asyncio +async def test_translategemma_image_extra_lang_code_fields( + client: openai.AsyncOpenAI, + stop_sign_image_url: str, +): + """en -> es OCR-translation through TranslateGemma's bundled chat + template's ``image`` content branch. Exercises the multimodal branch of + ``_collect_extra_fields`` in ``_parse_chat_message_content_part``: the + ``source_lang_code`` / ``target_lang_code`` extras must survive parsing + and end up alongside the ``{"type": "image"}`` placeholder so the + template's image branch (``content["type"] == 'image'``) renders the + right OCR-translation prompt.""" + completion = await client.chat.completions.create( + model=MODEL_NAME, + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": stop_sign_image_url}, + "source_lang_code": "en", + "target_lang_code": "es", + } + ], + } + ], + max_tokens=64, + temperature=0.0, + ) + + content = completion.choices[0].message.content + assert content is not None and content.strip(), ( + "expected a non-empty translation of the stop-sign image" + ) + # Common Spanish renderings of "STOP" on a road sign. If extras were + # stripped we would either fail the template or get back the original + # English word, neither of which match this set. + spanish_stop_terms = ("detén", "deteng", "deten", "alto", "pare", "para") + lowered = content.lower() + assert "stop" not in lowered, ( + f"output appears to echo the English source: {content!r}" + ) + assert any(t in lowered for t in spanish_stop_terms), ( + f"output is not a recognizable Spanish translation of 'STOP': {content!r}" + ) diff --git a/tests/entrypoints/openai/responses/test_parsable_context.py b/tests/entrypoints/openai/responses/test_parsable_context.py index 292edda9a7c4..8ff3a1eeae8c 100644 --- a/tests/entrypoints/openai/responses/test_parsable_context.py +++ b/tests/entrypoints/openai/responses/test_parsable_context.py @@ -38,7 +38,7 @@ def server(): "--reasoning-parser", "qwen3", "--max_model_len", - "5000", + "6000", "--structured-outputs-config.backend", "xgrammar", "--enable-auto-tool-choice", diff --git a/tests/entrypoints/pooling/basic/test_tiling_engine.py b/tests/entrypoints/pooling/basic/test_tiling_engine.py new file mode 100644 index 000000000000..329f91debb1e --- /dev/null +++ b/tests/entrypoints/pooling/basic/test_tiling_engine.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import weakref +from unittest import mock + +import pytest + +from vllm import LLM, PoolingParams +from vllm.distributed import cleanup_dist_env_and_memory + +MODEL_NAME = "intfloat/multilingual-e5-small" + + +@pytest.fixture(scope="module") +def llm(): + llm = LLM( + model=MODEL_NAME, + max_num_seqs=2, # small to trigger tiling + tensor_parallel_size=1, + gpu_memory_utilization=0.75, + enforce_eager=True, + seed=0, + ) + + yield weakref.proxy(llm) + + del llm + + cleanup_dist_env_and_memory() + + +@pytest.mark.skip_global_cleanup +def test_tiling_engine_basic(llm): + """ + Basic test with a small number of prompts (less than max_num_seqs). + No tiling should be triggered, but the engine still processes correctly. + """ + prompts = ["Hello", "World"] + outputs = llm.encode(prompts, pooling_task="embed") + assert len(outputs) == len(prompts) + + +@pytest.mark.skip_global_cleanup +def test_tiling_engine_many_requests(llm): + """ + Test with a large number of prompts that exceeds max_num_seqs. + This verifies that _run_tiling_engine correctly chunks requests, + processes all of them, and returns outputs in the correct order. + """ + num_prompts = 10 + prompts = [f"Prompt {i}" for i in range(num_prompts)] + outputs = llm.encode(prompts, pooling_task="embed") + assert len(outputs) == num_prompts + + +@pytest.mark.skip_global_cleanup +def test_tiling_engine_with_pooling_params(llm): + """ + Test the tiling engine when different PoolingParams are provided. + The engine must handle a list of params that matches the number of prompts. + """ + num_prompts = 10 + prompts = [f"Prompt {i}" for i in range(num_prompts)] + pooling_params = [PoolingParams() for _ in range(num_prompts)] + + outputs = llm.encode(prompts, pooling_params=pooling_params, pooling_task="embed") + assert len(outputs) == num_prompts + + # Single PoolingParams shared across all prompts + single_param = PoolingParams() + outputs = llm.encode(prompts, pooling_params=single_param, pooling_task="embed") + assert len(outputs) == num_prompts + + # None PoolingParams should fall back to default + outputs = llm.encode(prompts, pooling_params=None, pooling_task="embed") + assert len(outputs) == num_prompts + + +@pytest.mark.skip_global_cleanup +def test_tiling_engine_abort_on_exception(llm): + """ + Test that abort_request IS called with the correct arguments when an + exception occurs inside the engine's step() loop. + """ + prompts = ["Prompt 0", "Prompt 1", "Prompt 2"] + + # Mock the step method to throw an exception on the second call + original_step = llm.llm_engine.step + call_count = 0 + + def mocked_step(): + nonlocal call_count + call_count += 1 + if call_count == 2: + raise RuntimeError("Simulated engine error") + return original_step() + + with mock.patch.object(llm.llm_engine, "step", side_effect=mocked_step): + # We expect an exception to be raised from encode + with mock.patch.object(llm.llm_engine, "abort_request") as mock_abort: # noqa: SIM117 + with pytest.raises(RuntimeError, match="Simulated engine error"): + llm.encode(prompts, pooling_task="embed") + + args, kwargs = mock_abort.call_args + request_ids = args[0] + assert isinstance(request_ids, list) + assert len(request_ids) > 0 diff --git a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py index c6663dbdff00..8c1d75cd762c 100644 --- a/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py +++ b/tests/entrypoints/pooling/scoring/test_cross_encoder_online_vision.py @@ -47,6 +47,7 @@ # TRITON_ATTN: gfx942/ROCm 7.2 drifts ~0.008 abs on text-vs-text (~7.9% rel). BACKEND_ABS_TOL: dict[str, float] = { "default": 0.0, + "auto": 0.007, "ROCM_AITER_FA": 0.005, "TRITON_ATTN": 0.009, "FLEX_ATTENTION": 0.006, diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py new file mode 100644 index 000000000000..8e51ff5d42b3 --- /dev/null +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Round trip parity CI: render -> generate -> derender -> render == render. + +Pins the coupled parsing path (``/v1/chat/completions`` on a normal GPU +server which parses generated tokens incrementally as they stream out) +against the disaggregated path (``OnlineDerenderer.derender_chat`` via +``/v1/chat/completions/derender`` which parses the same tokens all at once, +out of process). A standard ``vllm serve`` GPU server mounts both, so one +real generation lets both parsers run on identical input. + +Both paths consume the same generated token IDs (extracted from the coupled +response's ``token_ids`` via ``return_token_ids=True``), so generation +nondeterminism is irrelevant. The only variable under test is whether the +two parsing code paths agree. Parity is asserted unconditionally as only the +stronger per case assertions (e.g. "a tool call was produced") are gated +behind the marker actually having been emitted since a 1.5B model is not +guaranteed to emit ```` / ````. +""" + +import json + +import httpx +import pytest +import pytest_asyncio + +from tests.utils import RemoteOpenAIServer + +MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" +ARGS = [ + "--enable-auto-tool-choice", + "--tool-call-parser", + "hermes", + "--reasoning-parser", + "deepseek_r1", +] + +TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + }, + } +] +FORCE_WEATHER_TOOL = {"type": "function", "function": {"name": "get_weather"}} + + +@pytest.fixture(scope="module") +def server(): + with RemoteOpenAIServer(MODEL, ARGS) as remote_server: + yield remote_server + + +@pytest_asyncio.fixture +async def client(server): + async with httpx.AsyncClient( + base_url=server.url_for(""), timeout=60.0 + ) as http_client: + yield http_client + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +async def _coupled(client: httpx.AsyncClient, messages: list[dict], **extra) -> dict: + resp = await client.post( + "/v1/chat/completions", + json={ + "model": MODEL, + "messages": messages, + "temperature": 0, + "max_tokens": 128, + "return_token_ids": True, + **extra, + }, + ) + assert resp.status_code == 200, resp.text + return resp.json() + + +async def _disagg( + client: httpx.AsyncClient, + output_ids: list[int], + prompt_tokens: int, + finish_reason: str, + chat_request: dict, + logprobs: dict | None = None, +) -> dict: + resp = await client.post( + "/v1/chat/completions/derender", + json={ + "model": MODEL, + "generate_response": { + "request_id": "parity", + "choices": [ + { + "index": 0, + "token_ids": output_ids, + "finish_reason": finish_reason, + "logprobs": logprobs, + } + ], + }, + "prompt_tokens": prompt_tokens, + "chat_request": chat_request, + }, + ) + assert resp.status_code == 200, resp.text + return resp.json() + + +def _tool_sig(response_choice: dict) -> list[tuple[str, dict]]: + """[(name, json normalized args)] so key ordering / whitespace don't + cause false negatives.""" + return [ + (tc["function"]["name"], json.loads(tc["function"]["arguments"])) + for tc in (response_choice["message"].get("tool_calls") or []) + ] + + +def _assert_parity(coupled: dict, disagg: dict) -> None: + """Both paths saw the same tokens, so they must agree unconditionally.""" + c, d = coupled["choices"][0], disagg["choices"][0] + assert d["message"]["content"] == c["message"]["content"] + assert d["message"].get("reasoning") == c["message"].get("reasoning") + assert _tool_sig(d) == _tool_sig(c) + assert d["finish_reason"] == c["finish_reason"] + assert disagg["usage"]["prompt_tokens"] == coupled["usage"]["prompt_tokens"] + assert disagg["usage"]["completion_tokens"] == len(c["token_ids"]) + + +async def _run_parity_case( + client: httpx.AsyncClient, messages: list[dict], **extra +) -> tuple[dict, dict]: + """Run the coupled request then feed its generated tokens into the + disaggregated derender endpoint. Returns (coupled, disagg).""" + coupled = await _coupled(client, messages, **extra) + ch = coupled["choices"][0] + chat_request = {"model": MODEL, "messages": messages, **extra} + disagg = await _disagg( + client, + ch["token_ids"], + coupled["usage"]["prompt_tokens"], + ch["finish_reason"], + chat_request, + ) + return coupled, disagg + + +# --------------------------------------------------------------------------- +# Parity cases +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_parity_plain(client): + """Plain detokenization parity. No reasoning/tool markers involved.""" + messages = [ + {"role": "user", "content": "What is 2+2? Answer in one short sentence."} + ] + coupled, disagg = await _run_parity_case(client, messages) + _assert_parity(coupled, disagg) + + +@pytest.mark.asyncio +async def test_parity_reasoning(client): + """Reasoning/content split parity for ... outputs.""" + messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}] + coupled, disagg = await _run_parity_case( + client, messages, include_reasoning=True, max_tokens=256 + ) + _assert_parity(coupled, disagg) + + if not coupled["choices"][0]["message"].get("reasoning"): + pytest.skip("Model did not emit a block") + assert disagg["choices"][0]["message"]["reasoning"] + + +@pytest.mark.asyncio +async def test_parity_tool_call(client): + """Tool call name+args parity.""" + messages = [{"role": "user", "content": "What's the weather in Paris?"}] + coupled, disagg = await _run_parity_case( + client, messages, tools=TOOLS, tool_choice=FORCE_WEATHER_TOOL, max_tokens=1024 + ) + _assert_parity(coupled, disagg) + + if not _tool_sig(coupled["choices"][0]): + pytest.skip("Model did not emit a tool call") + assert _tool_sig(disagg["choices"][0]) + + +@pytest.mark.asyncio +async def test_parity_reasoning_and_tool_call(client): + """Combined reasoning + tool call parity means the highest drift risk + since it exercises both parser branches on the same output.""" + messages = [{"role": "user", "content": "What's the weather in Paris?"}] + coupled, disagg = await _run_parity_case( + client, + messages, + tools=TOOLS, + tool_choice=FORCE_WEATHER_TOOL, + include_reasoning=True, + max_tokens=1024, + ) + _assert_parity(coupled, disagg) + + c_msg = coupled["choices"][0]["message"] + if not (c_msg.get("reasoning") and _tool_sig(coupled["choices"][0])): + pytest.skip("Model did not emit both a block and a tool call") + d_msg = disagg["choices"][0]["message"] + assert d_msg["reasoning"] + assert _tool_sig(disagg["choices"][0]) + + +@pytest.mark.asyncio +async def test_parity_logprobs(client): + """token_id:N resolution parity vs. the coupled server's real strings. + + A real disaggregated worker only has token IDs so it emits logprobs + with ``token_id:N`` placeholders (``return_tokens_as_token_ids=True`` + reproduces that shape here). ``/derender`` must resolve those + placeholders to the same token strings/bytes the coupled server + resolves them to directly. + """ + messages = [{"role": "user", "content": "What is 2+2?"}] + extra = {"logprobs": True, "top_logprobs": 3} + + # What a real GPU less worker would hand to /derender is token IDs plus + # logprobs still in token_id:N placeholder form + placeholder = await _coupled( + client, messages, return_tokens_as_token_ids=True, **extra + ) + ch = placeholder["choices"][0] + chat_request = {"model": MODEL, "messages": messages, **extra} + disagg = await _disagg( + client, + ch["token_ids"], + placeholder["usage"]["prompt_tokens"], + ch["finish_reason"], + chat_request, + logprobs=ch["logprobs"], + ) + + # The coupled server resolving the same greedy generation + # to real token strings itself. + resolved = await _coupled(client, messages, **extra) + assert resolved["choices"][0]["token_ids"] == ch["token_ids"], ( + "greedy (temperature=0) generation was expected to be deterministic " + "across the two coupled calls used to build this test's fixtures" + ) + _assert_parity(resolved, disagg) + + r_content = resolved["choices"][0]["logprobs"]["content"] + d_content = disagg["choices"][0]["logprobs"]["content"] + assert len(d_content) == len(r_content) + for d_entry, r_entry in zip(d_content, r_content): + assert d_entry["token"] == r_entry["token"] + assert d_entry["bytes"] == r_entry["bytes"] diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index 45f2a13fdd74..9c47826850bd 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -283,28 +283,36 @@ def evaluate_gsm8k_offline( max_tokens: int = 256, temperature: float = 0.0, gen_prefix: str = "", + use_chat_completions: bool = False, ) -> dict[str, float | int]: """Evaluate GSM8K accuracy using an offline vllm.LLM object. Same prompts and scoring as evaluate_gsm8k(), but runs generation directly via llm.generate() instead of calling a server over HTTP. + + When ``use_chat_completions=True``, prompts go through the chat template via + ``llm.chat()`` instead of raw completion (for instruction-tuned models). """ from vllm import SamplingParams prompts, labels = _build_gsm8k_prompts(num_questions, num_shots, gen_prefix) - sampling_params = SamplingParams( temperature=temperature, max_tokens=max_tokens, stop=["Question", "Assistant:", "<|separator|>"], ) - + mode = "chat" if use_chat_completions else "completion" print( - f"Running offline GSM8K evaluation: {len(prompts)} questions, {num_shots}-shot" + f"Running offline GSM8K evaluation: {len(prompts)} questions, " + f"{num_shots}-shot, {mode}" ) tic = time.perf_counter() - outputs = llm.generate(prompts, sampling_params) + if use_chat_completions: + conversations = [[{"role": "user", "content": p}] for p in prompts] + outputs = llm.chat(conversations, sampling_params) + else: + outputs = llm.generate(prompts, sampling_params) latency = time.perf_counter() - tic states = [o.outputs[0].text for o in outputs] diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 83eeab5ed0f2..8a2bd5892a2d 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -185,6 +185,43 @@ def _assert_topk_indices_equal_unordered( assert set(actual_row) == set(expected_row) +def _reference_decode_index_score( + idx_q: torch.Tensor, + index_kv_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + score_block_stride: int, +) -> torch.Tensor: + total_q, num_idx_heads, _ = idx_q.shape + out = torch.full( + (num_idx_heads, total_q, score_block_stride), + -float("inf"), + device=idx_q.device, + dtype=torch.float32, + ) + for req_id, seq_len in enumerate(seq_lens.tolist()): + num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE + token_start = req_id * decode_query_len + q = idx_q[token_start : token_start + decode_query_len].float() + pages = block_table[req_id, :num_blocks] + k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1).float() + score = torch.einsum("qhd,kd->hqk", q, k) + q_pos = ( + seq_len + - decode_query_len + + torch.arange(decode_query_len, device=idx_q.device) + ) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + score.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, token_start : token_start + decode_query_len, :num_blocks] = ( + score.reshape(num_idx_heads, decode_query_len, num_blocks, BLOCK_SIZE) + .max(dim=3) + .values + ) + return out + + def test_prefill_index_topk_correctness(): topk = 6 init_blocks = 0 @@ -383,8 +420,8 @@ def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype): _assert_topk_indices_equal_unordered(actual, expected) -# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score + -# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata +# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha/CuteDSL score +# + unified top-k) and MiniMaxM3IndexerTritonImpl through their real metadata # builders on the SAME CommonAttentionMetadata + index cache, and assert the # selected blocks agree. This exercises all the metadata the impl/kernels consume # (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather, @@ -395,7 +432,8 @@ def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype): reason="fmha_sm100 indexer requires SM100 (Blackwell).", ) @pytest.mark.parametrize("topk", [8, 16]) -def test_msa_indexer_impl_matches_triton(topk, monkeypatch): +@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_msa_indexer_impl_matches_triton(topk, index_dtype, monkeypatch): import vllm.models.minimax_m3.common.indexer as indexer_mod from tests.v1.attention.utils import ( BatchSpec, @@ -439,13 +477,13 @@ def test_msa_indexer_impl_matches_triton(topk, monkeypatch): block_table = common.block_table_tensor num_pages = int(block_table.max().item()) + 1 index_cache = torch.zeros( - num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE + num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype ) for r, seq_len in enumerate(batch.seq_lens): for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE): index_cache[block_table[r, b]] = float(b + 1) index_q = torch.ones( - num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE + num_tokens, num_idx_heads * head_dim, device=device, dtype=index_dtype ) spec = MLAAttentionSpec( @@ -588,9 +626,8 @@ def test_decode_index_topk_correctness( ) @pytest.mark.parametrize("num_idx_heads", [1, 4]) def test_decode_index_topk_fp8(num_idx_heads: int): - """The fp8 (e4m3) indexer cache feeds the Triton decode kernel on the MSA - path. The kernel must score in fp32 (no scaling) so its top-k matches a - reference computed from the dequantized fp8 values.""" + """The standalone Triton path must score FP8 inputs in FP32 so its top-k + matches a reference computed from the dequantized FP8 values.""" torch.manual_seed(0) topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128 decode_query_len = 1 @@ -641,6 +678,79 @@ def test_decode_index_topk_fp8(num_idx_heads: int): _assert_topk_indices_equal_unordered(actual, expected) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="CuteDSL index decode score requires Blackwell.", +) +@pytest.mark.parametrize( + ("dtype", "decode_query_len", "max_decode_query_len"), + [ + (torch.bfloat16, 1, 1), + (torch.bfloat16, 3, 8), + (torch.float8_e4m3fn, 1, 1), + (torch.float8_e4m3fn, 3, 8), + (torch.float8_e4m3fn, 8, 8), + ], +) +def test_decode_index_score_cutedsl_correctness( + dtype: torch.dtype, + decode_query_len: int, + max_decode_query_len: int, +): + pytest.importorskip("cutlass") + from vllm.models.minimax_m3.nvidia.ops import ( + minimax_m3_index_decode_score_cutedsl, + ) + + torch.manual_seed(0) + init_blocks, local_blocks = 0, 0 + num_idx_heads, head_dim = 4, 128 + active_seq_lens = torch.tensor((1025, 4097), device="cuda", dtype=torch.int32) + batch = active_seq_lens.numel() + total_q = batch * decode_query_len + max_seq_len = int(active_seq_lens.max()) + max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE + score_block_stride = ((max_blocks + 15) // 16) * 16 + num_pages = batch * max_blocks + block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape( + batch, max_blocks + ) + idx_q = torch.randn(total_q, num_idx_heads, head_dim, device="cuda").to(dtype) + index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to( + dtype + ) + unified_score = torch.full( + (total_q, num_idx_heads, score_block_stride), + -float("inf"), + device="cuda", + dtype=torch.float32, + ) + score = unified_score.transpose(0, 1) + + minimax_m3_index_decode_score_cutedsl( + idx_q, + index_kv_cache, + block_table, + active_seq_lens, + max_seq_len=max_seq_len, + init_blocks=init_blocks, + local_blocks=local_blocks, + num_kv_heads=num_idx_heads, + decode_query_len=decode_query_len, + max_decode_query_len=max_decode_query_len, + score_out=score, + ) + expected = _reference_decode_index_score( + idx_q, + index_kv_cache, + block_table, + active_seq_lens, + decode_query_len, + score_block_stride, + ) + torch.testing.assert_close(score, expected) + + # Sparse attention kernels. def _reference_sparse_attn( q: torch.Tensor, diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 41863c916318..77e068ab1717 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -11,22 +11,22 @@ ) -def _on_gfx950() -> bool: +def _on_split_decode_arch() -> bool: if not current_platform.is_rocm(): return False try: - from vllm.platforms.rocm import _ON_GFX950 + from vllm.platforms.rocm import _ON_GFX942, _ON_GFX950 - return bool(_ON_GFX950) + return bool(_ON_GFX942 or _ON_GFX950) except Exception: return False -# The flash-decode split-K decode path is only tuned for AMD gfx950; other +# The flash-decode split-K decode path is only tuned for AMD gfx942/gfx950; other # architectures take the fallback decode kernel, so its tests are skipped there. -requires_gfx950 = pytest.mark.skipif( - not _on_gfx950(), - reason="split-K decode kernel is only tuned for AMD gfx950", +requires_split_decode_arch = pytest.mark.skipif( + not _on_split_decode_arch(), + reason="split-K decode kernel is only tuned for AMD gfx942/gfx950", ) NOPE_HEAD_DIM = 448 @@ -91,9 +91,13 @@ def _ref_sparse_prefill_ragged( def _pack_fp8_ds_mla_cache( - kv: torch.Tensor, block_size: int, is_extra: bool = False + kv: torch.Tensor, block_size: int, use_fnuz: bool ) -> torch.Tensor: assert kv.shape[-1] == HEAD_DIM + from vllm.models.deepseek_v4.common.ops.cache_utils import ( + quantize_and_insert_k_cache, + ) + num_tokens = kv.shape[0] num_blocks = (num_tokens + block_size - 1) // block_size cache = torch.zeros( @@ -101,41 +105,34 @@ def _pack_fp8_ds_mla_cache( dtype=torch.uint8, device=kv.device, ) - cache_flat = cache.view(torch.uint8).flatten() - kv_nope_fp8 = ( - kv[:, :NOPE_HEAD_DIM] - .to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()) - .view(torch.uint8) - ) - kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8) - - for slot in range(num_tokens): - block_idx = slot // block_size - pos = slot % block_size - block_base = block_idx * cache.stride(0) - token_base = block_base + pos * 576 - scale_base = block_base + block_size * 576 + pos * 8 - cache_flat[token_base : token_base + NOPE_HEAD_DIM].copy_(kv_nope_fp8[slot]) - cache_flat[ - token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2 - ].copy_(kv_rope_u8[slot]) - cache_flat[scale_base : scale_base + 7].fill_(127) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=kv.device) + quantize_and_insert_k_cache( + kv, + cache, + slot_mapping, + block_size=block_size, + use_fnuz=use_fnuz, + ) return cache def _read_fp8_ds_mla_cache( - cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False + cache: torch.Tensor, slot: int, block_size: int, use_fnuz: bool ) -> torch.Tensor: cache_flat = cache.view(torch.uint8).flatten() block_idx = slot // block_size pos = slot % block_size block_base = block_idx * cache.stride(0) token_base = block_base + pos * 576 + scale_base = block_base + block_size * 576 + pos * 8 + fp8_dtype = torch.float8_e4m3fnuz if use_fnuz else torch.float8_e4m3fn nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM] - nope = nope_u8.view( - torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype() - ).to(torch.float32) + nope = nope_u8.view(fp8_dtype).to(torch.float32) + scales = torch.exp2( + cache_flat[scale_base : scale_base + 7].to(torch.float32) - 127.0 + ) + nope = nope * scales.repeat_interleave(64) rope_u8 = cache_flat[ token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2 ] @@ -152,19 +149,21 @@ def _ref_sparse_decode_ragged( block_size: int, extra_cache: torch.Tensor | None = None, extra_rows: list[list[int]] | None = None, + main_use_fnuz: bool = False, + extra_use_fnuz: bool = False, ) -> torch.Tensor: q_f32 = q.float() out = torch.empty_like(q_f32) for query_idx in range(q.shape[0]): row_kv = [ - _read_fp8_ds_mla_cache(main_cache, int(slot), block_size) + _read_fp8_ds_mla_cache(main_cache, int(slot), block_size, main_use_fnuz) for slot in main_rows[query_idx] ] if extra_cache is not None and extra_rows is not None: row_kv.extend( _read_fp8_ds_mla_cache( - extra_cache, int(slot), block_size, is_extra=True + extra_cache, int(slot), block_size, extra_use_fnuz ) for slot in extra_rows[query_idx] ) @@ -290,11 +289,12 @@ def test_sparse_attn_decode_ragged_kernel() -> None: device = torch.device("cuda") torch.manual_seed(1) block_size = 4 + main_use_fnuz = current_platform.is_fp8_fnuz() q = torch.randn(2, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 - main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size) - extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True) + main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size, use_fnuz=main_use_fnuz) + extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, use_fnuz=False) main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device) main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device) extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device) @@ -324,12 +324,13 @@ def test_sparse_attn_decode_ragged_kernel() -> None: block_size=block_size, extra_cache=extra_cache, extra_rows=[[1], [3, 0]], + main_use_fnuz=main_use_fnuz, ) torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) -@requires_gfx950 +@requires_split_decode_arch @torch.inference_mode() def test_decode_num_splits_heuristic(monkeypatch) -> None: """Split-count heuristic added with the flash-decode split-K decode path.""" @@ -353,7 +354,7 @@ def test_decode_num_splits_heuristic(monkeypatch) -> None: assert mod._decode_num_splits(2, 1, avg_main_len=0.0, avg_extra_len=0.0) >= 1 -@requires_gfx950 +@requires_split_decode_arch @pytest.mark.parametrize("num_splits", [1, 2, 3, 4, 8]) @pytest.mark.parametrize("with_extra", [True, False]) @pytest.mark.parametrize("with_sink", [True, False]) @@ -363,8 +364,8 @@ def test_sparse_attn_decode_split_k_kernel( ) -> None: """Flash-decode split-K decode path (partial + reduce kernels). - This path is the gfx950 production path (``_ON_GFX950``), so the test only - runs on gfx950. The split count is pinned so the partial/reduce kernels are + This path is the gfx942/gfx950 production path, so the test only runs on + those architectures. The split count is pinned so the partial/reduce kernels are exercised across split counts. ``num_splits=8`` drives splits past the shortest segment length, covering the empty-split edge case handled by the reduce kernel. @@ -375,6 +376,7 @@ def test_sparse_attn_decode_split_k_kernel( torch.manual_seed(7) block_size = 4 num_heads = 3 + main_use_fnuz = current_platform.is_fp8_fnuz() main_rows = [[0, 2, 4, 6, 1, 3, 7, 5], [4, 1, 6, 0, 2]] num_queries = len(main_rows) @@ -385,7 +387,7 @@ def test_sparse_attn_decode_split_k_kernel( * 0.125 ) main_kv = torch.randn(8, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 - main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size) + main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size, use_fnuz=main_use_fnuz) main_indices, main_indptr = _ragged_from_rows(main_rows, device) extra_rows: list[list[int]] | None = None @@ -396,7 +398,7 @@ def test_sparse_attn_decode_split_k_kernel( rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]] extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125 extra_rows = rows - extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True) + extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, use_fnuz=False) extra_indices, extra_indptr = _ragged_from_rows(rows, device) attn_sink = ( @@ -431,6 +433,7 @@ def test_sparse_attn_decode_split_k_kernel( block_size=block_size, extra_cache=extra_cache, extra_rows=extra_rows, + main_use_fnuz=main_use_fnuz, ) torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) diff --git a/tests/kernels/core/test_fused_rms_norm_gated.py b/tests/kernels/core/test_fused_rms_norm_gated.py index 793dd02a9f5a..69788e37721c 100644 --- a/tests/kernels/core/test_fused_rms_norm_gated.py +++ b/tests/kernels/core/test_fused_rms_norm_gated.py @@ -7,7 +7,7 @@ import pytest import torch -from vllm.model_executor.layers.fla.ops.kda import FusedRMSNormGated +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.utils.torch_utils import set_random_seed DTYPES = [torch.bfloat16] diff --git a/tests/kernels/mamba/test_gdn_forward_core_split.py b/tests/kernels/mamba/test_gdn_forward_core_split.py index f2bfc30abdb6..fa28f4701c61 100644 --- a/tests/kernels/mamba/test_gdn_forward_core_split.py +++ b/tests/kernels/mamba/test_gdn_forward_core_split.py @@ -53,11 +53,6 @@ create_vllm_config, ) from vllm.config import set_current_vllm_config # noqa: E402 -from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402 - prepare_chunk_indices, - prepare_chunk_offsets, -) -from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE # noqa: E402 from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: E402 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa: E402 ChunkGatedDeltaRule, @@ -66,6 +61,13 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( # noqa: E402 MambaStateShapeCalculator, ) +from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402 + prepare_chunk_indices, + prepare_chunk_offsets, +) +from vllm.third_party.flash_linear_attention.ops.utils import ( # noqa: E402 + FLA_CHUNK_SIZE, +) from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402 GDNAttentionMetadataBuilder, ) diff --git a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py index 8f371e79db47..1f5a24fd81ff 100644 --- a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py +++ b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py @@ -17,17 +17,17 @@ allow_module_level=True, ) -from vllm.model_executor.layers.fla.ops import ( # noqa: E402 +from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( # noqa: E402 + chunk_gated_delta_rule_cutedsl, + prepare_metadata_cutedsl, +) +from vllm.third_party.flash_linear_attention.ops import ( # noqa: E402 chunk_gated_delta_rule, ) -from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402 +from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402 prepare_chunk_indices, prepare_chunk_offsets, ) -from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( # noqa: E402 - chunk_gated_delta_rule_cutedsl, - prepare_metadata_cutedsl, -) @pytest.mark.parametrize("num_seqs", [1, 5, 257]) diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index 1b68213fafef..46ca934c1462 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -153,9 +153,9 @@ def test_fused_topk_softplus_sqrt_hash( # experts. hash_indices_table = torch.stack( [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)] - ).to(device="cuda", dtype=torch.int32) + ).to(device="cuda", dtype=torch.long) input_ids = torch.randint( - 0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda" + 0, vocab_size, (num_tokens,), dtype=torch.long, device="cuda" ) topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( diff --git a/tests/kernels/quantization/test_int4_emulation_moe.py b/tests/kernels/quantization/test_int4_emulation_moe.py new file mode 100644 index 000000000000..7016860d86f2 --- /dev/null +++ b/tests/kernels/quantization/test_int4_emulation_moe.py @@ -0,0 +1,1189 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for Int4EmulationTritonExperts MoE backend. + +Tests the weight dequantization helpers (_unpack_and_dequant_int4_gptq, +_unpack_and_dequant_int4_awq) and full MoE forward pass +(_process_weights_emulation_gptq, _process_weights_emulation_awq) +for both symmetric and asymmetric zero-point cases. + +Run `pytest tests/kernels/quantization/test_int4_emulation_moe.py`. +""" + +import numpy +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + int4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import ( + Int4EmulationTritonExperts, +) +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + _process_weights_emulation_awq, + _process_weights_emulation_gptq, + _unpack_and_dequant_int4_awq, + _unpack_and_dequant_int4_gptq, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + awq_pack, + gptq_pack, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="Int4EmulationTritonExperts requires CUDA.", +) + +device = "cuda" + +# (E, K, N, group_size) +SHAPES = [ + pytest.param(2, 64, 32, 32, id="tiny-gs32"), + pytest.param(4, 128, 64, 64, id="small-gs64"), + pytest.param(4, 256, 128, 128, id="medium-gs128"), +] + +# (E, K, N, top_k, group_size, num_tokens) +E2E_CONFIGS = [ + pytest.param(4, 64, 32, 2, 32, 8, id="tiny"), + pytest.param(8, 128, 64, 2, 64, 16, id="small"), +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _quantize_sym(w_fp: torch.Tensor, group_size: int): + """Quantize [K, N] float to int4 symmetric (uint4b8), return q and scale.""" + K, N = w_fp.shape + assert K % group_size == 0 + n_groups = K // group_size + w_grouped = w_fp.reshape(n_groups, group_size, N) + scale = w_grouped.abs().amax(dim=1) / 7.0 + scale = scale.clamp(min=1e-6) + w_quant = (w_grouped / scale.unsqueeze(1)).round().clamp(-8, 7) + q = (w_quant + 8).to(torch.int32).reshape(K, N) + return q, scale + + +def _quantize_asym(w_fp: torch.Tensor, group_size: int): + """Quantize [K, N] float to uint4 asymmetric, return q, scale, zero.""" + K, N = w_fp.shape + assert K % group_size == 0 + n_groups = K // group_size + w_grouped = w_fp.reshape(n_groups, group_size, N) + wmin = w_grouped.amin(dim=1) + wmax = w_grouped.amax(dim=1) + scale = (wmax - wmin) / 15.0 + scale = scale.clamp(min=1e-6) + zero = (-wmin / scale).round().clamp(0, 15).to(torch.int32) + w_quant = ((w_grouped - wmin.unsqueeze(1)) / scale.unsqueeze(1)).round() + q = w_quant.clamp(0, 15).to(torch.int32).reshape(K, N) + return q, scale, zero + + +def _dequantize_ref( + w_uint: torch.Tensor, + scale: torch.Tensor, + zero=None, + output_dtype: torch.dtype = torch.bfloat16, +): + """Reference dequant for a single [K, N] slice.""" + K, N = w_uint.shape + n_groups = scale.shape[0] + group_size = K // n_groups + w = w_uint.reshape(n_groups, group_size, N).to(output_dtype) + s = scale.unsqueeze(1).to(output_dtype) + if zero is None: + return ((w - 8) * s).reshape(K, N) + z = zero.unsqueeze(1).to(output_dtype) + return ((w - z) * s).reshape(K, N) + + +def _pack_gptq_zeros(zero: torch.Tensor, N: int) -> torch.Tensor: + """Pack [n_groups, N] zeros into GPTQ format [n_groups, N//8] int32.""" + n_groups, _ = zero.shape + z = zero.to(torch.int32).cpu().numpy().astype(numpy.uint32) + packed = numpy.zeros((n_groups, N // 8), dtype=numpy.uint32) + for i in range(8): + packed |= z[:, i::8] << (i * 4) + return torch.from_numpy(packed.astype(numpy.int32)).to(device) + + +def _pack_awq_zeros(zero: torch.Tensor, N: int) -> torch.Tensor: + """Pack [n_groups, N] zeros into AWQ column format [n_groups, N//8] int32.""" + n_groups, _ = zero.shape + interleave = numpy.array([0, 2, 4, 6, 1, 3, 5, 7]) + z = zero.to(torch.int32).cpu().numpy().astype(numpy.uint32) + z_interleaved = z.reshape(-1, 8)[:, interleave].reshape(n_groups, N) + packed = numpy.zeros((n_groups, N // 8), dtype=numpy.uint32) + for i in range(8): + packed |= z_interleaved[:, i::8] << (i * 4) + return torch.from_numpy(packed.astype(numpy.int32)).to(device) + + +def _make_moe_config(E, K, N): + return FusedMoEConfig( + num_experts=E, + experts_per_token=2, + hidden_dim=K, + intermediate_size=N, + num_local_experts=E, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + + +def _make_gptq_moe_weights(E, K, N, group_size, asym=False): + """Build GPTQ MoE weight tensors and per-expert float references.""" + torch.manual_seed(7) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) + + w13_list, w13s_list, w13z_list, w13_ref_list = [], [], [], [] + w2_list, w2s_list, w2z_list, w2_ref_list = [], [], [], [] + + for e in range(E): + if asym: + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size) + w13_list.append(gptq_pack(q13, 4, K, 2 * N)) + w13s_list.append(s13) + w13z_list.append(_pack_gptq_zeros(z13, 2 * N)) + w13_ref_list.append(_dequantize_ref(q13, s13, z13)) + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size) + w2_list.append(gptq_pack(q2, 4, N, K)) + w2s_list.append(s2) + w2z_list.append(_pack_gptq_zeros(z2, K)) + w2_ref_list.append(_dequantize_ref(q2, s2, z2)) + else: + q13, s13 = _quantize_sym(w13_fp[e], group_size) + w13_list.append(gptq_pack(q13, 4, K, 2 * N)) + w13s_list.append(s13) + w13z_list.append(None) + w13_ref_list.append(_dequantize_ref(q13, s13)) + q2, s2 = _quantize_sym(w2_fp[e], group_size) + w2_list.append(gptq_pack(q2, 4, N, K)) + w2s_list.append(s2) + w2z_list.append(None) + w2_ref_list.append(_dequantize_ref(q2, s2)) + + return ( + torch.stack(w13_list), + torch.stack(w13s_list), + torch.stack(w13z_list) if asym else None, + torch.stack(w13_ref_list), # [E, K, 2N] + torch.stack(w2_list), + torch.stack(w2s_list), + torch.stack(w2z_list) if asym else None, + torch.stack(w2_ref_list), # [E, N, K] + ) + + +def _make_awq_moe_weights(E, K, N, group_size, asym=False): + """Build AWQ MoE weight tensors and per-expert float references.""" + torch.manual_seed(8) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) + + w13_list, w13s_list, w13z_list, w13_ref_list = [], [], [], [] + w2_list, w2s_list, w2z_list, w2_ref_list = [], [], [], [] + + for e in range(E): + if asym: + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size) + w13_list.append(awq_pack(q13, 4, K, 2 * N)) + w13s_list.append(s13) + w13z_list.append(_pack_awq_zeros(z13, 2 * N)) + w13_ref_list.append(_dequantize_ref(q13, s13, z13)) + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size) + w2_list.append(awq_pack(q2, 4, N, K)) + w2s_list.append(s2) + w2z_list.append(_pack_awq_zeros(z2, K)) + w2_ref_list.append(_dequantize_ref(q2, s2, z2)) + else: + q13, s13 = _quantize_sym(w13_fp[e], group_size) + w13_list.append(awq_pack(q13, 4, K, 2 * N)) + w13s_list.append(s13) + w13z_list.append(None) + w13_ref_list.append(_dequantize_ref(q13, s13)) + q2, s2 = _quantize_sym(w2_fp[e], group_size) + w2_list.append(awq_pack(q2, 4, N, K)) + w2s_list.append(s2) + w2z_list.append(None) + w2_ref_list.append(_dequantize_ref(q2, s2)) + + return ( + torch.stack(w13_list), + torch.stack(w13s_list), + torch.stack(w13z_list) if asym else None, + torch.stack(w13_ref_list), # [E, K, 2N] + torch.stack(w2_list), + torch.stack(w2s_list), + torch.stack(w2z_list) if asym else None, + torch.stack(w2_ref_list), # [E, N, K] + ) + + +def _run_emulation_forward( + experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N +): + ws13_size = hidden_states.shape[0] * topk_ids.shape[1] * max(N, K) + ws2_size = hidden_states.shape[0] * topk_ids.shape[1] * max(2 * N, K) + workspace13 = torch.zeros(ws13_size, dtype=hidden_states.dtype, device=device) + workspace2 = torch.zeros(ws2_size, dtype=hidden_states.dtype, device=device) + output = torch.zeros( + hidden_states.shape[0], K, dtype=hidden_states.dtype, device=device + ) + experts.apply( + output=output, + hidden_states=hidden_states, + w1=w13_bf16, + w2=w2_bf16, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=E, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + return output + + +# --------------------------------------------------------------------------- +# Tests: _unpack_and_dequant_int4_gptq +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_gptq_unpack_symmetric(E, K, N, group_size): + """GPTQ symmetric unpacker matches reference.""" + torch.manual_seed(0) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list, ref_list = [], [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + packed_list.append(gptq_pack(q, 4, K, N)) + scale_list.append(s) + ref_list.append(_dequantize_ref(q, s, output_dtype=torch.float32)) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + ref = torch.stack(ref_list).to(device) + + out = _unpack_and_dequant_int4_gptq( + w_packed, scale, None, transpose_output=False, output_dtype=torch.float32 + ) + + assert out.shape == (E, K, N) + assert torch.allclose(out, ref, atol=0), ( + f"max diff: {(out - ref).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_gptq_unpack_asymmetric(E, K, N, group_size): + """GPTQ asymmetric unpacker matches reference.""" + torch.manual_seed(1) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list, zero_list, ref_list = [], [], [], [] + for e in range(E): + q, s, z = _quantize_asym(w_fp[e], group_size) + packed_list.append(gptq_pack(q, 4, K, N)) + scale_list.append(s) + zero_list.append(_pack_gptq_zeros(z, N)) + ref_list.append(_dequantize_ref(q, s, z, output_dtype=torch.float32)) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + qzeros = torch.stack(zero_list).to(device) + ref = torch.stack(ref_list).to(device) + + out = _unpack_and_dequant_int4_gptq( + w_packed, scale, qzeros, transpose_output=False, output_dtype=torch.float32 + ) + + assert out.shape == (E, K, N) + assert torch.allclose(out, ref, atol=0), ( + f"max diff: {(out - ref).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_gptq_unpack_transpose(E, K, N, group_size): + """GPTQ transpose_output=True gives [E, N, K].""" + torch.manual_seed(2) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list = [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + packed_list.append(gptq_pack(q, 4, K, N)) + scale_list.append(s) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + + out_normal = _unpack_and_dequant_int4_gptq(w_packed, scale, None, False) + out_transposed = _unpack_and_dequant_int4_gptq(w_packed, scale, None, True) + + assert out_transposed.shape == (E, N, K) + assert torch.allclose( + out_transposed, out_normal.permute(0, 2, 1).contiguous(), atol=0 + ) + + +# --------------------------------------------------------------------------- +# Tests: _unpack_and_dequant_int4_awq +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_awq_unpack_symmetric(E, K, N, group_size): + """AWQ symmetric unpacker matches reference.""" + torch.manual_seed(3) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list, ref_list = [], [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + packed_list.append(awq_pack(q, 4, K, N)) + scale_list.append(s) + ref_list.append(_dequantize_ref(q, s, output_dtype=torch.float32)) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + ref = torch.stack(ref_list).to(device) + + out = _unpack_and_dequant_int4_awq( + w_packed, scale, None, transpose_output=False, output_dtype=torch.float32 + ) + + assert out.shape == (E, K, N) + assert torch.allclose(out, ref, atol=0), ( + f"max diff: {(out - ref).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_awq_unpack_asymmetric(E, K, N, group_size): + """AWQ asymmetric unpacker matches reference.""" + torch.manual_seed(4) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list, zero_list, ref_list = [], [], [], [] + for e in range(E): + q, s, z = _quantize_asym(w_fp[e], group_size) + packed_list.append(awq_pack(q, 4, K, N)) + scale_list.append(s) + zero_list.append(_pack_awq_zeros(z, N)) + ref_list.append(_dequantize_ref(q, s, z, output_dtype=torch.float32)) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + qzeros = torch.stack(zero_list).to(device) + ref = torch.stack(ref_list).to(device) + + out = _unpack_and_dequant_int4_awq( + w_packed, scale, qzeros, transpose_output=False, output_dtype=torch.float32 + ) + + assert out.shape == (E, K, N) + assert torch.allclose(out, ref, atol=0), ( + f"max diff: {(out - ref).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_awq_unpack_transpose(E, K, N, group_size): + """AWQ transpose_output=True gives [E, N, K].""" + torch.manual_seed(5) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + packed_list, scale_list = [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + packed_list.append(awq_pack(q, 4, K, N)) + scale_list.append(s) + + w_packed = torch.stack(packed_list).to(device) + scale = torch.stack(scale_list).to(device) + + out_normal = _unpack_and_dequant_int4_awq(w_packed, scale, None, False) + out_transposed = _unpack_and_dequant_int4_awq(w_packed, scale, None, True) + + assert out_transposed.shape == (E, N, K) + assert torch.allclose( + out_transposed, out_normal.permute(0, 2, 1).contiguous(), atol=0 + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +def test_awq_gptq_unpack_agree(E, K, N, group_size): + """AWQ and GPTQ unpackers produce identical values for the same weights.""" + torch.manual_seed(6) + w_fp = torch.randn(E, K, N, dtype=torch.float16, device=device) + + gptq_list, awq_list, scale_list = [], [], [] + for e in range(E): + q, s = _quantize_sym(w_fp[e], group_size) + gptq_list.append(gptq_pack(q, 4, K, N)) + awq_list.append(awq_pack(q, 4, K, N)) + scale_list.append(s) + + scale = torch.stack(scale_list).to(device) + out_gptq = _unpack_and_dequant_int4_gptq( + torch.stack(gptq_list).to(device), scale, None, False, torch.float32 + ) + out_awq = _unpack_and_dequant_int4_awq( + torch.stack(awq_list).to(device), scale, None, False, torch.float32 + ) + + assert torch.allclose(out_gptq, out_awq, atol=0), ( + f"max diff: {(out_gptq - out_awq).abs().max().item()}" + ) + + +# --------------------------------------------------------------------------- +# Tests: _process_weights_emulation_{gptq,awq} +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"]) +def test_gptq_process_weights_shapes_and_values(E, K, N, group_size, asym): + """_process_weights_emulation_gptq shapes and values match reference.""" + w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_gptq_moe_weights( + E, K, N, group_size, asym + ) + result = _process_weights_emulation_gptq(w13, w2, w13s, w2s, w13z, w2z) + w13_out, w2_out = result[0], result[1] + + assert w13_out.shape == (E, 2 * N, K) + assert w2_out.shape == (E, K, N) + assert w13_out.dtype == torch.bfloat16 + assert w2_out.dtype == torch.bfloat16 + + expected_w13 = w13_ref.permute(0, 2, 1) + expected_w2 = w2_ref.permute(0, 2, 1) + + assert torch.allclose(w13_out.float(), expected_w13.float(), atol=0), ( + f"w13 max diff: {(w13_out.float() - expected_w13.float()).abs().max().item()}" + ) + assert torch.allclose(w2_out.float(), expected_w2.float(), atol=0), ( + f"w2 max diff: {(w2_out.float() - expected_w2.float()).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"]) +def test_awq_process_weights_shapes_and_values(E, K, N, group_size, asym): + """_process_weights_emulation_awq shapes and values match reference.""" + w13, w13s, w13z, w13_ref, w2, w2s, w2z, w2_ref = _make_awq_moe_weights( + E, K, N, group_size, asym + ) + result = _process_weights_emulation_awq(w13, w2, w13s, w2s, w13z, w2z) + w13_out, w2_out = result[0], result[1] + + assert w13_out.shape == (E, 2 * N, K) + assert w2_out.shape == (E, K, N) + assert w13_out.dtype == torch.bfloat16 + assert w2_out.dtype == torch.bfloat16 + + expected_w13 = w13_ref.permute(0, 2, 1) + expected_w2 = w2_ref.permute(0, 2, 1) + + assert torch.allclose(w13_out.float(), expected_w13.float(), atol=0), ( + f"w13 max diff: {(w13_out.float() - expected_w13.float()).abs().max().item()}" + ) + assert torch.allclose(w2_out.float(), expected_w2.float(), atol=0), ( + f"w2 max diff: {(w2_out.float() - expected_w2.float()).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, group_size", SHAPES) +@pytest.mark.parametrize("asym", [False, True], ids=["sym", "asym"]) +def test_gptq_awq_process_weights_agree(E, K, N, group_size, asym): + """AWQ and GPTQ process_weights produce identical dequantized tensors.""" + torch.manual_seed(9) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) + + g13_list, g13s_list, g13z_list = [], [], [] + a13_list, a13s_list, a13z_list = [], [], [] + g2_list, g2s_list, g2z_list = [], [], [] + a2_list, a2s_list, a2z_list = [], [], [] + + for e in range(E): + if asym: + q13, s13, z13 = _quantize_asym(w13_fp[e], group_size) + q2, s2, z2 = _quantize_asym(w2_fp[e], group_size) + g13z_list.append(_pack_gptq_zeros(z13, 2 * N)) + a13z_list.append(_pack_awq_zeros(z13, 2 * N)) + g2z_list.append(_pack_gptq_zeros(z2, K)) + a2z_list.append(_pack_awq_zeros(z2, K)) + else: + q13, s13 = _quantize_sym(w13_fp[e], group_size) + q2, s2 = _quantize_sym(w2_fp[e], group_size) + g13z_list.append(None) + a13z_list.append(None) + g2z_list.append(None) + a2z_list.append(None) + + g13_list.append(gptq_pack(q13, 4, K, 2 * N)) + a13_list.append(awq_pack(q13, 4, K, 2 * N)) + g13s_list.append(s13) + a13s_list.append(s13) + g2_list.append(gptq_pack(q2, 4, N, K)) + a2_list.append(awq_pack(q2, 4, N, K)) + g2s_list.append(s2) + a2s_list.append(s2) + + gptq_res = _process_weights_emulation_gptq( + torch.stack(g13_list), + torch.stack(g2_list), + torch.stack(g13s_list), + torch.stack(g2s_list), + torch.stack(g13z_list) if asym else None, + torch.stack(g2z_list) if asym else None, + ) + awq_res = _process_weights_emulation_awq( + torch.stack(a13_list), + torch.stack(a2_list), + torch.stack(a13s_list), + torch.stack(a2s_list), + torch.stack(a13z_list) if asym else None, + torch.stack(a2z_list) if asym else None, + ) + + assert torch.allclose(gptq_res[0].float(), awq_res[0].float(), atol=1e-3), ( + f"w13 max diff: {(gptq_res[0] - awq_res[0]).float().abs().max().item()}" + ) + assert torch.allclose(gptq_res[1].float(), awq_res[1].float(), atol=1e-3), ( + f"w2 max diff: {(gptq_res[1] - awq_res[1]).float().abs().max().item()}" + ) + + +# --------------------------------------------------------------------------- +# End-to-end MoE forward pass tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens", E2E_CONFIGS) +def test_gptq_vs_awq_forward_agree(E, K, N, top_k, group_size, num_tokens): + """GPTQ and AWQ emulation backends produce bit-identical forward outputs.""" + torch.manual_seed(42) + moe_config = _make_moe_config(E, K, N) + + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) + + g13_list, g13s_list, a13_list, a13s_list = [], [], [], [] + g2_list, g2s_list, a2_list, a2s_list = [], [], [], [] + + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e], group_size) + q2, s2 = _quantize_sym(w2_fp[e], group_size) + g13_list.append(gptq_pack(q13, 4, K, 2 * N)) + a13_list.append(awq_pack(q13, 4, K, 2 * N)) + g13s_list.append(s13) + a13s_list.append(s13.clone()) + g2_list.append(gptq_pack(q2, 4, N, K)) + a2_list.append(awq_pack(q2, 4, N, K)) + g2s_list.append(s2) + a2s_list.append(s2.clone()) + + gptq_res = _process_weights_emulation_gptq( + torch.stack(g13_list), + torch.stack(g2_list), + torch.stack(g13s_list), + torch.stack(g2s_list), + None, + None, + ) + awq_res = _process_weights_emulation_awq( + torch.stack(a13_list), + torch.stack(a2_list), + torch.stack(a13s_list), + torch.stack(a2s_list), + None, + None, + ) + w13_gptq, w2_gptq = gptq_res[0], gptq_res[1] + w13_awq, w2_awq = awq_res[0], awq_res[1] + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + experts_gptq = Int4EmulationTritonExperts( + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + experts_awq = Int4EmulationTritonExperts( + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + topk_ids = torch.stack( + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + out_gptq = _run_emulation_forward( + experts_gptq, w13_gptq, w2_gptq, hidden_states, topk_weights, topk_ids, E, K, N + ) + out_awq = _run_emulation_forward( + experts_awq, w13_awq, w2_awq, hidden_states, topk_weights, topk_ids, E, K, N + ) + + assert torch.allclose(out_gptq, out_awq, atol=0), ( + f"max diff: {(out_gptq - out_awq).abs().max().item()}" + ) + + +# --------------------------------------------------------------------------- +# EP (Expert Parallelism) tests +# --------------------------------------------------------------------------- + +# (E, K, N, top_k, group_size, num_tokens, ep_size) +EP_CONFIGS = [ + pytest.param(4, 64, 32, 2, 32, 8, 2, id="E4-ep2"), + pytest.param(8, 64, 32, 2, 32, 16, 4, id="E8-ep4"), + pytest.param(8, 128, 64, 2, 64, 16, 2, id="E8-ep2"), +] + + +def _make_expert_map(global_num_experts: int, start: int, end: int) -> torch.Tensor: + """Build expert_map for a rank that owns experts [start, end).""" + expert_map = torch.full((global_num_experts,), -1, dtype=torch.int32, device=device) + expert_map[start:end] = torch.arange(end - start, dtype=torch.int32, device=device) + return expert_map + + +def _run_emulation_forward_ep( + experts, + w13_bf16, + w2_bf16, + hidden_states, + topk_weights, + topk_ids, + global_num_experts, + expert_map, +): + """Run forward with EP expert_map; returns output tensor.""" + T, K = hidden_states.shape + N = w2_bf16.shape[2] + ws13_size = T * topk_ids.shape[1] * max(N, K) + ws2_size = T * topk_ids.shape[1] * max(2 * N, K) + workspace13 = torch.zeros(ws13_size, dtype=hidden_states.dtype, device=device) + workspace2 = torch.zeros(ws2_size, dtype=hidden_states.dtype, device=device) + output = torch.zeros(T, K, dtype=hidden_states.dtype, device=device) + experts.apply( + output=output, + hidden_states=hidden_states, + w1=w13_bf16, + w2=w2_bf16, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=global_num_experts, + expert_map=expert_map, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + return output + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS) +@pytest.mark.parametrize("fmt", ["gptq", "awq"]) +def test_ep_output_matches_no_ep(E, K, N, top_k, group_size, num_tokens, ep_size, fmt): + """EP simulation: sum of per-rank outputs equals the no-EP forward pass.""" + assert E % ep_size == 0 + num_local = E // ep_size + + torch.manual_seed(20) + + # Build all expert weights in BF16 (no-EP reference) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02 + + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + if fmt == "gptq": + packed13_list.append(gptq_pack(q13, 4, K, 2 * N)) + packed2_list.append(gptq_pack(q2, 4, N, K)) + else: + packed13_list.append(awq_pack(q13, 4, K, 2 * N)) + packed2_list.append(awq_pack(q2, 4, N, K)) + scales13_list.append(s13) + scales2_list.append(s2) + + process_fn = ( + _process_weights_emulation_gptq + if fmt == "gptq" + else _process_weights_emulation_awq + ) + res = process_fn( + torch.stack(packed13_list), + torch.stack(packed2_list), + torch.stack(scales13_list), + torch.stack(scales2_list), + None, + None, + ) + w13_all, w2_all = res[0], res[1] # [E, 2N, K], [E, K, N] + + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + topk_ids = torch.stack( + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + # No-EP reference + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + moe_config_full = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=E, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_ref = Int4EmulationTritonExperts( + moe_config_full, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_no_ep = _run_emulation_forward( + experts_ref, w13_all, w2_all, hidden_states, topk_weights, topk_ids, E, K, N + ) + + # EP simulation: sum contributions from each rank + out_ep_sum = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device) + for rank in range(ep_size): + start = rank * num_local + end = start + num_local + expert_map = _make_expert_map(E, start, end) + moe_config_ep = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=num_local, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_ep = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_rank = _run_emulation_forward_ep( + experts_ep, + w13_all[start:end], + w2_all[start:end], + hidden_states, + topk_weights, + topk_ids, + global_num_experts=E, + expert_map=expert_map, + ) + out_ep_sum = out_ep_sum + out_rank + + assert torch.allclose(out_ep_sum, out_no_ep, atol=1e-3), ( + f"[{fmt}] EP sum max diff: {(out_ep_sum - out_no_ep).abs().max().item():.6f}" + ) + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS) +def test_ep_gptq_awq_agree(E, K, N, top_k, group_size, num_tokens, ep_size): + """With EP, GPTQ and AWQ emulation produce the same outputs per rank.""" + assert E % ep_size == 0 + num_local = E // ep_size + + torch.manual_seed(21) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02 + + g13_list, g13s_list, g2_list, g2s_list = [], [], [], [] + a13_list, a13s_list, a2_list, a2s_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + g13_list.append(gptq_pack(q13, 4, K, 2 * N)) + a13_list.append(awq_pack(q13, 4, K, 2 * N)) + g13s_list.append(s13) + a13s_list.append(s13.clone()) + g2_list.append(gptq_pack(q2, 4, N, K)) + a2_list.append(awq_pack(q2, 4, N, K)) + g2s_list.append(s2) + a2s_list.append(s2.clone()) + + gptq_res = _process_weights_emulation_gptq( + torch.stack(g13_list), + torch.stack(g2_list), + torch.stack(g13s_list), + torch.stack(g2s_list), + None, + None, + ) + awq_res = _process_weights_emulation_awq( + torch.stack(a13_list), + torch.stack(a2_list), + torch.stack(a13s_list), + torch.stack(a2s_list), + None, + None, + ) + w13_gptq, w2_gptq = gptq_res[0], gptq_res[1] + w13_awq, w2_awq = awq_res[0], awq_res[1] + + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + topk_ids = torch.stack( + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + + for rank in range(ep_size): + start = rank * num_local + end = start + num_local + expert_map = _make_expert_map(E, start, end) + moe_config_ep = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=num_local, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_gptq = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + experts_awq = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_gptq = _run_emulation_forward_ep( + experts_gptq, + w13_gptq[start:end], + w2_gptq[start:end], + hidden_states, + topk_weights, + topk_ids, + E, + expert_map, + ) + out_awq = _run_emulation_forward_ep( + experts_awq, + w13_awq[start:end], + w2_awq[start:end], + hidden_states, + topk_weights, + topk_ids, + E, + expert_map, + ) + assert torch.allclose(out_gptq, out_awq, atol=0), ( + f"rank={rank} max diff: {(out_gptq - out_awq).abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS) +def test_ep_partial_rank_no_active_experts( + E, K, N, top_k, group_size, num_tokens, ep_size +): + """A rank that owns no token-selected experts produces an all-zero output.""" + assert E % ep_size == 0 + num_local = E // ep_size + + torch.manual_seed(22) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02 + + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + packed13_list.append(gptq_pack(q13, 4, K, 2 * N)) + packed2_list.append(gptq_pack(q2, 4, N, K)) + scales13_list.append(s13) + scales2_list.append(s2) + + res = _process_weights_emulation_gptq( + torch.stack(packed13_list), + torch.stack(packed2_list), + torch.stack(scales13_list), + torch.stack(scales2_list), + None, + None, + ) + w13_all, w2_all = res[0], res[1] + + # Force topk_ids to only use experts in [0, num_local) — rank 0's slice + topk_ids = torch.zeros(num_tokens, top_k, dtype=torch.int32, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + + # Last rank owns experts [E-num_local, E), tokens only route to [0, num_local) + last_rank = ep_size - 1 + start = last_rank * num_local + end = E + expert_map = _make_expert_map(E, start, end) + moe_config_ep = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=num_local, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_ep = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out = _run_emulation_forward_ep( + experts_ep, + w13_all[start:end], + w2_all[start:end], + hidden_states, + topk_weights, + topk_ids, + E, + expert_map, + ) + assert torch.all(out == 0), ( + f"Expected zeros for inactive rank, got max={out.abs().max().item()}" + ) + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens, ep_size", EP_CONFIGS) +def test_ep_sum_equals_full_forward(E, K, N, top_k, group_size, num_tokens, ep_size): + """With fixed routing, EP rank outputs sum to the single-rank full forward.""" + assert E % ep_size == 0 + num_local = E // ep_size + + torch.manual_seed(23) + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.float16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.float16, device=device) * 0.02 + + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + packed13_list.append(gptq_pack(q13, 4, K, 2 * N)) + packed2_list.append(gptq_pack(q2, 4, N, K)) + scales13_list.append(s13) + scales2_list.append(s2) + + res = _process_weights_emulation_gptq( + torch.stack(packed13_list), + torch.stack(packed2_list), + torch.stack(scales13_list), + torch.stack(scales2_list), + None, + None, + ) + w13_all, w2_all = res[0], res[1] + + # Fix routing so every token uses exactly 2 consecutive experts (round-robin) + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_ids = torch.stack( + [ + torch.tensor([(t * top_k + k) % E for k in range(top_k)], dtype=torch.int32) + for t in range(num_tokens) + ] + ).to(device) + topk_weights = torch.full((num_tokens, top_k), 1.0 / top_k, device=device) + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + + # Full (no-EP) reference + moe_config_full = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=E, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_full = Int4EmulationTritonExperts( + moe_config_full, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_full = _run_emulation_forward( + experts_full, w13_all, w2_all, hidden_states, topk_weights, topk_ids, E, K, N + ) + + # EP sum + out_ep_sum = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device) + for rank in range(ep_size): + start = rank * num_local + end = start + num_local + expert_map = _make_expert_map(E, start, end) + moe_config_ep = FusedMoEConfig( + num_experts=E, + experts_per_token=top_k, + hidden_dim=K, + intermediate_size=N, + num_local_experts=num_local, + num_logical_experts=E, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + activation=MoEActivation.SILU, + in_dtype=torch.bfloat16, + device=device, + routing_method=RoutingMethodType.TopK, + max_num_tokens=512, + ) + experts_ep = Int4EmulationTritonExperts( + moe_config_ep, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + out_rank = _run_emulation_forward_ep( + experts_ep, + w13_all[start:end], + w2_all[start:end], + hidden_states, + topk_weights, + topk_ids, + E, + expert_map, + ) + out_ep_sum = out_ep_sum + out_rank + + assert torch.allclose(out_ep_sum, out_full, atol=1e-3), ( + f"EP sum max diff: {(out_ep_sum - out_full).abs().max().item():.6f}" + ) + + +@pytest.mark.parametrize("E, K, N, top_k, group_size, num_tokens", E2E_CONFIGS) +@pytest.mark.parametrize("fmt", ["gptq", "awq"]) +def test_emulation_output_close_to_bf16_reference( + E, K, N, top_k, group_size, num_tokens, fmt +): + """Emulation output is close to a direct BF16 MoE forward.""" + torch.manual_seed(11) + moe_config = _make_moe_config(E, K, N) + + w13_fp = torch.randn(E, K, 2 * N, dtype=torch.bfloat16, device=device) * 0.02 + w2_fp = torch.randn(E, N, K, dtype=torch.bfloat16, device=device) * 0.02 + + packed13_list, scales13_list, packed2_list, scales2_list = [], [], [], [] + for e in range(E): + q13, s13 = _quantize_sym(w13_fp[e].float(), group_size) + q2, s2 = _quantize_sym(w2_fp[e].float(), group_size) + if fmt == "gptq": + packed13_list.append(gptq_pack(q13, 4, K, 2 * N)) + packed2_list.append(gptq_pack(q2, 4, N, K)) + else: + packed13_list.append(awq_pack(q13, 4, K, 2 * N)) + packed2_list.append(awq_pack(q2, 4, N, K)) + scales13_list.append(s13) + scales2_list.append(s2) + + process_fn = ( + _process_weights_emulation_gptq + if fmt == "gptq" + else _process_weights_emulation_awq + ) + res = process_fn( + torch.stack(packed13_list), + torch.stack(packed2_list), + torch.stack(scales13_list), + torch.stack(scales2_list), + None, + None, + ) + w13_bf16, w2_bf16 = res[0], res[1] + + dummy_scale = torch.ones(1, dtype=torch.float16, device=device) + experts = Int4EmulationTritonExperts( + moe_config, int4_w4a16_moe_quant_config(dummy_scale, dummy_scale) + ) + + hidden_states = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) + topk_weights = torch.softmax( + torch.randn(num_tokens, top_k, dtype=torch.float32, device=device), dim=-1 + ) + topk_ids = torch.stack( + [torch.randperm(E, device=device)[:top_k] for _ in range(num_tokens)] + ).to(torch.int32) + + out_emulation = _run_emulation_forward( + experts, w13_bf16, w2_bf16, hidden_states, topk_weights, topk_ids, E, K, N + ) + + ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device=device) + for m in range(num_tokens): + acc = torch.zeros(K, dtype=torch.float32, device=device) + for k in range(top_k): + e = topk_ids[m, k].item() + w = topk_weights[m, k].item() + gate_up = hidden_states[m] @ w13_bf16[e].T + gate, up = gate_up.chunk(2) + act = F.silu(gate) * up + acc += w * (act @ w2_bf16[e].T).float() + ref[m] = acc.bfloat16() + + rel_l2 = ( + torch.norm(out_emulation.float() - ref.float()) + / torch.norm(ref.float()).clamp(min=1e-6) + ).item() + assert rel_l2 < 0.15, f"[{fmt}] relative L2 = {rel_l2:.4f} (threshold 0.15)" diff --git a/tests/kernels/test_bf16x3_router_gemm_cutedsl.py b/tests/kernels/test_bf16x3_router_gemm_cutedsl.py new file mode 100644 index 000000000000..3c17165cadcb --- /dev/null +++ b/tests/kernels/test_bf16x3_router_gemm_cutedsl.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the experimental SM100 BF16x3 router GEMM.""" + +import pytest +import torch + +from vllm.utils.import_utils import has_cutedsl + + +def _requires_sm100_cutedsl(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + major, _ = torch.cuda.get_device_capability() + if major != 10: + pytest.skip("bf16x3 router GEMM requires SM100-class GPU") + if not has_cutedsl(): + pytest.skip("cutedsl (cutlass) not installed") + + +@pytest.mark.parametrize( + ("num_tokens", "hidden_dim", "num_experts"), + [(48, 6144, 128), (96, 3072, 256), (129, 3072, 17)], +) +def test_bf16x3_router_gemm_matches_reference( + num_tokens: int, hidden_dim: int, num_experts: int +): + _requires_sm100_cutedsl() + from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501 + bf16x3_router_gemm, + ) + + torch.manual_seed(42) + x = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda") + w = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device="cuda") + # Match the observed router weight scale + w *= 0.053 + out = bf16x3_router_gemm(x, w) + ref = torch.nn.functional.linear(x.float(), w) + + assert out.shape == (num_tokens, num_experts) + assert out.dtype == torch.float32 + assert torch.mean(torch.abs(out - ref)).item() < 5e-6 diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index 74dc01472a8e..b8f1cb8bfa73 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -9,6 +9,7 @@ C) Indexer: head_dim=128 (all FP8), quant_block=128 D) DeepseekV4 Attention magnitude range: correctness across small/large values E) Indexer fused Triton kernel: compress+norm+rope+quant+insert + F) Indexer fused two-stage Triton kernel: head=512 cr>=128 (no-overlap) """ import math @@ -24,7 +25,9 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( _fused_kv_compress_norm_rope_insert_indexer_attn, _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn, + _launch_two_stage_sparse_attn_compressor, ) +from vllm.platforms import current_platform from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4 @@ -816,3 +819,130 @@ def test_cutedsl_full_cache_store(compress_ratio: int, store_fp8: bool): torch.testing.assert_close(actual.float(), ref_fp8.float(), rtol=0.0, atol=0.3) else: torch.testing.assert_close(actual.float(), ref.float(), rtol=3e-2, atol=3e-2) + + +# ── Test F: DeepseekV4 Attention two-stage split compressor (Triton) ───────── +# +# Same full pipeline as Test E (state-cache gather -> softmax-weighted compress +# -> RMSNorm -> GPT-J RoPE -> quant -> paged insert), but for the head=512 +# fp8_ds_mla layout via the two-stage split + + +@pytest.mark.skipif( + not current_platform.is_rocm(), + reason="two-stage split compressor is only enabled for ROCm at the moment", +) +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17]) +@pytest.mark.parametrize("kv_block_size", [16, 64]) +def test_fused_kv_insert_split(num_tokens: int, kv_block_size: int): + """Two-stage split compress+norm+rope+quant+insert for the head=512 KV cache.""" + HEAD_DIM = 512 + NOPE_DIM = 448 + ROPE_DIM = 64 + HEAD_BYTES = 584 # 448 fp8 + 128 bf16 + 8 uint8 scale + RMS_EPS = 1e-6 + FP8_MAX = 448.0 + QUANT_BLOCK = 64 + TOKEN_STRIDE = 576 + SCALE_DIM = 8 + STATE_BLOCK_SIZE = 8 # CompressorStateCache block_size for cr=128 + + device = "cuda" + torch.manual_seed(42) + compress_ratio = 128 + overlap = 0 # no overlap for cr=128 + coff = 1 + overlap + + num_pages = (compress_ratio * num_tokens - 1) // STATE_BLOCK_SIZE + 2 + state_cache = torch.randn( + num_pages, + STATE_BLOCK_SIZE, + 2 * coff * HEAD_DIM, # kv_state + score_state + dtype=torch.float32, + device=device, + ) + block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0) + token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + positions = torch.arange( + compress_ratio - 1, + compress_ratio * num_tokens, + compress_ratio, + dtype=torch.int64, + device=device, + ) + rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device) + cos_sin_cache = torch.randn( + compress_ratio * num_tokens, ROPE_DIM, dtype=torch.float32, device=device + ) + + kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1 + kv_cache = torch.zeros( + kv_n_blocks, kv_block_size, HEAD_BYTES, dtype=torch.uint8, device=device + ) + compress_scratch = torch.empty( + num_tokens, HEAD_DIM, dtype=torch.float32, device=device + ) + + _launch_two_stage_sparse_attn_compressor( + state_cache, + token_to_req, + positions, + slot_mapping, + block_table, + STATE_BLOCK_SIZE, + coff * HEAD_DIM, + compress_ratio, + cos_sin_cache, + kv_cache, + slot_mapping, + rms_weight, + RMS_EPS, + QUANT_BLOCK, + TOKEN_STRIDE, + SCALE_DIM, + HEAD_DIM, + ROPE_DIM, + num_tokens, + compress_scratch, + ) + + # PyTorch reference: compress -> RMSNorm -> GPT-J RoPE (pre-quant bf16 row). + ref = _reference_kv_compress_norm_rope( + state_cache, + block_table, + positions, + rms_weight, + cos_sin_cache, + compress_ratio, + overlap, + rms_eps=RMS_EPS, + fp8_max=FP8_MAX, + return_full_cache=True, + ) # [num_tokens, HEAD_DIM] bf16 + + # Dequant + gather the fp8_ds_mla cache back to bf16 (Test B op). + out = torch.zeros(1, num_tokens, HEAD_DIM, dtype=torch.bfloat16, device=device) + seq_lens = torch.tensor([num_tokens], dtype=torch.int32, device=device) + gather_block_table = torch.arange( + kv_n_blocks, dtype=torch.int32, device=device + ).unsqueeze(0) + dequantize_and_gather_k_cache( + out, kv_cache, seq_lens, None, gather_block_table, kv_block_size, offset=0 + ) + recovered = out[0, :num_tokens] + + # NoPE (first 448): FP8 quantized, expect UE8M0 error (same bound as Test A). + nope_diff = (recovered[:, :NOPE_DIM].float() - ref[:, :NOPE_DIM].float()).abs() + for t in range(num_tokens): + _, scales = _ue8m0_reference(ref[t, :NOPE_DIM].float(), QUANT_BLOCK, FP8_MAX) + max_allowed = 16.0 * scales.max().item() + token_diff = nope_diff[t].max().item() + assert token_diff <= max_allowed, ( + f"Token {t} nope diff {token_diff} exceeds max_allowed " + f"{max_allowed} (scale={scales.max().item()})" + ) + + # RoPE (last 64): stored as bf16. The kernel recomputes the rotation, so it + # is bf16-close to the reference rather than bit-exact (cf. test_cutedsl). + torch.testing.assert_close(recovered[:, NOPE_DIM:], ref[:, NOPE_DIM:]) diff --git a/tests/kernels/test_cp_gather_fp8.py b/tests/kernels/test_cp_gather_fp8.py index d9ee8defdb27..36e7f0c9d4ed 100644 --- a/tests/kernels/test_cp_gather_fp8.py +++ b/tests/kernels/test_cp_gather_fp8.py @@ -27,8 +27,8 @@ def _build_test_case(seq_lens, block_size, seed=42): seed: Random seed for reproducibility. Returns: - Tuple of (cache, block_table, seq_lens_t, workspace_starts_t, - num_reqs, total_tokens, expected_output). + Tuple of (cache, block_table, workspace_starts_t, num_reqs, + total_tokens, expected_output). """ torch.manual_seed(seed) @@ -112,7 +112,6 @@ def _build_test_case(seq_lens, block_size, seed=42): # Expected output: exact copy expected[out_idx, NOPE_DIM:] = rope - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") workspace_starts_t = torch.tensor( workspace_starts, dtype=torch.int32, device="cuda" ) @@ -120,7 +119,6 @@ def _build_test_case(seq_lens, block_size, seed=42): return ( cache, block_table, - seq_lens_t, workspace_starts_t, num_reqs, total_tokens, @@ -197,7 +195,6 @@ def _build_test_case_fast(seq_lens, block_size, seed=42): flat_cache[:sl] = token_data[ws : ws + sl] block_start += nb - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") workspace_starts_t = torch.tensor( workspace_starts, dtype=torch.int32, device="cuda" ) @@ -205,7 +202,6 @@ def _build_test_case_fast(seq_lens, block_size, seed=42): return ( cache, block_table, - seq_lens_t, workspace_starts_t, num_reqs, total_tokens, @@ -232,7 +228,6 @@ def test_cp_gather_and_upconvert_fp8_kv_cache(seq_lens, block_size): ( cache, block_table, - seq_lens_t, workspace_starts_t, num_reqs, total_tokens, @@ -244,7 +239,7 @@ def test_cp_gather_and_upconvert_fp8_kv_cache(seq_lens, block_size): ) ops.cp_gather_and_upconvert_fp8_kv_cache( - cache, dst, block_table, seq_lens_t, workspace_starts_t, num_reqs + cache, dst, block_table, workspace_starts_t, num_reqs ) # NoPE: fp8 dequant has rounding error, so we allow small tolerance. @@ -277,8 +272,6 @@ def test_cp_gather_fp8_shuffled_blocks(): ) block_table = torch.tensor([[3, 1]], dtype=torch.int32, device="cuda") workspace_starts = torch.tensor([0], dtype=torch.int32, device="cuda") - seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") - expected = torch.zeros( total_tokens, NOPE_DIM + ROPE_DIM, dtype=torch.bfloat16, device="cuda" ) @@ -314,7 +307,7 @@ def test_cp_gather_fp8_shuffled_blocks(): ) ops.cp_gather_and_upconvert_fp8_kv_cache( - cache, dst, block_table, seq_lens_t, workspace_starts, len(seq_lens) + cache, dst, block_table, workspace_starts, len(seq_lens) ) torch.testing.assert_close( @@ -323,6 +316,64 @@ def test_cp_gather_fp8_shuffled_blocks(): assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) +@pytest.mark.parametrize( + "gather_seq_lens,seq_starts", + [ + ([6, 5, 3], [3, 4, 5]), + ([0, 5, 3], [12, 4, 5]), + ], +) +def test_cp_gather_fp8_with_sequence_starts(gather_seq_lens, seq_starts): + """Gather request slices beginning at arbitrary cache positions.""" + full_seq_lens = [12, 11, 9] + ( + cache, + block_table, + _workspace_starts_t, + num_reqs, + _total_tokens, + full_expected, + ) = _build_test_case(full_seq_lens, block_size=4) + + workspace_starts = torch.tensor( + [0, gather_seq_lens[0], sum(gather_seq_lens[:2])], + dtype=torch.int32, + device="cuda", + ) + seq_starts_t = torch.tensor(seq_starts, dtype=torch.int32, device="cuda") + dst = torch.empty( + sum(gather_seq_lens), + NOPE_DIM + ROPE_DIM, + dtype=torch.bfloat16, + device="cuda", + ) + + ops.cp_gather_and_upconvert_fp8_kv_cache( + cache, + dst, + block_table, + workspace_starts, + num_reqs, + seq_starts_t, + ) + + full_workspace_starts = [0, full_seq_lens[0], sum(full_seq_lens[:2])] + expected = torch.cat( + [ + full_expected[ + full_workspace_starts[i] + seq_starts[i] : full_workspace_starts[i] + + seq_starts[i] + + gather_seq_lens[i] + ] + for i in range(num_reqs) + ] + ) + torch.testing.assert_close( + dst[:, :NOPE_DIM], expected[:, :NOPE_DIM], atol=1e-3, rtol=1e-2 + ) + assert torch.equal(dst[:, NOPE_DIM:], expected[:, NOPE_DIM:]) + + @pytest.mark.parametrize( "seq_lens,block_size", [ @@ -342,7 +393,6 @@ def test_cp_gather_fp8_large_seqlens(seq_lens, block_size): ( cache, block_table, - seq_lens_t, workspace_starts_t, num_reqs, total_tokens, @@ -354,7 +404,7 @@ def test_cp_gather_fp8_large_seqlens(seq_lens, block_size): ) ops.cp_gather_and_upconvert_fp8_kv_cache( - cache, dst, block_table, seq_lens_t, workspace_starts_t, num_reqs + cache, dst, block_table, workspace_starts_t, num_reqs ) torch.testing.assert_close( diff --git a/tests/kernels/test_fla_layernorm_guard.py b/tests/kernels/test_fla_layernorm_guard.py index 4858ff2d7fe4..bb4de9320172 100644 --- a/tests/kernels/test_fla_layernorm_guard.py +++ b/tests/kernels/test_fla_layernorm_guard.py @@ -5,7 +5,7 @@ import torch import torch.nn.functional as F -from vllm.model_executor.layers.fla.ops.layernorm_guard import ( +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( layer_norm_fwd, layernorm_fn, rms_norm_ref, diff --git a/tests/kernels/test_fused_gdn_post_conv.py b/tests/kernels/test_fused_gdn_post_conv.py index ffc8ce281f90..fad77891a4dd 100644 --- a/tests/kernels/test_fused_gdn_post_conv.py +++ b/tests/kernels/test_fused_gdn_post_conv.py @@ -10,7 +10,7 @@ import torch import torch.nn.functional as F -from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import ( +from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import ( fused_post_conv_prep, ) diff --git a/tests/kernels/test_fused_recurrent_packed_decode.py b/tests/kernels/test_fused_recurrent_packed_decode.py index d63186bde118..128a02060043 100644 --- a/tests/kernels/test_fused_recurrent_packed_decode.py +++ b/tests/kernels/test_fused_recurrent_packed_decode.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm.model_executor.layers.fla.ops import ( +from vllm.third_party.flash_linear_attention.ops import ( fused_recurrent_gated_delta_rule, fused_recurrent_gated_delta_rule_packed_decode, ) diff --git a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py index 2b03e83c308a..82a5a6f4ca91 100644 --- a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py +++ b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py @@ -5,11 +5,11 @@ import torch import torch.nn.functional as F -from vllm.model_executor.layers.fla.ops import ( +from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops import ( fused_recurrent_gated_delta_rule, fused_sigmoid_gating_delta_rule_update, ) -from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed DEVICE = current_platform.device_type diff --git a/tests/kernels/test_kda.py b/tests/kernels/test_kda.py index 9ee816adbb82..75c553fb864b 100644 --- a/tests/kernels/test_kda.py +++ b/tests/kernels/test_kda.py @@ -10,12 +10,12 @@ import torch import torch.nn.functional as F -from vllm.model_executor.layers.fla.ops.kda import ( +from vllm.third_party.flash_linear_attention.ops.kda import ( chunk_kda, chunk_kda_with_fused_gate, fused_kda_gate, ) -from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd DEVICE = "cuda" diff --git a/tests/kernels/test_ll_bf16_gemm.py b/tests/kernels/test_ll_bf16_gemm.py index f530f6f7414b..e71c89aa0ddc 100644 --- a/tests/kernels/test_ll_bf16_gemm.py +++ b/tests/kernels/test_ll_bf16_gemm.py @@ -190,8 +190,8 @@ def test_arbitrary_N_splitk(N): @pytest.mark.parametrize( "N,K", - [(256, 7168), (256, 14400), (8, 4096), (384, 7168)], - ids=["DSV3", "DSV4-Flash", "Mixtral", "DSV4-Pro"], + [(256, 7168), (256, 14400), (8, 4096), (384, 7168), (264, 6144)], + ids=["DSV3", "DSV4-Flash", "Mixtral", "DSV4-Pro", "Inkling"], ) def test_single_token(N, K): torch.manual_seed(42) @@ -202,6 +202,15 @@ def test_single_token(N, K): _assert_close(out, _ref(a, b), context=f"M=1 {N}x{K}") +def test_inkling_max_tokens(): + torch.manual_seed(42) + a = torch.randn(64, 6144, dtype=torch.bfloat16, device="cuda") + b = torch.randn(264, 6144, dtype=torch.bfloat16, device="cuda") + out = _gemm(a, b) + assert out.shape == (64, 264) + _assert_close(out, _ref(a, b), context="Inkling M=64") + + # ================================================================= # Numerical robustness # ================================================================= diff --git a/tests/lora/test_gptoss_tp.py b/tests/lora/test_gptoss_tp.py index b9f99d7a8c2c..0f0d807441e9 100644 --- a/tests/lora/test_gptoss_tp.py +++ b/tests/lora/test_gptoss_tp.py @@ -42,6 +42,14 @@ ] +def reformat(text: str) -> str: + # Remove all spaces immediately before or after comma + text = ",".join(map(str.strip, text.split(","))) + # Remove duplicated blank spaces + text = " ".join(map(str.strip, text.split())) + return text + + def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: prompts = [ PROMPT_TEMPLATE.format( @@ -68,13 +76,18 @@ def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: generated_texts.append(generated_text) print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") for i in range(len(EXPECTED_LORA_OUTPUT)): - generated = " ".join(generated_texts[i].split()) - expected = " ".join(EXPECTED_LORA_OUTPUT[i].split()) - assert generated.startswith(expected) + # The generated text may have different numbers of blank space, + # so reformat to compare. + compactGeneratedStr = reformat(generated_texts[i]) + compactExpectedStr = reformat(EXPECTED_LORA_OUTPUT[i]) + if not generated_texts[i].startswith( + EXPECTED_LORA_OUTPUT[i] + ) and not compactGeneratedStr.startswith(compactExpectedStr): + raise AssertionError( + f"Generated: {generated_texts[i]}, Expected: {EXPECTED_LORA_OUTPUT[i]}" + ) -# TODO: make the Mxfp4MoeBackend.TRITON spawn-safe. -# For now just use TRITON_UNFUSED kernel @pytest.mark.parametrize( "mxfp4_use_marlin", [ diff --git a/tests/lora/test_peft_helper.py b/tests/lora/test_peft_helper.py index e3035b00e9e0..23a8aad1a58c 100644 --- a/tests/lora/test_peft_helper.py +++ b/tests/lora/test_peft_helper.py @@ -22,6 +22,8 @@ {"modules_to_save": ["lm_head"]}, "only supports modules_to_save being None", ), + ("test_rank_zero", {"r": 0}, "must be a positive integer"), + ("test_rank_negative", {"r": -8}, "must be a positive integer"), ] @@ -97,3 +99,18 @@ def test_peft_helper_error( PEFTHelper.from_local_dir( test_dir, max_position_embeddings=4096 ).validate_legal(lora_config) + + +@pytest.mark.parametrize("bad_rank", [0, -1, -8]) +def test_peft_helper_invalid_rank_direct(bad_rank: int): + """Regression test: constructing a PEFTHelper with a non-positive rank + must raise a clear ValueError instead of crashing with an unrelated + ZeroDivisionError (r=0) or silently succeeding with a sign-flipped + scaling factor that validate_legal() never catches (r<0, since its only + rank check is the upper bound against max_lora_rank). + + Network-free: constructs PEFTHelper directly rather than going through + from_local_dir(), which needs an on-disk adapter_config.json. + """ + with pytest.raises(ValueError, match="must be a positive integer"): + PEFTHelper(r=bad_rank, lora_alpha=16, target_modules=["q_proj"]) diff --git a/tests/models/inkling/test_contract_validation.py b/tests/models/inkling/test_contract_validation.py new file mode 100644 index 000000000000..42313e7d6775 --- /dev/null +++ b/tests/models/inkling/test_contract_validation.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest + +from vllm.config.compilation import CompilationConfig, CUDAGraphMode +from vllm.models.inkling.common.mm_preprocess import InklingMultiModalDataParser +from vllm.models.inkling.configs import ( + InklingAudioConfig, + InklingModelConfig, + InklingVisionConfig, +) +from vllm.models.inkling.nvidia.sconv_swa_attn import ( + InklingSconvMetadataBuilder, +) +from vllm.v1.attention.backend import AttentionCGSupport + + +@pytest.mark.parametrize( + ("config_cls", "kwargs", "missing"), + [ + (InklingAudioConfig, {"decoder_dmodel": 16}, "n_mel_bins"), + (InklingVisionConfig, {"decoder_dmodel": 16}, "vision_encoder_type"), + ], +) +def test_enabled_tower_requires_architecture_fields(config_cls, kwargs, missing): + with pytest.raises(ValueError, match=missing): + config_cls(**kwargs) + + +def test_inkling_raw_2d_audio_is_rejected_as_ambiguous(): + parser = InklingMultiModalDataParser(target_sr=16_000, target_channels=1) + with pytest.raises(ValueError, match="ambiguous channel layout"): + parser._parse_audio_data(np.zeros((2, 100), dtype=np.float32)) + + +def test_inkling_mtp_chain_norm_is_disabled_by_default(): + assert InklingModelConfig().chain_hidden_post_norm is False + + +def test_inkling_supports_piecewise_cudagraphs(): + support = InklingSconvMetadataBuilder.get_cudagraph_support + assert support(None, None) == AttentionCGSupport.UNIFORM_BATCH + + compilation_config = CompilationConfig( + cudagraph_mode=CUDAGraphMode.PIECEWISE, + splitting_ops=[], + ) + resolved_mode = compilation_config.resolve_cudagraph_mode_and_sizes( + AttentionCGSupport.UNIFORM_BATCH, + "InklingSconvBackend", + ) + + assert resolved_mode == CUDAGraphMode.PIECEWISE diff --git a/tests/models/inkling/test_fa4_rel_attention.py b/tests/models/inkling/test_fa4_rel_attention.py new file mode 100644 index 000000000000..1a11ac220c87 --- /dev/null +++ b/tests/models/inkling/test_fa4_rel_attention.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness test for the Inkling FA4 relative-attention score-mod kernel. + +Checks ``inkling_fa4_rel_attention`` against a pure-PyTorch reference that +implements the relative bias exactly as documented in the Inkling architecture +guide:: + + logit(i, j, h) = (1 / head_dim) * dot(q[i, h], k[j, h]) + rel_bias(i, j, h) + rel_bias(i, j, h) = rel_logits[i, h, i - j] if 0 <= i - j < rel_extent + = 0 otherwise + +with causal (and optionally sliding-window) masking handled by the backend. +""" + +import importlib + +import pytest +import torch + +from vllm.models.inkling.nvidia.attention import ( + InklingAttention, + compute_log_scaling_tau, +) +from vllm.models.inkling.nvidia.ops.fa4_rel_attention import ( + _use_sheared_bias, + inkling_fa4_num_splits, + inkling_fa4_rel_attention, +) +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability + +_cap = current_platform.get_device_capability() if current_platform.is_cuda() else None + +NUM_HEADS = [(4, 4), (8, 2)] # (num_heads, num_kv_heads) +GLOBAL_REL_EXTENTS = [128, 1024] +LOCAL_REL_EXTENTS = [128, 256] +HEAD_DIM = 128 +BLOCK_SIZE = 16 +DTYPE = torch.bfloat16 + + +def test_log_scaling_tau_matches_reference(): + positions = torch.tensor([0, 127999, 128000, 999999], dtype=torch.int64) + actual = compute_log_scaling_tau(positions, 128000, 0.1) + expected = 1.0 + 0.1 * torch.log( + torch.clamp((positions + 1).float() / 128000.0, min=1.0) + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_split_packed_kv_cache(): + attention = InklingAttention.__new__(InklingAttention) + torch.nn.Module.__init__(attention) + attention.head_dim = 8 + attention.kv_cache = torch.arange(2 * 3 * 4 * 16).reshape(2, 3, 4, 16) + + key_cache, value_cache = attention._split_kv_cache() + + assert key_cache.shape == value_cache.shape == (2, 4, 3, 8) + torch.testing.assert_close(key_cache, attention.kv_cache[..., :8].transpose(1, 2)) + torch.testing.assert_close(value_cache, attention.kv_cache[..., 8:].transpose(1, 2)) + + +def test_num_splits_hopper_is_unsplit(monkeypatch): + monkeypatch.setattr( + current_platform, + "get_device_capability", + lambda: DeviceCapability(major=9, minor=0), + ) + assert ( + inkling_fa4_num_splits( + is_local=False, + batch_size=1, + max_query_len=1, + num_heads=16, + num_kv_heads=2, + max_kv_len=1_048_576, + ) + == 1 + ) + + +@pytest.mark.parametrize( + ("major", "expected"), + [(9, False), (10, True), (11, True), (12, False)], +) +def test_sheared_bias_architecture_selection(monkeypatch, major, expected): + monkeypatch.setattr( + current_platform, + "get_device_capability", + lambda: DeviceCapability(major=major, minor=0), + ) + _use_sheared_bias.cache_clear() + try: + assert _use_sheared_bias() is expected + finally: + _use_sheared_bias.cache_clear() + + +@pytest.fixture +def blackwell_platform(monkeypatch): + monkeypatch.setattr( + current_platform, + "get_device_capability", + lambda: DeviceCapability(major=10, minor=0), + ) + + +@pytest.mark.parametrize( + ("batch_size", "max_query_len", "expected"), + [ + (1, 1, (16, 32, 128, 128)), + (8, 1, (2, 4, 8, 16)), + (32, 1, (1, 1, 2, 4)), + (1, 128, (2, 4, 8, 16)), + (1, 2048, (1, 1, 1, 1)), + ], +) +def test_num_splits_all_tp(blackwell_platform, batch_size, max_query_len, expected): + actual = tuple( + inkling_fa4_num_splits( + is_local=False, + batch_size=batch_size, + max_query_len=max_query_len, + num_heads=64 // tp, + num_kv_heads=8 // tp, + max_kv_len=131072, + ) + for tp in (1, 2, 4, 8) + ) + assert actual == expected + + +@pytest.mark.parametrize("tp", [1, 2, 4, 8]) +def test_num_splits_local_is_unsplit(tp): + assert ( + inkling_fa4_num_splits( + is_local=True, + batch_size=1, + max_query_len=1, + num_heads=64 // tp, + num_kv_heads=16 // tp, + max_kv_len=512, + ) + == 1 + ) + + +@pytest.mark.parametrize( + ("max_kv_len", "expected"), + [(8192, 32), (65536, 64), (1048576, 128)], +) +@pytest.mark.parametrize("tp", [4, 8]) +def test_num_splits_long_context_bound(blackwell_platform, tp, max_kv_len, expected): + assert ( + inkling_fa4_num_splits( + is_local=False, + batch_size=1, + max_query_len=1, + num_heads=64 // tp, + num_kv_heads=8 // tp, + max_kv_len=max_kv_len, + ) + == expected + ) + + +def _ref_rel_attn( + q: torch.Tensor, # [total_q, H, D] + key_cache: torch.Tensor, # [num_blocks, block, Hkv, D] + value_cache: torch.Tensor, + rel_logits: torch.Tensor, # [total_q, H, rel_extent] + *, + q_lens: list[int], + kv_lens: list[int], + block_table: torch.Tensor, + scale: float, + rel_extent: int, + window_left: int | None, +) -> torch.Tensor: + num_kv_heads = key_cache.shape[2] + num_heads = q.shape[1] + g = num_heads // num_kv_heads + bt = block_table.cpu().numpy() + out = torch.empty_like(q) + + start = 0 + for i, (ql, kl) in enumerate(zip(q_lens, kv_lens)): + qi = q[start : start + ql].float() # [ql, H, D] + rl = rel_logits[start : start + ql].float() # [ql, H, rel_extent] + + nblk = (kl + BLOCK_SIZE - 1) // BLOCK_SIZE + blk = bt[i, :nblk] + k = key_cache[blk].reshape(-1, num_kv_heads, HEAD_DIM)[:kl].float() + v = value_cache[blk].reshape(-1, num_kv_heads, HEAD_DIM)[:kl].float() + k = k.repeat_interleave(g, dim=1) # [kl, H, D] + v = v.repeat_interleave(g, dim=1) + + # [H, ql, kl] + scores = torch.einsum("qhd,khd->hqk", qi, k) * scale + + dev = q.device + qpos = torch.arange(ql, device=dev).view(ql, 1) + (kl - ql) # query pos + kpos = torch.arange(kl, device=dev).view(1, kl) + dist = qpos - kpos # [ql, kl] = i - j + + # Relative bias: rel_logits[i, h, dist] when 0 <= dist < rel_extent. + in_rng = (dist >= 0) & (dist < rel_extent) # [ql, kl] + idx = dist.clamp(0, rel_extent - 1) + # gather per head: bias[h, i, j] = rl[i, h, idx[i, j]] + bias = rl.permute(1, 0, 2).gather( # [H, ql, rel_extent] + 2, idx.unsqueeze(0).expand(num_heads, -1, -1) + ) # [H, ql, kl] + bias = torch.where(in_rng.unsqueeze(0), bias, torch.zeros_like(bias)) + scores = scores + bias + + mask = dist < 0 # causal + if window_left is not None: + mask = mask | (dist > window_left) + scores.masked_fill_(mask.unsqueeze(0), float("-inf")) + + probs = torch.softmax(scores, dim=-1) + out[start : start + ql] = torch.einsum("hqk,khd->qhd", probs, v).to(q.dtype) + start += ql + return out + + +def _run_case(seq_lens, num_heads, num_kv_heads, rel_extent, window_left, seed=0): + torch.manual_seed(seed) + device = "cuda" + q_lens = [s[0] for s in seq_lens] + kv_lens = [s[1] for s in seq_lens] + total_q = sum(q_lens) + num_seqs = len(seq_lens) + + scale = 1.0 / HEAD_DIM + + # q/k are RMS-normed in the model (unit-ish norm); normalize here so the + # logit magnitudes are realistic and the bias is not numerically dwarfed. + q = torch.randn(total_q, num_heads, HEAD_DIM, device=device, dtype=DTYPE) + q = torch.nn.functional.normalize(q.float(), dim=-1).to(DTYPE) + + # Paged KV cache. + max_blocks = (max(kv_lens) + BLOCK_SIZE - 1) // BLOCK_SIZE + num_blocks = num_seqs * max_blocks + 1 + key_cache = torch.randn( + num_blocks, BLOCK_SIZE, num_kv_heads, HEAD_DIM, device=device, dtype=DTYPE + ) + key_cache = torch.nn.functional.normalize(key_cache.float(), dim=-1).to(DTYPE) + value_cache = torch.randn( + num_blocks, BLOCK_SIZE, num_kv_heads, HEAD_DIM, device=device, dtype=DTYPE + ) + + # Distinct blocks per sequence (block 0 left as a never-referenced pad). + block_table = torch.zeros(num_seqs, max_blocks, dtype=torch.int32, device=device) + for i in range(num_seqs): + block_table[i] = torch.arange( + 1 + i * max_blocks, 1 + (i + 1) * max_blocks, dtype=torch.int32 + ) + + cu_seqlens_q = torch.tensor( + [0, *torch.cumsum(torch.tensor(q_lens), 0).tolist()], + dtype=torch.int32, + device=device, + ) + cache_seqlens = torch.tensor(kv_lens, dtype=torch.int32, device=device) + + rel_logits = torch.randn(total_q, num_heads, rel_extent, device=device, dtype=DTYPE) + + window_size = (-1, -1) if window_left is None else (window_left, 0) + + preallocated_out = torch.empty_like(q) + num_splits = inkling_fa4_num_splits( + is_local=window_left is not None, + batch_size=num_seqs, + max_query_len=max(q_lens), + num_heads=num_heads, + num_kv_heads=num_kv_heads, + max_kv_len=max(kv_lens), + ) + out = inkling_fa4_rel_attention( + q, + key_cache, + value_cache, + block_table=block_table, + cache_seqlens=cache_seqlens, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_q=max(q_lens), + softmax_scale=scale, + causal=True, + window_size=window_size, + rel_extent=rel_extent, + rel_logits=rel_logits, + num_splits=num_splits, + out=preallocated_out, + ) + assert out.data_ptr() == preallocated_out.data_ptr() + out = out.view(total_q, num_heads, HEAD_DIM) + + ref = _ref_rel_attn( + q, + key_cache, + value_cache, + rel_logits, + q_lens=q_lens, + kv_lens=kv_lens, + block_table=block_table, + scale=scale, + rel_extent=rel_extent, + window_left=window_left, + ) + + torch.testing.assert_close(out.float(), ref.float(), atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +@pytest.mark.skipif( + _cap is None or _cap.major < 9, + reason="FA4 score-mod requires Hopper+ (SM90+)", +) +@torch.inference_mode() +def test_score_mod_relative_attention(monkeypatch): + module = importlib.import_module("vllm.models.inkling.nvidia.ops.fa4_rel_attention") + monkeypatch.setattr(module, "_use_sheared_bias", lambda: False) + _run_case( + [(64, 64), (1, 80)], + num_heads=4, + num_kv_heads=4, + rel_extent=128, + window_left=None, + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +@pytest.mark.skipif( + _cap is None or _cap.major < 9, + reason="FA4 score-mod requires Hopper+ (SM90+)", +) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize( + "seq_lens", + [ + [(64, 64)], # single full prefill + [(64, 64), (33, 33), (17, 17)], # ragged prefill batch + [(512, 512)], # seq_len >> rel_extent (most keys get zero bias) + [(300, 300), (512, 512), (129, 129)], # large ragged batch + ], +) +@pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS) +@torch.inference_mode() +def test_full_attention(seq_lens, num_heads, rel_extent): + # rel_extent=128 exercises the out-of-range (zero bias) path; 1024 covers all. + # With the 512-token cases and rel_extent=128, query/seq lengths are far + # larger than rel_extent so the vast majority of (i, j) pairs are out of + # range and must contribute zero bias. + _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +@pytest.mark.skipif( + _cap is None or _cap.major < 9, + reason="FA4 score-mod requires Hopper+ (SM90+)", +) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize( + "seq_lens", + [ + [(200, 512)], # chunked prefill: q_len=200 (> rel_extent), 312 cached + [(200, 512), (50, 300), (1, 400)], # mixed chunked + decode + ], +) +@pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS) +@torch.inference_mode() +def test_chunked_prefill(seq_lens, num_heads, rel_extent): + # q_len < kv_len with q_len itself larger than rel_extent (for the 128 case): + # exercises the seqlen_k - seqlen_q offset together with the out-of-range path. + _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +@pytest.mark.skipif( + _cap is None or _cap.major < 9, + reason="FA4 score-mod requires Hopper+ (SM90+)", +) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize( + "seq_lens", + [ + [(64, 64), (40, 40)], # seq_len > window + [(512, 512), (300, 300)], # seq_len/query_len >> window + [(1, 512)], # decode with kv_len >> window + ], +) +@pytest.mark.parametrize("local_extent", LOCAL_REL_EXTENTS) +@torch.inference_mode() +def test_sliding_window(seq_lens, num_heads, local_extent): + # Local layers use window_size=(local_extent-1, 0) and rel_extent==local_extent. + # With the 512-token cases, query/seq lengths far exceed the window so most + # keys are masked out by the sliding window. + _run_case( + seq_lens, + num_heads[0], + num_heads[1], + rel_extent=local_extent, + window_left=local_extent - 1, + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +@pytest.mark.skipif( + _cap is None or _cap.major < 9, + reason="FA4 score-mod requires Hopper+ (SM90+)", +) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize( + "seq_lens", + [ + [(1, 50)], + [(1, 50), (1, 7), (1, 200)], + [(1, 512), (1, 333)], # kv_len >> rel_extent + ], +) +@pytest.mark.parametrize("rel_extent", GLOBAL_REL_EXTENTS) +@torch.inference_mode() +def test_decode(seq_lens, num_heads, rel_extent): + # q_len=1 with kv_len>q_len: the score-mod's seqlen_k - seqlen_q offset path. + _run_case(seq_lens, num_heads[0], num_heads[1], rel_extent, window_left=None) diff --git a/tests/models/inkling/test_fa4_warmup.py b/tests/models/inkling/test_fa4_warmup.py new file mode 100644 index 000000000000..eea9798aa93c --- /dev/null +++ b/tests/models/inkling/test_fa4_warmup.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.models.inkling.nvidia.ops.fa4_rel_attention import ( + bucket_max_seqlen_q, + inkling_fa4_num_splits, +) +from vllm.models.inkling.nvidia.ops.fa4_warmup import ( + InklingFA4WarmupConfig, + _iter_compile_units, + _num_warps_bucket, +) + + +def test_bucket_max_seqlen_q(): + assert [bucket_max_seqlen_q(n) for n in range(1, 10)] == [ + 1, + 2, + 4, + 4, + 8, + 8, + 8, + 8, + 16, + ] + + +def test_warmup_enumerates_every_runtime_compile_class(): + config = InklingFA4WarmupConfig( + num_heads=16, + num_kv_heads=2, + head_dim=128, + rel_extent=1024, + window_size=(-1, -1), + is_local=False, + max_kv_len=65536, + dtype=torch.bfloat16, + kv_dtype=torch.bfloat16, + block_size=16, + max_num_reqs=64, + max_num_batched_tokens=192, + ) + warmed_keys = {unit.key[-1] for unit in _iter_compile_units(config)} + + for query_len in range(1, config.max_num_batched_tokens + 1): + max_seqlen_q = bucket_max_seqlen_q(query_len) + max_num_reqs = min( + config.max_num_reqs, + config.max_num_batched_tokens - query_len + 1, + ) + for num_reqs in range(1, max_num_reqs + 1): + num_splits = inkling_fa4_num_splits( + is_local=config.is_local, + batch_size=num_reqs, + max_query_len=max_seqlen_q, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + max_kv_len=config.max_kv_len, + ) + runtime_key = ( + max_seqlen_q, + num_splits, + _num_warps_bucket(num_reqs) if num_splits > 1 else None, + num_reqs > 1024, + ) + assert runtime_key in warmed_keys diff --git a/tests/models/inkling/test_mm_towers.py b/tests/models/inkling/test_mm_towers.py new file mode 100644 index 000000000000..d969c8101a29 --- /dev/null +++ b/tests/models/inkling/test_mm_towers.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the fused Inkling vision/audio tower kernels. + +Each kernel is checked against a pure-PyTorch reference implementing the +towers' original op sequence (native ``rms_norm`` semantics: fp32 variance, +bf16-rounded weight multiply; exact-erf GELU; fp32-accumulated embedding sum). +The fused kernels keep the same accumulation dtype and rounding points, so +outputs differ from the reference only by reduction order — a few bf16 ulps. +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm.platforms import current_platform + +if not current_platform.is_cuda(): + pytest.skip("requires CUDA", allow_module_level=True) + +from vllm.models.inkling.common.towers import fold_timespace_to_depth +from vllm.models.inkling.nvidia.ops.mm_towers import dmel_embed_sum_norm, rmsnorm_gelu + +DTYPE = torch.bfloat16 + + +def _bf16_spacing(x: torch.Tensor) -> torch.Tensor: + return torch.exp2(torch.floor(torch.log2(x.float().abs().clamp(min=1e-30)))) * ( + 2**-7 + ) + + +def _assert_close_ulps( + got: torch.Tensor, + ref: torch.Tensor, + max_ulps: float = 4.0, + atol: float = 1e-3, + pre_act: torch.Tensor | None = None, +) -> None: + """Fused vs reference differ only by fp32 reduction order, which shows up + as a few bf16 ulps at any magnitude — a fixed rtol misrepresents that, so + compare against the reference's local ulp spacing. ``pre_act`` (the + reference pre-activation) adds derivative-propagated slack: near + ``gelu(x) ~ 0`` an ulp-level input flip legitimately moves the output by + many of ITS (tiny) ulps, bounded by |gelu'| <= 1.13 times the input ulps.""" + g, r = got.float(), ref.float() + tol = torch.clamp(max_ulps * _bf16_spacing(ref), min=atol) + if pre_act is not None: + tol = tol + 2.5 * _bf16_spacing(pre_act) + bad = (g - r).abs() > tol + assert not bad.any(), ( + f"{int(bad.sum())}/{ref.numel()} elements beyond tolerance; " + f"max abs diff {(g - r).abs().max().item():.3e}" + ) + + +def _ref_rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + # Matches ir.ops.rms_norm: fp32 normalize, bf16 weight multiply. + x32 = x.float() + var = x32.pow(2).mean(dim=-1, keepdim=True) + xn = x32 * torch.rsqrt(var + eps) + return (xn.to(weight.dtype) * weight).to(x.dtype) + + +@pytest.mark.parametrize("rows", [1, 7, 1200, 38400]) +@pytest.mark.parametrize("dim", [128, 320, 4800, 6144]) +@pytest.mark.parametrize("gelu", [True, False]) +def test_rmsnorm_gelu(rows: int, dim: int, gelu: bool) -> None: + torch.manual_seed(rows * dim) + x = torch.randn(rows, dim, device="cuda", dtype=DTYPE) + w = torch.randn(dim, device="cuda", dtype=DTYPE) + + h = _ref_rms_norm(x, w, 1e-5) + ref = F.gelu(h) if gelu else h + got = rmsnorm_gelu(x, w, 1e-5, gelu=gelu) + + _assert_close_ulps(got, ref, pre_act=h if gelu else None) + + +@pytest.mark.parametrize("n", [1, 5, 64]) +@pytest.mark.parametrize( + "shape,fold", + [ + ((2, 8, 8, 128), (1, 2)), # the real L0 -> L1 vision transition + ((2, 4, 4, 320), (1, 2)), + ((2, 8, 8, 128), (2, 2)), # temporal + spatial fold + ], +) +def test_rmsnorm_gelu_folded_store( + n: int, shape: tuple[int, ...], fold: tuple[int, int] +) -> None: + torch.manual_seed(n) + x = torch.randn(n, *shape, device="cuda", dtype=DTYPE) + w = torch.randn(shape[-1], device="cuda", dtype=DTYPE) + + plain = rmsnorm_gelu(x, w, 1e-5, gelu=True) + ref = fold_timespace_to_depth(plain, *fold) + got = rmsnorm_gelu(x, w, 1e-5, gelu=True, fold=fold) + + # The folded store is a pure permutation of the unfolded output. + assert got.shape == ref.shape + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + +@pytest.mark.parametrize("num_frames", [1, 3, 100, 4097]) +@pytest.mark.parametrize("with_norm", [True, False]) +def test_dmel_embed_sum_norm(num_frames: int, with_norm: bool) -> None: + torch.manual_seed(num_frames) + n_bins, vocab, dim = 80, 16, 6144 + idx = torch.randint( + 0, vocab, (num_frames, n_bins), device="cuda", dtype=torch.int32 + ) + table = torch.randn(n_bins * vocab, dim, device="cuda", dtype=DTYPE) + norm_w = torch.randn(dim, device="cuda", dtype=DTYPE) + + flat = (torch.arange(n_bins, device="cuda", dtype=torch.int32) * vocab).unsqueeze( + 0 + ) + idx + ref = ( + F.embedding(flat.reshape(-1).long(), table) + .reshape(num_frames, n_bins, dim) + .sum(dim=1) + ) + if with_norm: + ref = _ref_rms_norm(ref, norm_w, 1e-6) + + got = dmel_embed_sum_norm(idx, table, norm_w if with_norm else None, 1e-6) + _assert_close_ulps(got, ref) + + +def test_empty_inputs() -> None: + x = torch.empty(0, 320, device="cuda", dtype=DTYPE) + w = torch.randn(320, device="cuda", dtype=DTYPE) + assert rmsnorm_gelu(x, w, 1e-5).shape == (0, 320) + + idx = torch.empty(0, 80, device="cuda", dtype=torch.int32) + table = torch.randn(1280, 6144, device="cuda", dtype=DTYPE) + assert dmel_embed_sum_norm(idx, table, None, 0.0).shape == (0, 6144) diff --git a/tests/models/inkling/test_moe_weight_layout.py b/tests/models/inkling/test_moe_weight_layout.py new file mode 100644 index 000000000000..e6c35c57b4d0 --- /dev/null +++ b/tests/models/inkling/test_moe_weight_layout.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.lora.utils import get_supported_lora_modules +from vllm.models.inkling.nvidia import moe +from vllm.models.inkling.nvidia.model import _TmlForCausalLMBase +from vllm.platforms import current_platform + + +def test_gate_loads_directly_into_padded_runtime_weight() -> None: + gate = moe.InklingGate( + d_model=4, + n_routed_experts=5, + n_shared_experts=2, + experts_per_token=2, + route_scale=1.0, + ) + loaded = torch.arange(28, dtype=gate.weight.dtype).reshape(7, 4) + + gate.weight.weight_loader(gate.weight, loaded) + + assert gate.weight.shape == (8, 4) + torch.testing.assert_close(gate.weight[:7], loaded) + torch.testing.assert_close(gate.weight[7], torch.zeros(4)) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +@pytest.mark.parametrize(("num_tokens", "expected_calls"), [(1, 1), (64, 1), (65, 0)]) +def test_gate_uses_ll_bf16_gemm_through_token_limit( + monkeypatch, num_tokens, expected_calls +) -> None: + gate = moe.InklingGate( + d_model=8, + n_routed_experts=5, + n_shared_experts=2, + experts_per_token=2, + route_scale=1.0, + ).to(device="cuda", dtype=torch.bfloat16) + hidden_states = torch.randn(num_tokens, 8, device="cuda", dtype=torch.bfloat16) + calls = [] + + def fake_ll_bf16_gemm(x, weight): + calls.append((x, weight)) + return torch.ones(num_tokens, 8, device="cuda", dtype=torch.float32) + + monkeypatch.setattr( + moe.current_platform, "has_device_capability", lambda capability: True + ) + monkeypatch.setattr(moe.ll_bf16, "is_available", lambda: True) + monkeypatch.setattr(moe.ll_bf16, "ll_bf16_gemm", fake_ll_bf16_gemm) + + logits = gate.compute_logits(hidden_states) + + assert len(calls) == expected_calls + if calls: + assert calls[0][0] is hidden_states + assert calls[0][1] is gate.weight + assert logits.shape == (num_tokens, 8) + assert logits.dtype == torch.float32 + + +def test_gate_is_not_a_lora_target() -> None: + model = torch.nn.Module() + model.gate = moe.InklingGate( + d_model=4, + n_routed_experts=5, + n_shared_experts=2, + experts_per_token=2, + route_scale=1.0, + ) + + assert "gate" not in get_supported_lora_modules(model) + + +def test_custom_embedding_is_not_a_lora_target() -> None: + model = torch.nn.Module() + model.embedding_modules = _TmlForCausalLMBase.embedding_modules + + supported = get_supported_lora_modules(model) + + assert "embed_tokens" not in supported + assert "lm_head" in supported + + +@pytest.mark.parametrize(("projection", "amax"), [("w13", 4.375), ("w2", 2960.0)]) +def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None: + experts = SimpleNamespace( + w13_input_scale=torch.nn.Parameter(torch.empty(3, 2)), + w2_input_scale=torch.nn.Parameter(torch.empty(3)), + ) + layer = SimpleNamespace(experts=SimpleNamespace(routed_experts=experts)) + + loaded = moe.InklingMoE.load_expert_weight( + layer, + f"experts.{projection}_weight.input_amax", + torch.tensor([amax]), + ) + + scale = getattr(experts, f"{projection}_input_scale") + expected = torch.full_like(scale, amax / (448.0 * 6.0)) + torch.testing.assert_close(scale, expected) + assert loaded == [f"experts.routed_experts.{projection}_input_scale"] + + +def test_sink_down_projection_is_packed_during_load(monkeypatch) -> None: + monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1) + sink = moe.InklingSinkExperts(n_experts=2, d_model=3, d_mlp=8) + loaded = torch.arange(48, dtype=sink.w2_weight.dtype).reshape(2, 3, 8) + + sink.load_weight("w2_weight", loaded) + + expected = loaded[:, :, 4:].permute(1, 0, 2).reshape(3, 8) + torch.testing.assert_close(sink.w2_weight, expected) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires CUDA") +def test_sink_packed_weight_forward_matches_expert_sum(monkeypatch) -> None: + monkeypatch.setattr(moe, "get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(moe, "get_tensor_model_parallel_rank", lambda: 1) + sink = moe.InklingSinkExperts(n_experts=2, d_model=3, d_mlp=8).to( + device="cuda", dtype=torch.bfloat16 + ) + torch.manual_seed(1) + w13 = torch.randn(2, 16, 3, dtype=torch.bfloat16) + w2 = torch.randn(2, 3, 8, dtype=torch.bfloat16) + sink.load_weight("w13_weight", w13) + sink.load_weight("w2_weight", w2) + x = torch.randn(5, 3, device="cuda", dtype=torch.bfloat16) + gammas = torch.randn(5, 2, device="cuda") + + output = sink(x, gammas) + + raw = torch.einsum("td,efd->tef", x, w13[:, 8:].to("cuda")) + hidden = torch.nn.functional.silu(raw[:, :, 0::2].float()) + hidden = (hidden * raw[:, :, 1::2] * gammas[:, :, None]).to(torch.bfloat16) + expected = torch.einsum("tef,edf->td", hidden, w2[:, :, 4:].to("cuda")) + torch.testing.assert_close(output, expected, rtol=0, atol=0) diff --git a/tests/models/inkling/test_mtp_input_fusion.py b/tests/models/inkling/test_mtp_input_fusion.py new file mode 100644 index 000000000000..7ab33dd45ae8 --- /dev/null +++ b/tests/models/inkling/test_mtp_input_fusion.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bit-exactness tests for the fused MTP depth-layer input kernel. + +``embed_dual_rmsnorm_cat`` must match the unfused module sequence exactly: +each rmsnorm computes in fp32 and rounds to bf16 at the same points as the +vendored ``rmsnorm`` kernel (including the bf16 round-trip between the +chained backbone embed_norm and the depth embed_norm), and the fused row +gather matches ``F.embedding``. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_cuda(): + pytest.skip("requires CUDA", allow_module_level=True) + +from vllm.models.inkling.nvidia.ops.norm import ( + embed_dual_rmsnorm_cat, + embed_rmsnorm, + rmsnorm, +) + +EPS = 1e-6 +VOCAB = 4096 + + +def _ref(hidden, w_h, w_e, emb, w_pre=None): + if w_pre is not None: + emb = rmsnorm(emb, w_pre, EPS) + return torch.cat([rmsnorm(hidden, w_h, EPS), rmsnorm(emb, w_e, EPS)], dim=-1) + + +@pytest.mark.parametrize("n", [1536, 6144]) +@pytest.mark.parametrize("t", [0, 1, 7, 256]) +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) +def test_embed_dual_rmsnorm_cat(n: int, t: int, ids_dtype: torch.dtype) -> None: + torch.manual_seed(0) + dev = "cuda" + table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) + w_h = torch.randn(n, device=dev).to(torch.bfloat16) + w_e = (1 + 0.01 * torch.randn(n, device=dev)).to(torch.bfloat16) + w_pre = torch.randn(n, device=dev).to(torch.bfloat16) + ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) + hidden = (torch.randn(t, n, device=dev) * 2).to(torch.bfloat16) + emb = table[ids.long()] + + # Fused gather + chained backbone pre-norm (the decode draft-step path). + out = embed_dual_rmsnorm_cat( + hidden, + w_h, + w_e, + EPS, + input_ids=ids, + embed_table=table, + pre_norm_weight=w_pre, + ) + assert out.shape == (t, 2 * n) + assert torch.equal(out, _ref(hidden, w_h, w_e, emb, w_pre)) + + # Precomputed embeds, no pre-norm (draft prefill with target-merged MM + # embeddings, already backbone-normed). + out = embed_dual_rmsnorm_cat(hidden, w_h, w_e, EPS, embeds=emb) + assert torch.equal(out, _ref(hidden, w_h, w_e, emb)) + + # Fused gather, no pre-norm (use_embed_norm=False). + out = embed_dual_rmsnorm_cat( + hidden, w_h, w_e, EPS, input_ids=ids, embed_table=table + ) + assert torch.equal(out, _ref(hidden, w_h, w_e, emb)) + + +@pytest.mark.parametrize("n", [1536, 6144]) +@pytest.mark.parametrize("t", [0, 1, 7, 256]) +@pytest.mark.parametrize("ids_dtype", [torch.int32, torch.int64]) +def test_embed_rmsnorm(n: int, t: int, ids_dtype: torch.dtype) -> None: + torch.manual_seed(0) + dev = "cuda" + table = (torch.randn(VOCAB, n, device=dev) * 0.3).to(torch.bfloat16) + w = torch.randn(n, device=dev).to(torch.bfloat16) + ids = torch.randint(0, VOCAB, (t,), device=dev, dtype=ids_dtype) + ref_emb = table[ids.long()] + + # Gather + embed_norm (base model / MTP prefill embed path). + out = embed_rmsnorm(ids, table, w, EPS) + assert out.shape == (t, n) + assert torch.equal(out, rmsnorm(ref_emb, w, EPS) if t else ref_emb) + + # Pure gather (use_embed_norm=False / replicated module forward). + out = embed_rmsnorm(ids, table, None, EPS) + assert torch.equal(out, ref_emb) + + # Chained first-layer attn_norm (the target text-path forward): one launch + # emits both the residual and layer 0's normed attention input. + w_chain = (1 + 0.05 * torch.randn(n, device=dev)).to(torch.bfloat16) + res, attn_in = embed_rmsnorm(ids, table, w, EPS, chain_weight=w_chain) + ref_res = rmsnorm(ref_emb, w, EPS) if t else ref_emb + assert torch.equal(res, ref_res) + assert torch.equal(attn_in, rmsnorm(ref_res, w_chain, EPS) if t else ref_res) + + # Chained without embed_norm (use_embed_norm=False). + res, attn_in = embed_rmsnorm(ids, table, None, EPS, chain_weight=w_chain) + assert torch.equal(res, ref_emb) + assert torch.equal(attn_in, rmsnorm(ref_emb, w_chain, EPS) if t else ref_emb) diff --git a/tests/models/inkling/test_qkvr_prep.py b/tests/models/inkling/test_qkvr_prep.py new file mode 100644 index 000000000000..1e0bf49bb6f3 --- /dev/null +++ b/tests/models/inkling/test_qkvr_prep.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm.models.inkling.nvidia.ops.sconv as sconv +from vllm.models.inkling.nvidia.ops import qkvr_prep +from vllm.platforms import current_platform + +_cap = current_platform.get_device_capability() if current_platform.is_cuda() else None + + +def _make_inputs(*, is_local: bool, tokens: int = 33, tp_size: int = 4): + torch.manual_seed(0) + heads = 64 // tp_size + kv_heads = (16 if is_local else 8) // tp_size + head_dim = 128 + d_rel = 16 + rel_extent = 512 if is_local else 1024 + page_size = 16 + num_blocks = (tokens + page_size - 1) // page_size + q_width = heads * head_dim + kv_width = kv_heads * head_dim + r_width = heads * d_rel + device = "cuda" + + qkvr = torch.randn( + tokens, + q_width + 2 * kv_width + r_width, + device=device, + dtype=torch.bfloat16, + ) + k_weight = torch.randn(kv_width, 4, device=device, dtype=torch.bfloat16) + v_weight = torch.randn_like(k_weight) + q_norm_weight = torch.randn(head_dim, device=device, dtype=torch.bfloat16) + k_norm_weight = torch.randn_like(q_norm_weight) + rel_proj = torch.randn(d_rel, rel_extent, device=device, dtype=torch.bfloat16) + conv_cache = torch.zeros( + num_blocks, + kv_heads, + page_size, + 2 * head_dim, + device=device, + dtype=torch.bfloat16, + ) + key_cache = torch.empty( + num_blocks, + page_size, + kv_heads, + head_dim, + device=device, + dtype=torch.bfloat16, + ) + value_cache = torch.empty_like(key_cache) + positions = torch.arange(tokens, device=device, dtype=torch.int64) + block_table = torch.arange(num_blocks, device=device, dtype=torch.int32)[None] + seq_idx = torch.zeros(tokens, device=device, dtype=torch.int32) + slots = torch.arange(tokens, device=device, dtype=torch.int64) + query_start = torch.zeros(tokens, device=device, dtype=torch.int32) + log_scaling_n_floor = None if is_local else 128000 + log_scaling = None + if log_scaling_n_floor is not None: + log_scaling = torch.linspace( + 1.0, + 1.1, + tokens, + device=device, + dtype=torch.float32, + ) + return ( + qkvr, + k_weight, + v_weight, + q_norm_weight, + k_norm_weight, + rel_proj, + 1e-6, + heads, + kv_heads, + head_dim, + d_rel, + conv_cache, + key_cache, + value_cache, + positions, + block_table, + seq_idx, + slots, + query_start, + slots, + 0, + head_dim, + page_size, + log_scaling, + ) + + +def _reference(args): + qkvr = args[0] + heads, kv_heads, head_dim, d_rel = args[7:11] + q_width = heads * head_dim + kv_width = kv_heads * head_dim + q, k, v, r = qkvr.split((q_width, kv_width, kv_width, heads * d_rel), dim=1) + conv_cache = args[11].clone() + k = sconv.fused_sconv( + k.contiguous(), + args[1], + conv_cache, + args[14], + args[15], + args[16], + args[17], + args[18], + args[20], + head_dim, + args[22], + ) + v = sconv.fused_sconv( + v.contiguous(), + args[2], + conv_cache, + args[14], + args[15], + args[16], + args[17], + args[18], + args[21], + head_dim, + args[22], + ) + + def rms_norm(x, weight): + x = x.reshape(-1, head_dim).float() + rstd = torch.rsqrt(x.square().mean(1, keepdim=True) + args[6]) + return (x * rstd * weight.float()).to(qkvr.dtype) + + q = rms_norm(q, args[3]).view(qkvr.shape[0], heads, head_dim) + k = rms_norm(k, args[4]).view(qkvr.shape[0], kv_heads, head_dim) + v = v.view(qkvr.shape[0], kv_heads, head_dim) + rel = torch.mm(r.reshape(-1, d_rel), args[5]).view(qkvr.shape[0], heads, -1) + if args[23] is not None: + q = (q.float() * args[23][:, None, None]).to(q.dtype) + rel = (rel.float() * args[23][:, None, None]).to(rel.dtype) + + key_cache = args[12].clone() + value_cache = args[13].clone() + slots = args[19] + valid = slots >= 0 + key_cache.view(-1, kv_heads, head_dim)[slots[valid]] = k[valid] + value_cache.view(-1, kv_heads, head_dim)[slots[valid]] = v[valid] + return q.flatten(1), rel, key_cache, value_cache + + +@pytest.mark.skipif( + _cap is None, + reason="Inkling QKVR prep kernels require CUDA", +) +@pytest.mark.parametrize( + ("is_local", "tokens", "tp_size"), + [ + (True, 9, 4), + (False, 33, 4), + (True, 33, 8), + (True, 128, 4), + (False, 128, 8), + (True, 512, 4), + (False, 640, 8), + ], +) +@torch.inference_mode() +def test_qkvr_prep_matches_reference(is_local, tokens, tp_size): + args = _make_inputs(is_local=is_local, tokens=tokens, tp_size=tp_size) + ref_q, ref_rel, ref_key, ref_value = _reference(args) + + q, rel = qkvr_prep.fused_qkvr_prep(*args) + + torch.testing.assert_close(q, ref_q, rtol=0.01, atol=0.02) + torch.testing.assert_close(rel, ref_rel, rtol=0.02, atol=0.125) + torch.testing.assert_close(args[12], ref_key, rtol=0.01, atol=0.01) + torch.testing.assert_close(args[13], ref_value, rtol=0, atol=0) + + +@pytest.mark.skipif( + _cap is None, + reason="Inkling QKVR prep kernels require CUDA", +) +@pytest.mark.parametrize("tokens", [9, 128]) +@torch.inference_mode() +def test_qkvr_log_scaling_preserves_bf16_norm_output(tokens): + args = list(_make_inputs(is_local=False, tokens=tokens, tp_size=8)) + tau = args[23] + assert tau is not None + + scaled_q, scaled_rel = qkvr_prep.fused_qkvr_prep(*args) + args[23] = None + unscaled_q, unscaled_rel = qkvr_prep.fused_qkvr_prep(*args) + + expected_q = (unscaled_q.float() * tau[:, None]).to(unscaled_q.dtype) + expected_rel = (unscaled_rel.float() * tau[:, None, None]).to(unscaled_rel.dtype) + torch.testing.assert_close(scaled_q, expected_q, rtol=0, atol=0) + torch.testing.assert_close(scaled_rel, expected_rel, rtol=0, atol=0) + + +@pytest.mark.skipif( + _cap is None, + reason="Inkling QKVR prep kernels require CUDA", +) +@pytest.mark.parametrize("tokens", [9, 128]) +@torch.inference_mode() +def test_qkvr_prep_negative_slots_skip_cache_writes(tokens): + args = list(_make_inputs(is_local=True, tokens=tokens)) + args[17].fill_(-1) + args[19].fill_(-1) + conv_cache = args[11].clone() + key_cache = args[12].clone() + value_cache = args[13].clone() + + qkvr_prep.fused_qkvr_prep(*args) + + torch.testing.assert_close(args[11], conv_cache, rtol=0, atol=0) + torch.testing.assert_close(args[12], key_cache, rtol=0, atol=0) + torch.testing.assert_close(args[13], value_cache, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@torch.inference_mode() +def test_fused_sconv_negative_slots_skip_cache_writes(): + args = list(_make_inputs(is_local=True, tokens=9)) + args[17].fill_(-1) + cache = args[11].clone() + + sconv.fused_sconv( + args[0][:, 16 * 128 : 16 * 128 + 4 * 128].contiguous(), + args[1], + args[11], + args[14], + args[15], + args[16], + args[17], + args[18], + 0, + 128, + args[22], + ) + + torch.testing.assert_close(args[11], cache, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("rel_extent", "last_latency_rows", "first_throughput_rows"), + [(512, 8191, 8192), (1024, 2047, 2048)], +) +def test_rel_projection_schedule_crossover( + rel_extent, last_latency_rows, first_throughput_rows +): + assert not qkvr_prep.use_rel_proj_throughput(last_latency_rows, rel_extent) + assert qkvr_prep.use_rel_proj_throughput(first_throughput_rows, rel_extent) diff --git a/tests/models/inkling/test_sconv_metadata.py b/tests/models/inkling/test_sconv_metadata.py new file mode 100644 index 000000000000..70a091d382c8 --- /dev/null +++ b/tests/models/inkling/test_sconv_metadata.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Parity test for the fused sconv seq-metadata kernel. + +``sconv_seq_metadata`` fills the per-token ``seq_idx`` (owning request) and +``query_start`` (first x-row of that request) buffers in a single launch. Actual +tokens must match the searchsorted-based reference, while CUDA graph padding +must be initialized to safe zero values. +""" + +import pytest +import torch + +from vllm.models.inkling.nvidia.ops.sconv import sconv_seq_metadata +from vllm.models.inkling.nvidia.sconv_swa_attn import InklingSconvMetadataBuilder + +CASES = [ + # (query_lens, extra_pad_tokens) + ([1], 0), # bsz1 decode + ([1] * 8, 0), # uniform decode + ([1] * 8, 3), # uniform decode, padded tokens past the last request + ([2] * 4, 0), # uniform spec-decode + ([2048], 0), # single prefill + ([517, 1, 1, 33, 1, 128], 0), # mixed prefill/decode + ([517, 1, 1, 33, 1, 128], 5), # mixed, padded + ([1] * 500, 0), # many requests (deep binary search) +] + + +def _ref(query_start_loc: torch.Tensor, num_reqs: int, num_tokens: int): + cu_seqlens = query_start_loc[: num_reqs + 1].to(torch.int64) + token_idx = torch.arange(num_tokens, device=cu_seqlens.device, dtype=torch.int64) + seq_idx = (torch.searchsorted(cu_seqlens, token_idx, right=True) - 1).clamp( + max=num_reqs - 1 + ) + return seq_idx.to(torch.int32), cu_seqlens[seq_idx].to(torch.int32) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("query_lens,extra_pad", CASES) +def test_sconv_seq_metadata_matches_searchsorted(query_lens, extra_pad): + device = "cuda" + num_reqs = len(query_lens) + query_start_loc = torch.tensor( + [0] + list(torch.tensor(query_lens).cumsum(0)), dtype=torch.int32 + ).to(device) + num_actual_tokens = int(query_start_loc[-1]) + num_padded_tokens = num_actual_tokens + extra_pad + + ref_seq, ref_qs = _ref(query_start_loc, num_reqs, num_actual_tokens) + + seq_idx = torch.full((num_padded_tokens,), -1, dtype=torch.int32, device=device) + query_start = torch.full_like(seq_idx, -1) + sconv_seq_metadata( + query_start_loc, + num_reqs, + num_actual_tokens, + seq_idx, + query_start, + num_padded_tokens, + ) + + torch.testing.assert_close(seq_idx[:num_actual_tokens], ref_seq, rtol=0, atol=0) + torch.testing.assert_close(query_start[:num_actual_tokens], ref_qs, rtol=0, atol=0) + assert torch.count_nonzero(seq_idx[num_actual_tokens:]) == 0 + assert torch.count_nonzero(query_start[num_actual_tokens:]) == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_sconv_metadata_reuses_padded_static_buffers(): + device = torch.device("cuda") + builder = object.__new__(InklingSconvMetadataBuilder) + builder.seq_idx_buffer = torch.empty(8, dtype=torch.int32, device=device) + builder.query_start_buffer = torch.empty(8, dtype=torch.int32, device=device) + + class CommonMetadata: + num_reqs = 1 + num_actual_tokens = 8 + query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device) + query_start_loc_cpu = torch.tensor([0, 3], dtype=torch.int32) + block_table_tensor = torch.zeros((1, 1), dtype=torch.int32, device=device) + slot_mapping = torch.tensor( + [0, 1, 2, -1, -1, -1, -1, -1], + dtype=torch.int64, + device=device, + ) + + common = CommonMetadata() + first = builder.build(0, common) + pointers = tuple( + tensor.data_ptr() + for tensor in ( + first.block_table, + first.slot_mapping, + first.seq_idx, + first.query_start, + ) + ) + + common.query_start_loc[1] = 5 + common.query_start_loc_cpu[1] = 5 + common.slot_mapping[:5] = torch.arange(5, dtype=torch.int64, device=device) + second = builder.build(0, common) + + assert second.slot_mapping.shape == (8,) + assert second.seq_idx.shape == (8,) + assert second.query_start.shape == (8,) + assert pointers == tuple( + tensor.data_ptr() + for tensor in ( + second.block_table, + second.slot_mapping, + second.seq_idx, + second.query_start, + ) + ) + assert torch.count_nonzero(second.seq_idx[5:]) == 0 + assert torch.count_nonzero(second.query_start[5:]) == 0 + assert torch.all(second.slot_mapping[5:] == -1) diff --git a/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py b/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py index 4119e1d5e00d..251bc8c1d8a4 100644 --- a/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py +++ b/tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py @@ -5,6 +5,7 @@ from transformers import AutoModel from tests.models.utils import check_embeddings_close +from tests.utils import VLLM_PATH from vllm import TokensPrompt from vllm.config import PoolerConfig @@ -53,3 +54,53 @@ def test_embed_models(hf_runner, vllm_runner, model: str): name_1="vllm", tol=1e-2, ) + + +_RERANKER_HF_OVERRIDES = { + "architectures": ["Qwen3ForSequenceClassification"], + "classifier_from_token": ["no", "yes"], + "is_original_qwen3_reranker": True, +} +_RERANKER_TEMPLATE = VLLM_PATH / "examples/pooling/score/template/qwen3_reranker.jinja" + + +@pytest.mark.parametrize("model", ["Qwen/Qwen3-Reranker-0.6B"]) +@torch.inference_mode +def test_last_pool_score_chunked_prefill_matches_unchunked(vllm_runner, model: str): + """LAST-pooling score must not depend on whether the prompt is chunked. + + Regression test for a wrong-score bug where, under ``torch.compile`` (i.e. + not ``enforce_eager``), a query+document pair long enough to be split across + prefill chunks produced a corrupted last-token hidden state and thus a + wrong relevance score, while the same input scored correctly unchunked. + The existing pooling+chunked-prefill coverage runs ``enforce_eager=True`` + and so never exercised the compiled path. + """ + chat_template = _RERANKER_TEMPLATE.read_text() + query = "What organelle produces energy in the cell?" + # A long, matching document so query + doc exceeds the chunk size below. + document = ( + "The mitochondria is the powerhouse of the cell. It generates most of " + "the cell chemical energy through oxidative phosphorylation. " + ) * 400 + + def score_with(max_num_batched_tokens: int) -> float: + with vllm_runner( + model, + runner="pooling", + hf_overrides=_RERANKER_HF_OVERRIDES, + max_model_len=16384, + max_num_batched_tokens=max_num_batched_tokens, + enable_chunked_prefill=True, + enable_prefix_caching=False, + ) as vllm_model: + return vllm_model.score(query, document, chat_template=chat_template)[0] + + # Chunk size forces the prompt across several prefill chunks; the large + # budget keeps it in a single chunk as the reference. + chunked = score_with(2048) + unchunked = score_with(16384) + + assert chunked == pytest.approx(unchunked, abs=5e-2), ( + f"chunked score {chunked} diverged from unchunked {unchunked}" + ) diff --git a/tests/models/language/pooling/test_multi_vector_retrieval.py b/tests/models/language/pooling/test_multi_vector_retrieval.py index 3161271be091..d72bf2088180 100644 --- a/tests/models/language/pooling/test_multi_vector_retrieval.py +++ b/tests/models/language/pooling/test_multi_vector_retrieval.py @@ -12,7 +12,7 @@ "model", ["BAAI/bge-m3"], ) -@pytest.mark.parametrize("dtype", ["half"]) +@pytest.mark.parametrize("dtype", ["float"]) @torch.inference_mode def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype: str): with vllm_runner( @@ -20,13 +20,11 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype runner="pooling", pooler_config=PoolerConfig(task="token_embed"), max_model_len=None, + dtype=dtype, ) as vllm_model: vllm_outputs = vllm_model.token_embed(example_prompts) - with hf_runner( - model, - auto_cls=AutoModel, - ) as hf_model: + with hf_runner(model, auto_cls=AutoModel, dtype=dtype) as hf_model: tokenizer = hf_model.tokenizer hf_outputs = [] for prompt in example_prompts: @@ -34,7 +32,6 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype inputs = hf_model.wrap_device(inputs) output = hf_model.model(**inputs) embedding = output.last_hidden_state[0].float() - # normal hf_outputs.append(embedding.cpu()) for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): diff --git a/tests/models/language/pooling/test_token_classification.py b/tests/models/language/pooling/test_token_classification.py index 0fdeeea4b384..2e95fb8f70a4 100644 --- a/tests/models/language/pooling/test_token_classification.py +++ b/tests/models/language/pooling/test_token_classification.py @@ -272,6 +272,11 @@ def test_bert_for_masked_lm( if current_platform.is_rocm(): hf_model_kwargs["attn_implementation"] = "eager" + # Run hf_runner reference with "highest" fp32 precision to match + # default behvior of vLLM. This is needed on ROCm since the + # pooling tests set matmul precision to "high" in conftest.py + prev_matmul_precision = torch.get_float32_matmul_precision() + torch.set_float32_matmul_precision("highest") with hf_runner( model, dtype=dtype, @@ -285,6 +290,7 @@ def test_bert_for_masked_lm( inputs = hf_model.wrap_device(inputs) output = hf_model.model(**inputs) hf_outputs.append(softmax(output.logits[0])) + torch.set_float32_matmul_precision(prev_matmul_precision) # Compare the per-token vocabulary distributions position by position. for hf_output, vllm_output in zip(hf_outputs, vllm_outputs): diff --git a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py index 4eb4d03bfe5d..af12a14a45a0 100644 --- a/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py +++ b/tests/models/multimodal/processing/test_qwen2_5_omni_embed.py @@ -19,6 +19,7 @@ check_interleaved_audio_video, merge_interleaved_embeddings, ) +from vllm.multimodal.utils import set_mm_embedding_modality # Fake token IDs AUDIO_TOKEN_ID = 1001 @@ -27,6 +28,10 @@ TEXT_TOKEN_ID = 0 +def _mm_embed(shape: tuple[int, ...], value: float, modality: str) -> torch.Tensor: + return set_mm_embedding_modality(torch.full(shape, value), modality) + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -116,6 +121,24 @@ def test_interleaved(self): is_video, is_audio, is_video.sum().item(), is_audio.sum().item() ) + def test_multi_video_with_boundary_tokens(self): + """Two interleaved videos separated by boundary tokens → still True. + + use_audio_in_video expands each video into a local V/A span bounded by + non-pad tokens. A global density check would fail across those gaps. + """ + # [text][V A V A][text boundary][V A V A] + first_ids, _ = make_interleaved_seq([2, 2], [2, 2], text_prefix=1) + second_ids, _ = make_interleaved_seq([2, 2], [2, 2], text_prefix=0) + boundary = torch.tensor([TEXT_TOKEN_ID, TEXT_TOKEN_ID]) + input_ids = torch.cat([first_ids, boundary, second_ids]) + is_multimodal = (input_ids == VIDEO_TOKEN_ID) | (input_ids == AUDIO_TOKEN_ID) + is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID) + is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID) + assert check_interleaved_audio_video( + is_video, is_audio, is_video.sum().item(), is_audio.sum().item() + ) + def test_batched_non_interleaved_no_false_positive(self): """ Regression test for https://github.com/vllm-project/vllm/issues/35394. @@ -323,13 +346,16 @@ def test_interleaved_use_audio_in_video(self): # mm_embeds come in [video, audio] order (video feature first in # mm_features when positions are the same for use_audio_in_video) mm_embeds = [ - torch.full((video_n, hidden), video_val), - torch.full((audio_n, hidden), audio_val), + _mm_embed((video_n, hidden), video_val, "video"), + _mm_embed((audio_n, hidden), audio_val, "audio"), ] model, _ = make_mock_model(hidden) + # Modalities are attached on the embedding tensors (as in encoder gather). result = model.embed_input_ids( - input_ids, mm_embeds, is_multimodal=is_multimodal + input_ids, + mm_embeds, + is_multimodal=is_multimodal, ) video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0] @@ -362,8 +388,8 @@ def test_basic_interleaved(self): inputs_embeds = torch.zeros(len(input_ids), hidden) mm_embeds = [ - torch.full((num_video, hidden), 30.0), - torch.full((num_audio, hidden), 10.0), + _mm_embed((num_video, hidden), 30.0, "video"), + _mm_embed((num_audio, hidden), 10.0, "audio"), ] result = merge_interleaved_embeddings( @@ -372,8 +398,6 @@ def test_basic_interleaved(self): is_video, is_audio, is_multimodal, - num_video, - num_audio, ) video_pos = is_video.nonzero(as_tuple=True)[0] @@ -381,6 +405,62 @@ def test_basic_interleaved(self): assert result[video_pos].allclose(torch.full((num_video, hidden), 30.0)) assert result[audio_pos].allclose(torch.full((num_audio, hidden), 10.0)) + def test_image_and_video_mixed(self): + """Image embeddings must not be misclassified as video.""" + hidden = 4 + # [text][I I][V A V A] + tokens = ( + [TEXT_TOKEN_ID] * 2 + + [IMAGE_TOKEN_ID] * 2 + + [VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, VIDEO_TOKEN_ID, AUDIO_TOKEN_ID] + ) + input_ids = torch.tensor(tokens) + is_multimodal = ( + (input_ids == IMAGE_TOKEN_ID) + | (input_ids == VIDEO_TOKEN_ID) + | (input_ids == AUDIO_TOKEN_ID) + ) + is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID) + is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID) + is_image = is_multimodal & (input_ids == IMAGE_TOKEN_ID) + + inputs_embeds = torch.zeros(len(input_ids), hidden) + mm_embeds = [ + _mm_embed((2, hidden), 20.0, "image"), + _mm_embed((2, hidden), 30.0, "video"), + _mm_embed((2, hidden), 10.0, "audio"), + ] + result = merge_interleaved_embeddings( + inputs_embeds, + mm_embeds, + is_video, + is_audio, + is_multimodal, + ) + assert result[is_image.nonzero(as_tuple=True)[0]].allclose( + torch.full((2, hidden), 20.0) + ) + assert result[is_video.nonzero(as_tuple=True)[0]].allclose( + torch.full((2, hidden), 30.0) + ) + assert result[is_audio.nonzero(as_tuple=True)[0]].allclose( + torch.full((2, hidden), 10.0) + ) + + def test_missing_modality_raises(self): + hidden = 2 + input_ids, is_multimodal = make_interleaved_seq([2], [2]) + is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID) + is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID) + with pytest.raises(ValueError, match="Missing modality"): + merge_interleaved_embeddings( + torch.zeros(len(input_ids), hidden), + [torch.zeros(2, hidden), torch.zeros(2, hidden)], + is_video, + is_audio, + is_multimodal, + ) + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/models/registry.py b/tests/models/registry.py index 8d6b45d4796c..fa23e14d9654 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -330,6 +330,12 @@ def check_available_online( "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B", min_transformers_version="5.9.0", ), + "InklingForCausalLM": _HfExamplesInfo( + "thinkingmachines/Inkling-NVFP4", + tokenizer_mode="inkling", + trust_remote_code=True, + max_model_len=4096, + ), "InternLM2ForCausalLM": _HfExamplesInfo( "internlm/internlm2-chat-7b", trust_remote_code=True ), @@ -956,6 +962,12 @@ def check_available_online( "HuggingFaceM4/Idefics3-8B-Llama3", extras={"tiny": "HuggingFaceTB/SmolVLM-256M-Instruct"}, ), + "InklingForConditionalGeneration": _HfExamplesInfo( + "thinkingmachines/Inkling-NVFP4", + tokenizer_mode="inkling", + trust_remote_code=True, + max_model_len=4096, + ), "IsaacForConditionalGeneration": _HfExamplesInfo( "PerceptronAI/Isaac-0.1", trust_remote_code=True, @@ -1366,7 +1378,15 @@ def check_available_online( "HuggingFaceTB/SmolVLM2-2.2B-Instruct" ), "Step3VLForConditionalGeneration": _HfExamplesInfo( - "stepfun-ai/step3", trust_remote_code=True + "stepfun-ai/step3", + trust_remote_code=True, + max_transformers_version="5.3", + transformers_version_reason={ + "hf": ( + "Transformers v5.4 removed the ignore_keys param from " + "validate_rope(); vLLM has vendored the config and is unaffected" + ) + }, ), "StepVLForConditionalGeneration": _HfExamplesInfo( "stepfun-ai/Step3-VL-10B", trust_remote_code=True @@ -1457,6 +1477,12 @@ def check_available_online( is_available_online=False, use_original_num_layers=True, # DSpark backbone requires all layers ), + "Gemma4DSparkModel": _HfExamplesInfo( + "google/gemma-4-12B-it", + speculative_model="deepseek-ai/dspark_gemma4_12b_block7", + is_available_online=False, + use_original_num_layers=True, + ), # [Eagle] "EagleCohereForCausalLM": _HfExamplesInfo( "/host/engines/cohere-moe", @@ -1636,6 +1662,13 @@ def check_available_online( "tencent/Hy3-preview", speculative_model="tencent/Hy3-preview", ), + "InklingMTPModel": _HfExamplesInfo( + "thinkingmachines/Inkling-NVFP4", + speculative_model="thinkingmachines/Inkling-NVFP4", + tokenizer_mode="inkling", + trust_remote_code=True, + max_model_len=4096, + ), "LongCatFlashMTPModel": _HfExamplesInfo( "meituan-longcat/LongCat-Flash-Chat", trust_remote_code=True, diff --git a/tests/parser/engine/test_delegating_replay.py b/tests/parser/engine/test_delegating_replay.py index 1f2086fda3a2..b922e71e3982 100644 --- a/tests/parser/engine/test_delegating_replay.py +++ b/tests/parser/engine/test_delegating_replay.py @@ -88,6 +88,11 @@ def _discover_pairings() -> list[_PairingInfo]: if cfg.name not in _BUILDERS: missing_builders.append(f"{engine_cls.__name__} (config.name={cfg.name!r})") continue + if cfg.name == "inkling": + # Inkling uses typed structural blocks and opts out of token-id + # terminal matching; combined-parser replay coverage lives in + # test_inkling.py. + continue parser_cls = type( f"_Delegating{engine_cls.__name__}", diff --git a/tests/parser/engine/test_inkling.py b/tests/parser/engine/test_inkling.py new file mode 100644 index 000000000000..c0b9877db34a --- /dev/null +++ b/tests/parser/engine/test_inkling.py @@ -0,0 +1,506 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the engine-based Inkling parser. + +Inkling output is a sequence of typed content blocks delimited by dedicated +special tokens; the tool-call payload is ``{"name":...,"args":{...}}`` +between ``<|content_invoke_tool_json|>`` and ``<|end_message|>``. The +cases mirror the Rust unified parser's tests +(``rust/src/parser/src/unified/inkling.rs``) where applicable. +""" + +import json + +import pytest + +from tests.parser.engine.conftest import make_mock_tokenizer +from tests.parser.engine.streaming_helpers import ( + collect_content, + collect_function_name, + collect_tool_arguments, +) +from vllm.parser.engine.parser_engine_config import ParserState +from vllm.parser.inkling import InklingParser, _inkling_arg_converter +from vllm.parser.parser_manager import ParserManager + +MSG_MODEL = "<|message_model|>" +TEXT_START = "<|content_text|>" +THINK_START = "<|content_thinking|>" +TOOL_JSON = "<|content_invoke_tool_json|>" +TOOL_TEXT = "<|content_invoke_tool_text|>" +TOOL_ERROR = "<|content_tool_error|>" +END_MESSAGE = "<|end_message|>" +END_SAMPLING = "<|content_model_end_sampling|>" + +_TML_VOCAB = { + MSG_MODEL: 200001, + TEXT_START: 200004, + END_SAMPLING: 200006, + THINK_START: 200008, + END_MESSAGE: 200010, + TOOL_ERROR: 200022, + TOOL_JSON: 200049, + TOOL_TEXT: 200057, +} + + +@pytest.fixture +def mock_tokenizer(): + return make_mock_tokenizer(_TML_VOCAB) + + +@pytest.fixture +def parser(mock_tokenizer): + return InklingParser(mock_tokenizer) + + +def _tool_block(name: str, args: str) -> str: + return f'{TOOL_JSON}{{"name":"{name}","args":{args}}}{END_MESSAGE}' + + +_MARKERS = sorted(_TML_VOCAB, key=len, reverse=True) + + +def _tokenize(text: str) -> list[tuple[int, str]]: + """Tokenize like the real stream: markers are atomic special tokens, + plain text becomes one token per character (matching the mock + tokenizer's ``chr``-based decode).""" + tokens: list[tuple[int, str]] = [] + i = 0 + while i < len(text): + for marker in _MARKERS: + if text.startswith(marker, i): + tokens.append((_TML_VOCAB[marker], marker)) + i += len(marker) + break + else: + tokens.append((ord(text[i]), text[i])) + i += 1 + return tokens + + +def _stream(parser, request, text: str, chunk_size: int): + """Stream production-shaped deltas: ``chunk_size`` tokens per delta, + with delta_token_ids covering every token (specials and text).""" + tokens = _tokenize(text) + results = [] + previous_text = "" + previous_token_ids: list[int] = [] + for start in range(0, len(tokens), chunk_size): + batch = tokens[start : start + chunk_size] + delta_text = "".join(t for _, t in batch) + delta_token_ids = [tid for tid, _ in batch] + current_text = previous_text + delta_text + current_token_ids = previous_token_ids + delta_token_ids + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=tuple(previous_token_ids), + current_token_ids=tuple(current_token_ids), + delta_token_ids=tuple(delta_token_ids), + request=request, + ) + results.append((delta, current_text)) + previous_text = current_text + previous_token_ids = current_token_ids + finish = parser.finish_streaming() + if finish is not None: + results.append((finish, text)) + return results + + +def _stream_text_only(parser, request, text: str, chunk_size: int): + """Stream text-only deltas (no token ids), chunked at arbitrary + character boundaries — exercises the text-lexing fallback path, + including markers split across chunks.""" + results = [] + previous_text = "" + for start in range(0, len(text), chunk_size): + delta_text = text[start : start + chunk_size] + current_text = previous_text + delta_text + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=delta_text, + previous_token_ids=(), + current_token_ids=(), + delta_token_ids=(), + request=request, + ) + results.append((delta, current_text)) + previous_text = current_text + finish = parser.finish_streaming() + if finish is not None: + results.append((finish, text)) + return results + + +def _collect_reasoning(results) -> str: + return "".join(d.reasoning for d, _ in results if d and d.reasoning) + + +class TestArgConverter: + def test_complete_wrapper(self): + raw = '{"name":"get_weather","args":{"city":"SF"}}' + assert _inkling_arg_converter(raw, False) == '{"city":"SF"}' + + def test_partial_before_args(self): + assert _inkling_arg_converter('{"name":"get_w', True) == "" + + def test_partial_inside_args(self): + raw = '{"name":"x","args":{"a":1' + assert _inkling_arg_converter(raw, True) == '{"a":1' + + def test_prefix_stability(self): + full = '{"name":"x","args":{"a":{"b":[1,2]},"c":"d"}}' + prev = "" + for end in range(len(full)): + out = _inkling_arg_converter(full[:end], True) + assert out.startswith(prev) or prev.startswith(out) or not prev + if out.startswith(prev): + prev = out + + def test_args_value_appearing_in_name(self): + raw = '{"name":"args","args":{"k":1}}' + assert _inkling_arg_converter(raw, False) == '{"k":1}' + + def test_whitespace_tolerated(self): + raw = '{ "name" : "x" , "args" : {"a": 1} }' + assert _inkling_arg_converter(raw, False) == '{"a": 1}' + + def test_missing_args_defaults_empty(self): + assert _inkling_arg_converter('{"name":"x"}', False) == "{}" + + def test_non_object_args_rejected(self): + with pytest.raises(ValueError, match="JSON object"): + _inkling_arg_converter('{"name":"x","args":[1]}', False) + + +class TestNonStreaming: + def test_plain_text(self, parser, mock_request): + reasoning, content, tools = parser.parse( + f"{TEXT_START}hello world{END_MESSAGE}", mock_request + ) + assert reasoning is None + assert content == "hello world" + assert tools is None + + def test_reasoning_text_tool(self, parser, mock_request): + text = ( + f"{THINK_START}I should check the weather.{END_MESSAGE}" + f"{MSG_MODEL}{TEXT_START}Let me check.{END_MESSAGE}" + f"{MSG_MODEL}" + _tool_block("get_weather", '{"city":"SF"}') + ) + reasoning, content, tools = parser.parse(text, mock_request) + assert reasoning == "I should check the weather." + assert content == "Let me check." + assert [t.name for t in tools] == ["get_weather"] + assert json.loads(tools[0].arguments) == {"city": "SF"} + + def test_tool_header_name_is_not_visible_content(self, parser, mock_request): + text = "get_weather" + _tool_block("get_weather", '{"city":"SF"}') + _, content, tools = parser.parse(text, mock_request) + assert content is None + assert [tool.name for tool in tools] == ["get_weather"] + + def test_parallel_tool_calls(self, parser, mock_request): + text = _tool_block("a", "{}") + MSG_MODEL + _tool_block("b", '{"x":[1,2]}') + _, _, tools = parser.parse(text, mock_request) + assert [t.name for t in tools] == ["a", "b"] + assert json.loads(tools[0].arguments) == {} + assert json.loads(tools[1].arguments) == {"x": [1, 2]} + + def test_nested_args(self, parser, mock_request): + args = '{"q":{"deep":{"list":[{"k":"v"}]}},"s":"a}b"}' + _, _, tools = parser.parse(_tool_block("f", args), mock_request) + assert json.loads(tools[0].arguments) == json.loads(args) + + def test_invoke_tool_text_is_visible_text(self, parser, mock_request): + reasoning, content, tools = parser.parse( + f"{TOOL_TEXT}do something{END_MESSAGE}", mock_request + ) + assert content == "do something" + assert tools is None + + def test_tool_error_is_visible_text(self, parser, mock_request): + _, content, tools = parser.parse(f"{TOOL_ERROR}boom{END_MESSAGE}", mock_request) + assert content == "boom" + assert tools is None + + def test_end_sampling_closes_blocks(self, parser, mock_request): + reasoning, content, _ = parser.parse( + f"{THINK_START}hm{END_MESSAGE}{MSG_MODEL}{TEXT_START}hi{END_SAMPLING}", + mock_request, + ) + assert reasoning == "hm" + assert content == "hi" + + def test_multiple_reasoning_blocks_concatenate(self, parser, mock_request): + text = ( + f"{THINK_START}one{END_MESSAGE}" + f"{MSG_MODEL}{TEXT_START}mid{END_MESSAGE}" + f"{MSG_MODEL}{THINK_START}two{END_MESSAGE}" + ) + reasoning, content, _ = parser.parse(text, mock_request) + assert reasoning == "onetwo" + assert content == "mid" + + def test_text_after_tool_call(self, parser, mock_request): + text = _tool_block("f", "{}") + f"{MSG_MODEL}{TEXT_START}done{END_MESSAGE}" + _, content, tools = parser.parse(text, mock_request) + assert [t.name for t in tools] == ["f"] + assert content == "done" + + def test_incomplete_tool_call_at_eos(self, parser, mock_request): + # Engine convention: best-effort with what arrived. (The Rust + # parser instead errors with "incomplete Inkling tool call".) + _, _, tools = parser.parse( + f'{TOOL_JSON}{{"name":"d","args":{{"k":"v"', mock_request + ) + assert [t.name for t in tools] == ["d"] + + def test_prose_marker_without_token_ids_is_structural(self, parser, mock_request): + # Inkling opts into text-lexer terminal recognition so held-back + # structural marker text from the detokenizer is still parsed. + _, content, _ = parser.parse( + f"{TEXT_START}see {TEXT_START} token{END_MESSAGE}", mock_request + ) + assert content == "see token" + + +class TestStreaming: + @pytest.mark.parametrize("chunk_size", [1, 3, 7, 64, 4096]) + def test_chunk_invariance_tool_call(self, mock_tokenizer, mock_request, chunk_size): + parser = InklingParser(mock_tokenizer) + text = f"{TEXT_START}Check this.{END_MESSAGE}{MSG_MODEL}" + _tool_block( + "get_weather", '{"city":"San Francisco"}' + ) + results = _stream(parser, mock_request, text, chunk_size) + assert collect_content(results) == "Check this." + assert collect_function_name(results) == "get_weather" + assert json.loads(collect_tool_arguments(results)) == {"city": "San Francisco"} + + @pytest.mark.parametrize("chunk_size", [1, 3, 7, 64]) + def test_chunk_invariance_tool_call_text_only( + self, mock_tokenizer, mock_request, chunk_size + ): + # Same case through the text-lexing fallback (no token ids), + # with markers split at arbitrary character boundaries. + parser = InklingParser(mock_tokenizer) + text = f"{TEXT_START}Check this.{END_MESSAGE}{MSG_MODEL}" + _tool_block( + "get_weather", '{"city":"San Francisco"}' + ) + results = _stream_text_only(parser, mock_request, text, chunk_size) + assert collect_content(results) == "Check this." + assert collect_function_name(results) == "get_weather" + assert json.loads(collect_tool_arguments(results)) == {"city": "San Francisco"} + + @pytest.mark.parametrize("chunk_size", [1, 5, 11]) + def test_chunk_invariance_reasoning(self, mock_tokenizer, mock_request, chunk_size): + parser = InklingParser(mock_tokenizer) + text = ( + f"{THINK_START}thinking...{END_MESSAGE}" + f"{MSG_MODEL}{TEXT_START}answer{END_MESSAGE}" + ) + results = _stream(parser, mock_request, text, chunk_size) + assert _collect_reasoning(results) == "thinking..." + assert collect_content(results) == "answer" + + def test_split_marker_held_across_chunks(self, parser, mock_request): + # Mirrors Rust `inkling_streaming_holds_split_markers`. + text = f"{TEXT_START}hello{END_MESSAGE}" + results = _stream_text_only(parser, mock_request, text, 9) + assert collect_content(results) == "hello" + + def test_name_streams_before_args_complete(self, parser, mock_request): + # Feed only up to the name's closing quote — the name delta must + # already be emitted before any args arrive. + prefix = f'{TOOL_JSON}{{"name":"get_weather",' + results = _stream(parser, mock_request, prefix, 4096) + assert collect_function_name(results) == "get_weather" + + def test_combined_parser_reasoning_to_tool_handoff_uses_text_markers( + self, mock_tokenizer, mock_request + ): + parser_cls = ParserManager.get_parser( + tool_parser_name="inkling", + reasoning_parser_name="inkling", + enable_auto_tools=True, + ) + parser = parser_cls(mock_tokenizer, []) + + first = parser.parse_delta( + THINK_START, + [_TML_VOCAB[THINK_START]], + mock_request, + prompt_token_ids=[_TML_VOCAB[MSG_MODEL]], + finished=False, + ) + assert first is None + + second = parser.parse_delta( + "thinking", + [ord(c) for c in "thinking"], + mock_request, + finished=False, + ) + assert second is not None + assert second.reasoning == "thinking" + + # Mirrors the DelegatingParser handoff after reasoning closes: the + # tool pass receives reconstructed text that starts at the Inkling + # tool marker, while the token-id slice has already moved past it. + body = ( + "get_weather" + f'{TOOL_JSON}{{"name":"get_weather","args":{{"city":"Seattle"}}}}' + f"{END_MESSAGE}" + ) + third = parser.parse_delta( + body, + [_TML_VOCAB[END_MESSAGE], _TML_VOCAB[END_SAMPLING]], + mock_request, + finished=True, + ) + assert third is not None + assert third.tool_calls + assert third.tool_calls[0].function.name == "get_weather" + assert third.tool_calls[0].function.arguments == '{"city":"Seattle"}' + assert TOOL_JSON not in ((third.content or "") + (third.reasoning or "")) + + def test_streamed_args_are_object_only(self, parser, mock_request): + # The streamed `arguments` must be the bare args object, never + # the `{"name":...}` wrapper. + text = _tool_block("f", '{"a":1}') + results = _stream(parser, mock_request, text, 3) + args = collect_tool_arguments(results) + assert json.loads(args) == {"a": 1} + assert "name" not in args + + @pytest.mark.parametrize("chunk_size", [1, 9]) + def test_parallel_calls_streaming(self, mock_tokenizer, mock_request, chunk_size): + parser = InklingParser(mock_tokenizer) + text = _tool_block("a", '{"i":1}') + MSG_MODEL + _tool_block("b", '{"i":2}') + results = _stream(parser, mock_request, text, chunk_size) + indexed: dict[int, dict[str, str]] = {} + for delta, _ in results: + if not (delta and delta.tool_calls): + continue + for tc in delta.tool_calls: + slot = indexed.setdefault(tc.index, {"name": "", "args": ""}) + if tc.function and tc.function.name: + slot["name"] = tc.function.name + if tc.function and tc.function.arguments: + slot["args"] += tc.function.arguments + assert indexed[0]["name"] == "a" + assert indexed[1]["name"] == "b" + assert json.loads(indexed[0]["args"]) == {"i": 1} + assert json.loads(indexed[1]["args"]) == {"i": 2} + + +class TestPromptSeededState: + def test_prompt_ending_in_thinking_starts_reasoning(self, parser, mock_request): + parser.adjust_initial_state_from_prompt([200001, _TML_VOCAB[THINK_START]]) + assert parser._engine.state == ParserState.REASONING + + def test_prompt_ending_in_text_starts_content(self, parser): + parser.adjust_initial_state_from_prompt([200001, _TML_VOCAB[TEXT_START]]) + assert parser._engine.state == ParserState.CONTENT + + def test_generation_prompt_tail_starts_message_header(self, parser): + parser.adjust_initial_state_from_prompt( + [_TML_VOCAB[END_MESSAGE], _TML_VOCAB[MSG_MODEL]] + ) + assert parser._engine.state == ParserState.MESSAGE_HEADER + + def test_generation_prompt_header_hides_tool_name(self, parser, mock_request): + text = "get_weather" + _tool_block("get_weather", '{"city":"SF"}') + delta = parser.parse_delta( + text, + [token_id for token_id, _ in _tokenize(text)], + mock_request, + prompt_token_ids=[_TML_VOCAB[END_MESSAGE], _TML_VOCAB[MSG_MODEL]], + finished=True, + ) + assert delta is not None + assert delta.content is None + assert delta.tool_calls[0].function.name == "get_weather" + + +class TestToolCallFiltering: + """Inkling equivalents of the generic tool-call-filtering replay tests + (Inkling is excluded from those in test_replay.py: its structural + role/kind tokens and shared block-end token don't fit the generic + reasoning/tool split model).""" + + def test_skip_tool_parsing_round_trip(self, mock_tokenizer, mock_request): + # First pass (reasoning adapter, skip_tool_parsing): reasoning is + # classified as reasoning while tool markup survives in content; + # second pass (tool adapter) re-extracts the calls from it. + text = ( + f"{THINK_START}plan{END_MESSAGE}{MSG_MODEL}" + + _tool_block("f", '{"a":1}') + + MSG_MODEL + + _tool_block("g", '{"b":[2]}') + ) + first = InklingParser(mock_tokenizer) + first.skip_tool_parsing = True + reasoning, content = first.extract_reasoning(text, mock_request) + assert reasoning == "plan" + assert content.count(TOOL_JSON) == 2 + + second = InklingParser(mock_tokenizer) + result = second.extract_tool_calls_from_content(content, mock_request) + assert result.tools_called + assert [tc.function.name for tc in result.tool_calls] == ["f", "g"] + assert json.loads(result.tool_calls[0].function.arguments) == {"a": 1} + assert json.loads(result.tool_calls[1].function.arguments) == {"b": [2]} + + @pytest.fixture + def none_request(self, mock_request): + mock_request.tools = [{"type": "function", "function": {"name": "f"}}] + mock_request.tool_choice = "none" + return mock_request + + def test_tool_choice_none_non_streaming(self, mock_tokenizer, none_request): + parser = InklingParser(mock_tokenizer) + text = ( + f"{THINK_START}plan{END_MESSAGE}" + f"{MSG_MODEL}{TEXT_START}visible{END_MESSAGE}" + f"{MSG_MODEL}" + _tool_block("f", '{"a":1}') + ) + reasoning, content, tools = parser.parse(text, none_request) + assert reasoning == "plan" + assert content == "visible" + assert not tools + + def test_tool_choice_none_streaming(self, mock_tokenizer, none_request): + parser = InklingParser(mock_tokenizer) + text = f"{TEXT_START}visible{END_MESSAGE}{MSG_MODEL}" + _tool_block( + "f", '{"a":1}' + ) + results = _stream(parser, none_request, text, 3) + assert collect_content(results) == "visible" + assert all(not (d and d.tool_calls) for d, _ in results) + + +class TestRegisteredAdapters: + def test_adapters_resolve(self): + from vllm.reasoning import ReasoningParserManager + from vllm.tool_parsers import ToolParserManager + + reasoning_cls = ReasoningParserManager.get_reasoning_parser("inkling") + tool_cls = ToolParserManager.get_tool_parser("inkling") + assert reasoning_cls._parser_engine_cls is InklingParser + assert tool_cls._parser_engine_cls is InklingParser + assert tool_cls.supports_required_and_named is False + + def test_adapter_round_trip(self, mock_tokenizer, mock_request): + from vllm.tool_parsers import ToolParserManager + + tool_cls = ToolParserManager.get_tool_parser("inkling") + adapter = tool_cls(mock_tokenizer) + result = adapter.extract_tool_calls(_tool_block("f", '{"a":1}'), mock_request) + assert result.tools_called + assert result.tool_calls[0].function.name == "f" + assert json.loads(result.tool_calls[0].function.arguments) == {"a": 1} diff --git a/tests/parser/engine/test_replay.py b/tests/parser/engine/test_replay.py index 1626135f7866..064e1b623679 100644 --- a/tests/parser/engine/test_replay.py +++ b/tests/parser/engine/test_replay.py @@ -68,6 +68,10 @@ def _discover_parsers() -> list[_ParserInfo]: if cfg.name not in _BUILDERS: missing_builders.append(f"{obj.__name__} (config.name={cfg.name!r})") continue + if cfg.name == "inkling": + # Inkling opts out of token-id terminal matching and has typed + # structural blocks; its replay coverage lives in test_inkling.py. + continue tool_end = cfg.token_id_terminals.get("TOOL_END") if not tool_end: raise RuntimeError( @@ -299,7 +303,12 @@ def test_adjust_request_disables_skip_special_tokens(self, parser_cls, sample): (p.parser_cls, s, p.think_end, p.tool_start) for p in _PARSERS for s in p.samples - if s.expected_tool_calls and s.expected_reasoning + if s.expected_tool_calls + and s.expected_reasoning + # Inkling's typed-block format (structural role/kind tokens, shared + # block-end token) doesn't fit the generic reasoning/tool split + # below; its filtering modes are covered in test_inkling.py instead. + and p.name != "inkling" ] diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index 533c64408f78..f24e3318f975 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -33,6 +33,7 @@ DeepSeekV32Parser, Gemma4Parser, Glm47MoeParser, + InklingParser, KimiK2Parser, MinimaxM2Parser, NemotronV3Parser, @@ -948,6 +949,77 @@ def _build_kimi_k2( ] +# ── Inkling (typed content blocks, JSON tool payloads) ─────────────────── + +_TML_VOCAB: dict[str, int] = { + "<|message_model|>": 200001, + "<|content_text|>": 200004, + "<|content_model_end_sampling|>": 200006, + "<|content_thinking|>": 200008, + "<|end_message|>": 200010, + "<|content_tool_error|>": 200022, + "<|content_invoke_tool_json|>": 200049, + "<|content_invoke_tool_text|>": 200057, +} + + +def _inkling_block( + segs: list[tuple[str, bool]], + kind_token: str, + body: str, +) -> None: + """Append one Inkling content block; blocks after the first start with + the ``<|message_model|>`` role token (the first block continues the + generation prompt directly).""" + if segs: + segs.append(("<|message_model|>", True)) + segs.append((kind_token, True)) + if body: + segs.append((body, False)) + segs.append(("<|end_message|>", True)) + + +def _inkling_segments(scenario: Scenario) -> list[tuple[str, bool]]: + segs: list[tuple[str, bool]] = [] + if scenario.reasoning is not None: + _inkling_block(segs, "<|content_thinking|>", scenario.reasoning) + if scenario.tool_calls is not None and not scenario.tool_calls: + _inkling_block(segs, "<|content_invoke_tool_json|>", "") + if scenario.content is not None: + _inkling_block(segs, "<|content_text|>", scenario.content) + if scenario.tool_calls: + for tc in scenario.tool_calls: + args = json.dumps(tc.arguments, ensure_ascii=False, separators=(",", ":")) + payload = f'{{"name":"{tc.name}","args":{args}}}' + _inkling_block(segs, "<|content_invoke_tool_json|>", payload) + return segs + + +def _build_inkling(scenario: Scenario, validate: bool = True) -> Sample: + prompt_token_ids = None + if scenario.after_tool_response: + # Prompt ends with a closed tool-response block and the + # generation-prompt role token. + prompt_token_ids = [ + _TML_VOCAB["<|end_message|>"], + _TML_VOCAB["<|message_model|>"], + ] + sample = _make_sample( + sample_id=f"inkling-{scenario.id}", + description=scenario.description, + vocab=_TML_VOCAB, + segments=_inkling_segments(scenario), + expected_reasoning=scenario.reasoning, + expected_content=_qwen3_expected_content(scenario), + expected_tool_calls=_expected_tc(scenario), + tools=_expected_tools(scenario), + prompt_token_ids=prompt_token_ids, + ) + if validate: + _validate_sample(sample, InklingParser) + return sample + + # ── Registry and public API ────────────────────────────────────────── _BUILDERS: dict[str, Any] = { @@ -960,6 +1032,7 @@ def _build_kimi_k2( "glm47_moe": _build_glm47_moe, "kimi_k2": _build_kimi_k2, "qwen3": _build_qwen3, + "inkling": _build_inkling, } diff --git a/tests/quantization/test_fp8.py b/tests/quantization/test_fp8.py index 571e180d6c74..abb995257e34 100644 --- a/tests/quantization/test_fp8.py +++ b/tests/quantization/test_fp8.py @@ -17,6 +17,10 @@ from vllm.model_executor.kernels.linear.scaled_mm import ( MarlinFP8ScaledMMLinearKernel, ) +from vllm.model_executor.layers.attention.attention import ( + Attention, + set_default_quant_scales, +) from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.quantization.fp8 import ( Fp8Config, @@ -24,6 +28,7 @@ Fp8LinearMethod, Fp8MoEMethod, ) +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, ) @@ -422,6 +427,42 @@ def test_fp8_reloading( method.process_weights_after_loading(layer) +@pytest.mark.parametrize("source", ["checkpoint", "runtime_calc"]) +def test_kv_cache_scale_sync_to_host_copies(source): + """Test device-to-host sync of the k/v quantization scales, for both the + checkpoint-load and runtime-calc paths that produce them. + """ + layer = torch.nn.Module() + set_default_quant_scales(layer, register_buffer=True) + layer.kv_cache_dtype = "fp8" + + if source == "checkpoint": + # Scales come from the checkpoint, so runtime calc is disabled. + layer.calculate_kv_scales = False + method = BaseKVCacheMethod(quant_config=None) + method.create_weights(layer) + # 0.3 stays != 1.0 even after the fp8_fnuz x2 rescale. + checkpoint_scale = torch.tensor(0.3, dtype=torch.float32) + layer.k_scale.weight_loader(layer.k_scale, checkpoint_scale) + layer.v_scale.weight_loader(layer.v_scale, checkpoint_scale) + method.process_weights_after_loading(layer) + else: + # First forward computes distinct, non-unity scales from live k/v. + layer.calculate_kv_scales = True + query = torch.full((4, 8), 10.0) + key = torch.full((4, 8), 60.0) + value = torch.full((4, 8), 50.0) + Attention.calc_kv_scales(layer, query, key, value) + + assert layer._k_scale_float != 1.0 + assert layer._v_scale_float != 1.0 + # Host copy must mirror both the float and the device scale tensor. + assert layer._k_scale_cpu.item() == pytest.approx(layer._k_scale_float) + assert layer._v_scale_cpu.item() == pytest.approx(layer._v_scale_float) + assert layer._k_scale_cpu.item() == pytest.approx(layer._k_scale.item()) + assert layer._v_scale_cpu.item() == pytest.approx(layer._v_scale.item()) + + @pytest.mark.skipif( not is_quant_method_supported("fp8"), reason="FP8 is not supported on this GPU type.", diff --git a/tests/quantization/test_humming_ignore.py b/tests/quantization/test_humming_ignore.py new file mode 100644 index 000000000000..53554b1af312 --- /dev/null +++ b/tests/quantization/test_humming_ignore.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for humming's is_layer_skipped handling of compressed-tensors +regex ("re:") ignore entries. + +Regression test: compressed-tensors checkpoints (e.g. Kimi-K2.6) list ignored +layers as regex patterns prefixed with "re:" (e.g. "re:vision_tower.*"). +humming previously substring-matched these literals, so ignored layers were +never skipped and got (incorrectly) quantized -- which then diverged between +the real weight-load path (falls back to unquantized) and the dummy-load path +(stays quantized), breaking any consumer that byte-compares the two. +""" + +import pytest + +from vllm.model_executor.layers.quantization.humming import HummingConfig + +# The real ignore list from Kimi-K2.6's compressed-tensors quantization_config. +K26_IGNORE = [ + "re:.*self_attn.*", + "re:.*shared_experts.*", + r"re:.*mlp\.(gate|up|gate_up|down)_proj.*", + "re:.*lm_head.*", + "re:vision_tower.*", + "re:mm_projector.*", +] + + +@pytest.mark.parametrize( + "prefix,expected", + [ + # ignored layers (must be skipped -> stay unquantized) + ("vision_tower.encoder.blocks.0.mlp.fc0", True), + ("vision_tower.encoder.blocks.16.wo", True), + ("mm_projector.linear_1", True), + ("language_model.model.layers.0.self_attn.o_proj", True), + ("language_model.model.layers.0.mlp.gate_up_proj", True), + ("language_model.model.layers.0.mlp.down_proj", True), + ("language_model.model.layers.3.mlp.shared_experts.gate_proj", True), + ("model.lm_head", True), + # routed experts are NOT ignored -> must be quantized + ("language_model.model.layers.3.mlp.experts.383.up_proj", False), + ("language_model.model.layers.3.mlp.experts.0.down_proj", False), + ], +) +def test_is_layer_skipped_regex_ignore(prefix, expected): + cfg = HummingConfig(full_config={"ignore": K26_IGNORE}) + assert cfg.is_layer_skipped({"ignore": K26_IGNORE}, prefix) is expected + + +def test_plain_substring_entries_still_work(): + # bitsandbytes-style modules_to_not_convert use plain substrings. + cfg = HummingConfig() + cfg_dict = {"modules_to_not_convert": ["lm_head", "vision"]} + assert cfg.is_layer_skipped(cfg_dict, "model.vision.encoder.fc") is True + assert cfg.is_layer_skipped(cfg_dict, "model.layers.0.mlp.up_proj") is False diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 995df7946008..3c21441ed65b 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -16,7 +16,11 @@ Fp8PerTensorOnlineLinearMethod, Fp8PerTensorOnlineMoEMethod, ) +from vllm.model_executor.layers.quantization.online.nvfp4 import ( + Nvfp4OnlineMoEMethod, +) from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe @pytest.mark.skipif( @@ -145,6 +149,38 @@ def check_model(model): print(outputs[0][1]) +@pytest.mark.skipif( + not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and has_flashinfer_trtllm_fused_moe() + ), + reason="nvfp4_per_token needs a Blackwell (SM100) GPU + FlashInfer TRTLLM MoE.", +) +def test_online_nvfp4_per_token_moe(vllm_runner, monkeypatch) -> None: + """Online NVFP4 quantizes the MoE and leaves dense layers unquantized.""" + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + + with vllm_runner( + "ibm-granite/granite-3.0-1b-a400m-base", + quantization="nvfp4_per_token", + enforce_eager=True, + ) as llm: + + def check_model(model): + layer = model.model.layers[0] + assert isinstance( + layer.block_sparse_moe.experts._quant_method, Nvfp4OnlineMoEMethod + ) + assert isinstance( + layer.self_attn.o_proj.quant_method, UnquantizedLinearMethod + ) + + llm.apply_model(check_model) + outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4) + print(outputs[0][1]) + + @pytest.mark.skipif( not is_quant_method_supported("fp8"), reason="FP8 is not supported on this GPU type.", diff --git a/tests/renderers/test_inkling.py b/tests/renderers/test_inkling.py new file mode 100644 index 000000000000..2753417b42f3 --- /dev/null +++ b/tests/renderers/test_inkling.py @@ -0,0 +1,484 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Golden tests for the native Inkling renderer encoding. + +The message fixtures mirror the Rust renderer's fixture tests +(``rust/src/chat/src/renderer/inkling/tests.rs``) so the two frontends stay +in token-level parity. Image blocks emit the bare ``<|content_image|>`` marker; +``InklingMultiModalProcessor`` inserts the per-patch placeholder run. +""" + +import pytest + +from vllm.renderers.inkling import ( + InklingRenderer, + _HfBackedTmlTokenizer, + _resolve_reasoning_effort, +) +from vllm.renderers.inkling_encoding import ( + SPECIAL_TOKEN_SPELLINGS, + render_inkling_messages, +) +from vllm.renderers.params import ChatParams + + +@pytest.fixture() +def should_do_global_cleanup_after_test() -> bool: + # These tests touch no distributed or device state; the global + # cleanup fixture is unnecessary (and trips a torch MPS allocator + # assert on macOS dev machines). + return False + + +# One id per special token, mirroring the real Inkling vocab layout; plain +# text encodes one token per character so decoded output is exact. +_SPECIAL_VOCAB = { + "<|message_user|>": 200000, + "<|message_model|>": 200001, + "<|message_system|>": 200002, + "<|message_tool|>": 200003, + "<|content_text|>": 200004, + "<|content_image|>": 200005, + "<|content_model_end_sampling|>": 200006, + "<|content_thinking|>": 200008, + "<|end_message|>": 200010, + "<|content_audio_input|>": 200020, + # The HF vocab spells CONTENT_XML as an unused slot. + "<|unused_200024|>": 200024, + "<|audio_end|>": 200043, + "<|content_invoke_tool_json|>": 200049, +} + +_ID_TO_SPECIAL = {v: k for k, v in _SPECIAL_VOCAB.items()} + + +class FakeHfTokenizer: + def get_vocab(self): + return dict(_SPECIAL_VOCAB) + + def encode(self, text, add_special_tokens=False): + assert not add_special_tokens + return [ord(ch) for ch in text] + + +def decode(token_ids): + return "".join( + _ID_TO_SPECIAL.get(tid, chr(tid) if tid < 200000 else f"<{tid}>") + for tid in token_ids + ) + + +@pytest.fixture +def inkling_tokenizer(): + return _HfBackedTmlTokenizer(FakeHfTokenizer()) + + +def render_text(inkling_tokenizer, messages, **kwargs): + return decode(render_inkling_messages(messages, inkling_tokenizer, **kwargs)) + + +class TestRustFixtureParity: + def test_tool_round_trip(self, inkling_tokenizer): + messages = [ + { + "role": "assistant", + "reasoning_content": "think", + "content": "answer", + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "get_weather", + "arguments": '{"city":"SF"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + assert render_text( + inkling_tokenizer, messages, add_generation_prompt=False + ) == ( + "<|message_model|><|content_thinking|>think<|end_message|>" + "<|message_model|><|content_text|>answer<|end_message|>" + "<|message_model|>get_weather<|content_invoke_tool_json|>" + '{"name":"get_weather","args":{"city":"SF"}}<|end_message|>' + "<|content_model_end_sampling|>" + "<|message_tool|>get_weather<|content_text|>sunny<|end_message|>" + ) + + def test_tool_declare(self, inkling_tokenizer): + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather information", + "parameters": { + "type": "object", + "required": ["city"], + "properties": {"city": {"type": "string"}}, + }, + }, + } + ] + messages = [ + { + "role": "developer", + "content": "rules", + "tools": [ + { + "type": "function", + "function": { + "name": "local_tool", + "parameters": {"z": 1, "a": {"b": 2}}, + }, + } + ], + }, + {"role": "user", "content": "hi"}, + ] + assert render_text(inkling_tokenizer, messages, tools=tools) == ( + "<|message_system|>tool_declare<|unused_200024|>" + '[{"description":"Get weather information","name":"get_weather",' + '"parameters":{"properties":{"city":{"type":"string"}},' + '"required":["city"],"type":"object"},"type":"function"},' + '{"description":"","name":"local_tool",' + '"parameters":{"a":{"b":2},"z":1},"type":"function"}]' + "<|end_message|>" + "<|message_system|><|content_text|>rules<|end_message|>" + "<|message_user|><|content_text|>hi<|end_message|>" + "<|message_model|>" + ) + + def test_text_image(self, inkling_tokenizer): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "look"}, + {"type": "image_url", "image_url": "data:image/png;base64,"}, + ], + } + ] + # Multimodal preprocessing expands this bare marker after rendering. + assert render_text(inkling_tokenizer, messages) == ( + "<|message_user|><|content_text|>look<|end_message|>" + "<|message_user|><|content_image|><|end_message|>" + "<|message_model|>" + ) + + +class TestRenderingSemantics: + @pytest.mark.parametrize( + ("value", "expected"), + [ + ("none", 0.0), + ("minimal", 0.1), + ("low", 0.2), + ("medium", 0.7), + ("high", 0.9), + ("xhigh", 0.99), + ("max", 0.99), + (None, 0.9), + (0.8, 0.8), + (True, None), + ("invalid", None), + ], + ) + def test_resolve_reasoning_effort(self, value, expected): + assert _resolve_reasoning_effort(value) == expected + + def test_generation_prompt_default(self, inkling_tokenizer): + text = render_text(inkling_tokenizer, [{"role": "user", "content": "hi"}]) + assert text.endswith("<|end_message|><|message_model|>") + + def test_developer_folds_into_system(self, inkling_tokenizer): + assert ( + render_text( + inkling_tokenizer, + [{"role": "developer", "content": "be nice"}], + add_generation_prompt=False, + ) + == "<|message_system|><|content_text|>be nice<|end_message|>" + ) + + def test_empty_string_content_skipped(self, inkling_tokenizer): + assert ( + render_text( + inkling_tokenizer, + [{"role": "user", "content": ""}], + add_generation_prompt=False, + ) + == "" + ) + + def test_empty_reasoning_skipped(self, inkling_tokenizer): + assert ( + render_text( + inkling_tokenizer, + [{"role": "assistant", "reasoning_content": "", "content": "hi"}], + add_generation_prompt=False, + ) + == "<|message_model|><|content_text|>hi<|end_message|>" + "<|content_model_end_sampling|>" + ) + + def test_reasoning_field(self, inkling_tokenizer): + messages = [ + { + "role": "assistant", + "reasoning": "think", + "content": "answer", + } + ] + assert render_text( + inkling_tokenizer, messages, add_generation_prompt=False + ) == ( + "<|message_model|><|content_thinking|>think<|end_message|>" + "<|message_model|><|content_text|>answer<|end_message|>" + "<|content_model_end_sampling|>" + ) + + def test_audio_part(self, inkling_tokenizer): + messages = [ + { + "role": "user", + "content": [{"type": "input_audio", "input_audio": {}}], + } + ] + assert render_text( + inkling_tokenizer, messages, add_generation_prompt=False + ) == ("<|message_user|><|content_audio_input|><|audio_end|><|end_message|>") + + def test_tool_response_name_from_message(self, inkling_tokenizer): + assert ( + render_text( + inkling_tokenizer, + [{"role": "tool", "name": "my_tool", "content": "ok"}], + add_generation_prompt=False, + ) + == "<|message_tool|>my_tool<|content_text|>ok<|end_message|>" + ) + + def test_tool_call_args_object_form(self, inkling_tokenizer): + messages = [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "x", + "function": { + "name": "f", + # dict-form arguments, unsorted keys + "arguments": {"z": 1, "a": 2}, + }, + } + ], + } + ] + assert render_text( + inkling_tokenizer, messages, add_generation_prompt=False + ) == ( + "<|message_model|>f<|content_invoke_tool_json|>" + '{"name":"f","args":{"a":2,"z":1}}<|end_message|>' + "<|content_model_end_sampling|>" + ) + + def test_tool_call_empty_args(self, inkling_tokenizer): + messages = [ + { + "role": "assistant", + "tool_calls": [{"id": "x", "function": {"name": "f", "arguments": ""}}], + } + ] + assert render_text( + inkling_tokenizer, messages, add_generation_prompt=False + ) == ( + "<|message_model|>f<|content_invoke_tool_json|>" + '{"name":"f","args":{}}<|end_message|>' + "<|content_model_end_sampling|>" + ) + + def test_tool_call_non_object_args_rejected(self, inkling_tokenizer): + messages = [ + { + "role": "assistant", + "tool_calls": [ + {"id": "x", "function": {"name": "f", "arguments": "[1]"}} + ], + } + ] + with pytest.raises(TypeError, match="decode to an object"): + render_inkling_messages( + messages, inkling_tokenizer, add_generation_prompt=False + ) + + def test_unsupported_role_rejected(self, inkling_tokenizer): + with pytest.raises(ValueError, match="unsupported Inkling message role"): + render_inkling_messages( + [{"role": "narrator", "content": "hi"}], inkling_tokenizer + ) + + +class TestReasoningEffort: + def test_frontend_defaults_to_high(self, inkling_tokenizer): + renderer = InklingRenderer.__new__(InklingRenderer) + renderer._inkling_tokenizer = inkling_tokenizer + + text = decode( + renderer._render( + [ + {"role": "system", "content": "rules"}, + {"role": "user", "content": "hi"}, + ], + ChatParams(), + ) + ) + + assert text.startswith( + "<|message_system|><|content_text|>rules<|end_message|>" + "<|message_system|><|content_text|>Thinking effort level: 0.9" + "<|end_message|>" + ) + + @pytest.mark.parametrize("value", ["none", 0, 0.0, -0.0]) + def test_zero_effort_has_one_canonical_spelling(self, inkling_tokenizer, value): + renderer = InklingRenderer.__new__(InklingRenderer) + renderer._inkling_tokenizer = inkling_tokenizer + + text = decode( + renderer._render( + [{"role": "user", "content": "hi"}], + ChatParams(chat_template_kwargs={"reasoning_effort": value}), + ) + ) + + assert text.startswith( + "<|message_system|><|content_text|>Thinking effort level: 0.0" + "<|end_message|>" + ) + + def test_emits_one_effort_after_initial_prefix(self, inkling_tokenizer): + effort_block = ( + "<|message_system|><|content_text|>Thinking effort level: 0.7" + "<|end_message|>" + ) + text = render_text( + inkling_tokenizer, + [ + {"role": "system", "content": "rules"}, + {"role": "developer", "content": "policy"}, + {"role": "user", "content": "user1"}, + {"role": "assistant", "content": "assistant1"}, + {"role": "user", "content": "user2"}, + ], + tools=[{"type": "function", "function": {"name": "f"}}], + reasoning_effort=0.7, + ) + + assert text.count(effort_block) == 1 + assert ( + text.index("tool_declare") + < text.index("rules") + < text.index("policy") + < text.index(effort_block) + < text.index("user1") + < text.index("assistant1") + < text.index("user2") + ) + + def test_renders_after_tool_declare(self, inkling_tokenizer): + tools = [{"type": "function", "function": {"name": "f"}}] + text = render_text( + inkling_tokenizer, + [{"role": "user", "content": "hi"}], + tools=tools, + reasoning_effort=0.8, + ) + declare_end = text.index("<|end_message|>") + len("<|end_message|>") + assert text[declare_end:].startswith( + "<|message_system|><|content_text|>Thinking effort level: 0.8" + "<|end_message|>" + ) + + @pytest.mark.parametrize( + ("value", "expected"), + [(0.2, "0.2"), (0.7, "0.7"), (0.9, "0.9"), (0.99, "0.99")], + ) + def test_formats_at_most_two_decimals(self, inkling_tokenizer, value, expected): + text = render_text( + inkling_tokenizer, + [{"role": "user", "content": "hi"}], + reasoning_effort=value, + ) + assert f"Thinking effort level: {expected}<|end_message|>" in text + + def test_absent_emits_no_block(self, inkling_tokenizer): + text = render_text(inkling_tokenizer, [{"role": "user", "content": "hi"}]) + assert "Thinking effort" not in text + + @pytest.mark.parametrize("value", [0.9900001, 1, 1.5, -0.1]) + def test_out_of_range_rejected(self, inkling_tokenizer, value): + with pytest.raises(ValueError, match="must be in"): + render_inkling_messages( + [{"role": "user", "content": "hi"}], + inkling_tokenizer, + reasoning_effort=value, + ) + + +class TestSpecialTokenResolution: + def test_missing_special_token_raises(self): + class IncompleteTokenizer(FakeHfTokenizer): + def get_vocab(self): + vocab = dict(_SPECIAL_VOCAB) + del vocab["<|content_invoke_tool_json|>"] + return vocab + + with pytest.raises(ValueError, match="missing special tokens"): + _HfBackedTmlTokenizer(IncompleteTokenizer()) + + def test_semantic_spelling_preferred(self): + class SemanticSpellingTokenizer(FakeHfTokenizer): + def get_vocab(self): + vocab = dict(_SPECIAL_VOCAB) + del vocab["<|unused_200024|>"] + vocab["<|content_xml|>"] = 200024 + return vocab + + inkling_tokenizer = _HfBackedTmlTokenizer(SemanticSpellingTokenizer()) + tools = [{"type": "function", "function": {"name": "f"}}] + ids = render_inkling_messages( + [{"role": "user", "content": "hi"}], inkling_tokenizer, tools=tools + ) + assert 200024 in ids + + def test_all_spellings_covered(self): + # Every semantic token must resolve from the reference vocab. + for token, spellings in SPECIAL_TOKEN_SPELLINGS.items(): + assert any(s in _SPECIAL_VOCAB for s in spellings), token + + +def test_render_inkling_messages_direct_protocol(): + """The encoding core only needs the structural tokenizer protocol.""" + + class ProtocolTokenizer: + def encode_text(self, text): + return [ord(c) for c in text] + + def encode_special(self, token): + return { + "<|message_user|>": 200000, + "<|message_model|>": 200001, + "<|content_text|>": 200004, + "<|content_model_end_sampling|>": 200006, + "<|end_message|>": 200010, + }[token] + + ids = render_inkling_messages( + [{"role": "user", "content": "hi"}], + ProtocolTokenizer(), + add_generation_prompt=True, + ) + assert ids == [200000, 200004, ord("h"), ord("i"), 200010, 200001] diff --git a/tests/test_config.py b/tests/test_config.py index 6f2be35ae86f..a785297997b1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -211,6 +211,26 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1( ), True, ), + ( + SimpleNamespace( + model="thinkingmachines/Inkling", + architectures=["InklingForCausalLM"], + runner_type="generate", + is_moe=True, + is_quantized=False, + ), + True, + ), + ( + SimpleNamespace( + model="thinkingmachines/Inkling", + architectures=["InklingForConditionalGeneration"], + runner_type="generate", + is_moe=True, + is_quantized=False, + ), + True, + ), ( SimpleNamespace( model="mistralai/Mixtral-8x7B-Instruct-v0.1", diff --git a/tests/test_jit_monitor.py b/tests/test_jit_monitor.py index 9f3285ddec0e..50261a479d92 100644 --- a/tests/test_jit_monitor.py +++ b/tests/test_jit_monitor.py @@ -320,6 +320,26 @@ def compile_fn(*args, **kwargs): with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"): cute.compile(lambda: None, "arg", option=True) + def test_subscripted_compile_is_monitored(self): + """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work.""" + + class FakeCompileCallable: + def __getitem__(self, options): + return self + + def __call__(self, *args, **kwargs): + return "compiled" + + with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()): + import cutlass.cute as cute + + jit_monitor.activate() + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: + result = cute.compile[("opt_level", 3)](lambda: None, "arg") + + assert result == "compiled" + warning_once.assert_called_once() + class TestTileLangHook: def test_jit_kernel_logs_warning(self): diff --git a/tests/v1/attention/test_attention_backends.py b/tests/v1/attention/test_attention_backends.py index abbdeac1ef8c..3110e4b4ee17 100644 --- a/tests/v1/attention/test_attention_backends.py +++ b/tests/v1/attention/test_attention_backends.py @@ -646,6 +646,33 @@ class MockLayer: assert impl.get_xqa_bmm1_scale(MockLayer, torch.float8_e4m3fn) == 3.0 +@pytest.mark.skipif( + AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, + reason="FlashInfer is not available.", +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_flashinfer_attention_sinks_refreshed_after_reload(dtype): + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + source_sinks = torch.tensor([1.0, 2.0], dtype=dtype) + impl = object.__new__(flashinfer_backend.FlashInferImpl) + impl._sinks_source = source_sinks + impl.sinks = source_sinks + + impl.process_weights_after_loading(dtype) + + assert impl.sinks is not None + sinks_ptr = impl.sinks.data_ptr() + assert impl.sinks.dtype == torch.float32 + torch.testing.assert_close(impl.sinks, source_sinks.float()) + + source_sinks.copy_(torch.tensor([3.0, 4.0], dtype=dtype)) + impl.process_weights_after_loading(dtype) + + assert impl.sinks.data_ptr() == sinks_ptr + torch.testing.assert_close(impl.sinks, source_sinks.float()) + + @pytest.mark.skipif( AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST, reason="FlashInfer is not available.", diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 3afa8eca9255..aca09c3db3d6 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -42,7 +42,10 @@ ) from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported -from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm.v1.kv_cache_interface import ( + KVQuantMode, + MLAAttentionSpec, +) BACKENDS_TO_TEST = [ AttentionBackendEnum.CUTLASS_MLA, @@ -55,6 +58,31 @@ DEVICE_TYPE = current_platform.device_type + +@pytest.mark.parametrize( + ("cache_dtype", "expected_quant_mode"), + [ + ("auto", KVQuantMode.NONE), + ("fp8_ds_mla", KVQuantMode.FP8_PER_TENSOR), + ], +) +def test_mla_kv_cache_spec_uses_layer_cache_dtype( + cache_dtype: str, expected_quant_mode: KVQuantMode +): + layer = SimpleNamespace(kv_cache_dtype=cache_dtype, head_size=576) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(block_size=64), model_config=None + ) + + spec = MLAAttention.get_kv_cache_spec(layer, vllm_config) + + assert isinstance(spec, MLAAttentionSpec) + assert spec.cache_dtype_str == cache_dtype + assert spec.kv_quant_mode == expected_quant_mode + if cache_dtype == "fp8_ds_mla": + assert spec.page_size_bytes == 64 * 656 + + # Remove sm100 backends from the list if not using sm100 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10: BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA) @@ -77,6 +105,45 @@ BACKENDS_TO_TEST.remove(AttentionBackendEnum.TOKENSPEED_MLA) +def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch): + layer = MLAAttention.__new__(MLAAttention) + torch.nn.Module.__init__(layer) + layer.kv_lora_rank = 2 + layer.num_heads = 2 + layer.qk_nope_head_dim = 3 + layer.v_head_dim = 4 + layer.kv_b_proj = torch.nn.Module() + layer.kv_b_proj.weight = torch.nn.Parameter( + torch.arange(28.0, dtype=torch.float16).reshape(14, 2) + ) + layer.kv_b_proj.quant_method = None + layer.is_aiter_triton_fp4_bmm_enabled = False + layer.is_aiter_triton_fp8_bmm_enabled = False + layer.quant_config = None + layer.layer_name = "test" + + monkeypatch.setattr( + mla_attention_module, "set_default_quant_scales", lambda *_, **__: None + ) + + with torch.no_grad(): + layer.process_weights_after_loading(torch.float32) + assert isinstance(layer.W_UV, torch.nn.Parameter) + assert isinstance(layer.W_UK_T, torch.nn.Parameter) + w_uv_ptr = layer.W_UV.data_ptr() + w_uk_t_ptr = layer.W_UK_T.data_ptr() + old_w_uv = layer.W_UV.clone() + old_w_uk_t = layer.W_UK_T.clone() + + layer.kv_b_proj.weight.add_(100) + layer.process_weights_after_loading(torch.float32) + + assert layer.W_UV.data_ptr() == w_uv_ptr + assert layer.W_UK_T.data_ptr() == w_uk_t_ptr + torch.testing.assert_close(layer.W_UV, old_w_uv + 100) + torch.testing.assert_close(layer.W_UK_T, old_w_uk_t + 100) + + # Filtered per-test via validate_configuration (capability/deps/dims). PREFILL_BACKENDS_TO_TEST = [ MLAPrefillBackendEnum.FLASH_ATTN, diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index 8677197b7a54..d82397591f77 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -125,8 +125,8 @@ def test_explicit_backend_import_error_raises(self): @pytest.mark.parametrize( ("qk_nope_head_dim", "v_head_dim"), - [(128, 128), (192, 256)], - ids=["deepseek", "glm"], + [(128, 128), (192, 256), (64, 128)], + ids=["deepseek", "glm", "mistral_s4"], ) def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int): try: diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index d88f107ecee8..6ec9fc25ed3e 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -40,10 +40,16 @@ ) from vllm.v1.attention.backends.mla.flashmla_sparse import ( FlashMLASparseBackend, + FlashMLASparseImpl, + FlashMLASparseMetadata, + FlashMLASparseMetadataBuilder, triton_convert_req_index_to_global_index, ) from vllm.v1.attention.backends.mla.indexer import split_indexer_prefill_chunks -from vllm.v1.attention.backends.utils import split_prefill_chunks +from vllm.v1.attention.backends.utils import ( + split_decodes_and_prefills, + split_prefill_chunks, +) from vllm.v1.attention.ops import flashmla SPARSE_BACKEND_BATCH_SPECS = { @@ -721,6 +727,120 @@ def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_s torch.testing.assert_close(result, reference_result, rtol=0, atol=0) +@pytest.mark.skipif( + torch.cuda.get_device_capability() < (9, 0), + reason="FlashMLASparseBackend requires CUDA 9.0 or higher", +) +def test_triton_convert_rejects_req_id_longer_than_token_indices(): + """Guard against the #47327 regression: the kernel grid is sized by + req_id but the output is allocated like token_indices, so a full-batch + req_id combined with an MQA-subset token_indices wrote past the end of + the output buffer. The wrapper must reject the length mismatch instead + of corrupting memory.""" + device = torch.device(DEVICE_TYPE) + num_topk_tokens = 128 + block_size = 64 + block_table = torch.arange(40, dtype=torch.int32, device=device).view(4, 10) + + # Full batch: 2 decode tokens + 10 prefill tokens + req_id_full = torch.tensor( + [0, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3], dtype=torch.int32, device=device + ) + num_mqa_tokens = 2 + token_indices = torch.randint( + 0, + block_size * 10, + (num_mqa_tokens, num_topk_tokens), + dtype=torch.int32, + device=device, + ) + + with pytest.raises(AssertionError, match="must cover the same tokens"): + triton_convert_req_index_to_global_index( + req_id_full, + block_table, + token_indices, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=num_topk_tokens, + ) + + # The sliced call is the intended usage and must match the reference. + result = triton_convert_req_index_to_global_index( + req_id_full[:num_mqa_tokens], + block_table, + token_indices, + BLOCK_SIZE=block_size, + NUM_TOPK_TOKENS=num_topk_tokens, + ) + reference = _triton_convert_reference_impl( + req_id_full[:num_mqa_tokens], + block_table, + token_indices, + block_size, + num_topk_tokens, + ) + torch.testing.assert_close(result, reference, rtol=0, atol=0) + + +@pytest.mark.skipif( + torch.cuda.get_device_capability() < (9, 0), + reason="FlashMLASparseBackend requires CUDA 9.0 or higher", +) +def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): + """Guard against the #47327 regression: when the dense-MHA prefill split + is active, forward_mqa only receives the leading decode tokens, but + _forward_bf16_kv passed the full-batch req_id_per_token to the index + conversion, making it write past the end of its output buffer. The call + site must slice req_id_per_token to the MQA tokens.""" + device = torch.device(DEVICE_TYPE) + num_topk_tokens = 128 + block_size = 64 + num_batch_tokens = 12 + num_mqa_tokens = 2 + + attn_metadata = SimpleNamespace( + req_id_per_token=torch.tensor( + [0, 1] + [2] * 5 + [3] * 5, dtype=torch.int32, device=device + ), + block_table=torch.arange(40, dtype=torch.int32, device=device).view(4, 10), + block_size=block_size, + ) + assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens + + q = torch.zeros(num_mqa_tokens, 4, 576, dtype=torch.bfloat16, device=device) + kv_cache = torch.zeros(40 * block_size, 576, dtype=torch.bfloat16, device=device) + topk_indices = torch.randint( + 0, + block_size * 10, + (num_mqa_tokens, num_topk_tokens), + dtype=torch.int32, + device=device, + ) + + captured = {} + + def _stub_kernel(q, kv, indices, lengths): + captured["indices"] = indices + return torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype, device=q.device) + + stub_impl = SimpleNamespace(_bf16_flash_mla_kernel=_stub_kernel) + + out = FlashMLASparseImpl._forward_bf16_kv( + stub_impl, q, kv_cache, topk_indices, attn_metadata + ) + + assert out.shape[0] == num_mqa_tokens + assert captured["indices"].shape[0] == num_mqa_tokens + reference = _triton_convert_reference_impl( + attn_metadata.req_id_per_token[:num_mqa_tokens], + attn_metadata.block_table, + topk_indices, + block_size, + num_topk_tokens, + ) + torch.testing.assert_close(captured["indices"], reference, rtol=0, atol=0) + + @pytest.mark.parametrize( "seq_lens,max_buf,expected", [ @@ -1124,3 +1244,203 @@ def test_triton_convert_returns_valid_counts(): ) assert isinstance(result_only, torch.Tensor) torch.testing.assert_close(result_only, result, rtol=0, atol=0) + + +def test_flashmla_cache_dtype_aliases_use_ds_layout(): + from vllm.model_executor.layers.attention.mla_attention import ( + _canonicalize_sparse_mla_kv_cache_dtype, + ) + + # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer + # stores kv_cache_dtype, so they cannot bypass the gate. + for alias in ("fp8", "fp8_e4m3"): + assert ( + _canonicalize_sparse_mla_kv_cache_dtype(FlashMLASparseBackend, alias) + == "fp8_ds_mla" + ) + + +def test_flashmla_fp8_metadata_reuses_common_batch_split(): + builder = SimpleNamespace( + device=torch.device(DEVICE_TYPE), + vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=8)), + ) + common_metadata = SimpleNamespace( + num_actual_tokens=1, + seq_lens_cpu_upper_bound=torch.tensor([1]), + query_start_loc_cpu=torch.tensor([0, 1]), + block_table_tensor=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + ) + metadata = FlashMLASparseMetadata( + num_reqs=1, + max_query_len=1, + max_seq_len=1, + num_actual_tokens=1, + query_start_loc=torch.tensor([0, 1], device=DEVICE_TYPE), + slot_mapping=torch.tensor([0], device=DEVICE_TYPE), + block_table=torch.zeros(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + req_id_per_token=torch.zeros(1, dtype=torch.int32, device=DEVICE_TYPE), + num_decodes=0, + num_prefills=1, + num_decode_tokens=0, + ) + + fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode( + builder, common_metadata, metadata + ) + + assert fp8_metadata.num_decodes == 0 + assert fp8_metadata.num_prefills == 1 + assert fp8_metadata.num_decode_tokens == 0 + assert fp8_metadata.num_prefill_tokens == 1 + + +def test_flashmla_common_metadata_requires_uniform_decodes(): + common_metadata = SimpleNamespace( + max_query_len=3, + num_reqs=3, + num_actual_tokens=6, + query_start_loc_cpu=torch.tensor([0, 1, 3, 6]), + is_prefilling=None, + ) + + split = split_decodes_and_prefills( + common_metadata, + decode_threshold=128, + require_uniform=FlashMLASparseMetadataBuilder.require_uniform_decodes, + ) + + assert split == (1, 2, 1, 5) + + +def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch): + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.flashmla_sparse.get_mla_metadata", + lambda: (object(), None), + ) + builder = SimpleNamespace( + device=torch.device(DEVICE_TYPE), + dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE), + max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE), + ) + query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660]) + common_metadata = SimpleNamespace( + num_actual_tokens=660, + query_start_loc_cpu=query_start_loc_cpu, + seq_lens=torch.arange(7, device=DEVICE_TYPE), + ) + metadata = FlashMLASparseMetadata( + num_reqs=7, + max_query_len=110, + max_seq_len=110, + num_actual_tokens=660, + query_start_loc=query_start_loc_cpu.to(DEVICE_TYPE), + slot_mapping=torch.arange(660, device=DEVICE_TYPE), + block_table=torch.zeros(7, 1, dtype=torch.int32, device=DEVICE_TYPE), + req_id_per_token=torch.zeros(660, dtype=torch.int32, device=DEVICE_TYPE), + num_decodes=7, + num_prefills=0, + num_decode_tokens=660, + ) + + fp8_metadata = FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode( + builder, common_metadata, metadata + ) + + assert fp8_metadata.num_decodes == 6 + assert fp8_metadata.num_decode_tokens == 660 + assert fp8_metadata.decode is not None + assert fp8_metadata.decode.decode_query_len == 110 + torch.testing.assert_close( + fp8_metadata.decode.seq_lens, torch.arange(6, device=DEVICE_TYPE) + ) + + +@pytest.mark.parametrize("use_mixed_batch", [False, True]) +def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: bool): + num_decode_tokens = 2 + num_batch_tokens = 5 + q = torch.empty(num_decode_tokens, 2, 3, device=DEVICE_TYPE) + topk_indices = torch.empty(num_decode_tokens, 4, device=DEVICE_TYPE) + kernel_q_shapes = [] + + def convert_indices(*args, **kwargs): # noqa: ARG001 + assert not kwargs.get("HAS_PREFILL_WORKSPACE", False) + if not kwargs.get("return_valid_counts", False): + return topk_indices + valid_counts = torch.full( + (num_decode_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE + ) + return topk_indices, valid_counts + + monkeypatch.setattr( + "vllm.v1.attention.backends.mla.flashmla_sparse." + "triton_convert_req_index_to_global_index", + convert_indices, + ) + + def run_kernel(**kwargs): + kernel_q_shapes.append(kwargs["q"].shape) + return kwargs["q"][..., :1], None + + if use_mixed_batch: + fp8_metadata = FlashMLASparseMetadata.FP8KernelMetadata( + scheduler_metadata=object(), # type: ignore[arg-type] + dummy_block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + cache_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), + ) + else: + FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode + fp8_metadata = FP8Meta( + num_decodes=1, + num_prefills=1, + num_decode_tokens=num_decode_tokens, + num_prefill_tokens=num_batch_tokens - num_decode_tokens, + decode=FP8Meta.Decode( + seq_lens=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), + kernel_metadata=object(), # type: ignore[arg-type] + decode_query_len=num_decode_tokens, + ), + prefill=FP8Meta.Prefill( + request_ids=torch.empty( + num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE + ), + workspace_starts=torch.empty(1, dtype=torch.int32, device=DEVICE_TYPE), + chunks=[], + ), + ) + metadata = SimpleNamespace( + fp8_extra_metadata=fp8_metadata, + fp8_use_mixed_batch=use_mixed_batch, + num_actual_tokens=num_batch_tokens, + req_id_per_token=torch.empty( + num_batch_tokens, dtype=torch.int32, device=DEVICE_TYPE + ), + block_table=torch.empty(1, 1, dtype=torch.int32, device=DEVICE_TYPE), + block_size=64, + ) + impl = SimpleNamespace( + kv_cache_dtype="fp8_ds_mla", + topk_indices_buffer=topk_indices, + num_heads=2, + kv_lora_rank=1, + _fp8_flash_mla_kernel=run_kernel, + ) + impl._forward_fp8_kv_mixed_batch = MethodType( + FlashMLASparseImpl._forward_fp8_kv_mixed_batch, impl + ) + impl._forward_fp8_kv_separate_prefill_decode = MethodType( + FlashMLASparseImpl._forward_fp8_kv_separate_prefill_decode, impl + ) + + output, lse = FlashMLASparseImpl.forward_mqa( + impl, + q, + torch.empty(0, device=DEVICE_TYPE), + metadata, + None, + ) + + assert kernel_q_shapes == [(1, num_decode_tokens, 2, 3)] + assert output.shape == (num_decode_tokens, 2, 1) + assert lse is None diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 157170400b0d..da1e0c5e76c5 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -43,6 +43,65 @@ pytestmark = pytest.mark.cpu_test +def test_make_scheduled_encoder_input_stats_output_embeddings(): + scheduler = create_scheduler() + mm_features = [ + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="image", + identifier="image-0", + mm_position=PlaceholderRange(offset=0, length=196), + ), + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="video", + identifier="video-0", + mm_position=PlaceholderRange(offset=200, length=196), + ), + MultiModalFeatureSpec( + data=MultiModalKwargsItem.dummy(), + modality="audio", + identifier="audio-0", + mm_position=PlaceholderRange(offset=400, length=49), + ), + ] + scheduler.requests["req"] = Mock(mm_features=mm_features) + + stats = scheduler._make_scheduled_encoder_input_stats({"req": [0, 1, 2]}) + + assert stats is not None + assert stats.num_inputs == 3 + assert stats.output_tokens == 441 + + +def test_scheduled_encoder_input_stats_disabled_without_iteration_logging( + monkeypatch: pytest.MonkeyPatch, +): + scheduler = create_scheduler() + make_stats = Mock(side_effect=AssertionError("stats should not be computed")) + monkeypatch.setattr(scheduler, "_make_scheduled_encoder_input_stats", make_stats) + + scheduler_output = scheduler.schedule() + + make_stats.assert_not_called() + assert scheduler_output.scheduled_encoder_input_stats is None + + +def test_scheduled_encoder_input_stats_disabled_without_log_stats( + monkeypatch: pytest.MonkeyPatch, +): + scheduler = create_scheduler() + scheduler.log_stats = False + scheduler.observability_config.enable_logging_iteration_details = True + make_stats = Mock(side_effect=AssertionError("stats should not be computed")) + monkeypatch.setattr(scheduler, "_make_scheduled_encoder_input_stats", make_stats) + + scheduler_output = scheduler.schedule() + + make_stats.assert_not_called() + assert scheduler_output.scheduled_encoder_input_stats is None + + def test_add_requests(): scheduler = create_scheduler() requests = create_requests(num_requests=10) @@ -108,6 +167,29 @@ def test_schedule(enable_prefix_caching: bool, prompt_logprobs: int | None): assert scheduler.running[i] == request +def test_scheduler_stats_route_to_existing_output_client(): + scheduler = create_scheduler() + request = create_requests(num_requests=1)[0] + request.client_index = 1 + scheduler.add_request(request) + + scheduler_output = scheduler.schedule() + model_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[1000]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + engine_core_outputs = scheduler.update_from_output(scheduler_output, model_output) + + assert 0 not in engine_core_outputs + assert engine_core_outputs[1].scheduler_stats is not None + assert len(engine_core_outputs[1].outputs) == 1 + + def test_schedule_multimodal_requests(): scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf") mm_positions = [[PlaceholderRange(offset=i, length=100)] for i in range(10)] diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index f856d91b6395..742aafd3890e 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -8,6 +8,8 @@ import os import threading +from contextlib import nullcontext +from unittest.mock import patch import pytest import torch @@ -15,6 +17,54 @@ os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1" +def test_piecewise_capture_builds_fresh_metadata_for_both_passes(): + from vllm.config import CUDAGraphMode + from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, + ) + + manager = CudaGraphManager.__new__(CudaGraphManager) + desc = BatchExecutionDescriptor(CUDAGraphMode.PIECEWISE, 8, None) + manager.device = torch.device("cpu") + manager._capture_descs = {CUDAGraphMode.PIECEWISE: [desc]} + manager._graphs_captured = False + manager.use_breakable_cg = True + + create_calls = [] + forward_calls = [] + + def create_forward_fn(desc_arg, warmup): + assert desc_arg == desc + metadata = {"layer": object()} + create_calls.append((warmup, metadata)) + + def forward_fn(cg_mode): + assert metadata + forward_calls.append((warmup, cg_mode, metadata)) + + return forward_fn + + with ( + patch( + "vllm.v1.worker.gpu.cudagraph_utils.graph_capture", + return_value=nullcontext(), + ), + patch( + "vllm.v1.worker.gpu.cudagraph_utils.is_global_first_rank", + return_value=False, + ), + ): + manager.capture(create_forward_fn) + + assert [warmup for warmup, _ in create_calls] == [True, False] + assert [mode for _, mode, _ in forward_calls] == [ + CUDAGraphMode.NONE, + CUDAGraphMode.PIECEWISE, + ] + assert create_calls[0][1] is not create_calls[1][1] + + @pytest.fixture(autouse=True) def _reset_breakable_tls(): """Defensively clear thread-local capture state between tests so a diff --git a/tests/v1/e2e/spec_decode/test_spec_decode.py b/tests/v1/e2e/spec_decode/test_spec_decode.py index 96d6684594cc..01bec69640b8 100644 --- a/tests/v1/e2e/spec_decode/test_spec_decode.py +++ b/tests/v1/e2e/spec_decode/test_spec_decode.py @@ -286,6 +286,73 @@ def test_suffix_decoding_acceptance( cleanup_dist_env_and_memory() +@pytest.mark.slow_test +@large_gpu_mark(min_gb=80) +def test_gemma4_dspark_correctness_and_acceptance_rate( + monkeypatch: pytest.MonkeyPatch, +): + """ + E2E test for Gemma4 DSpark speculative decoding: acceptance rate/length + regression coverage plus GSM8K correctness, at temperature=1.0 to exercise + the probabilistic draft-sampling/rejection-sampling path (not just greedy). + + Uses google/gemma-4-12B-it as target with the dspark_gemma4_12b_block7 draft + model. Exercises the full Gemma4 DSpark path (draft build over the reused + base classes DFlashQwen3Model / Qwen3DSparkForCausalLM / Gemma4MTP* / + DSparkMarkovHead, the fused context-KV precompute, the Markov head, and + rejection sampling), so it fails if any base class drifts out of sync. + + gemma-4-12B is instruct-tuned, so GSM8K is run through the chat template + (use_chat_completions=True; raw few-shot completion collapses to a few + percent). Reference: measured over 5 runs of 200 GSM8K questions at + temperature=1.0 (prefix caching disabled): + accuracy: min=0.900 max=0.955 mean=0.937 + acceptance_rate: min=0.578 max=0.595 mean=0.588 + acceptance_len: min=5.044 max=5.167 mean=5.116 + Thresholds set conservatively to 10% to avoid flaking due to unlucky sampling + """ + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") + + spec_llm = LLM( + model="google/gemma-4-12B-it", + trust_remote_code=True, + speculative_config={ + "method": "dspark", + "model": "deepseek-ai/dspark_gemma4_12b_block7", + "num_speculative_tokens": 7, + "draft_sample_method": "probabilistic", + }, + max_model_len=8192, + max_num_seqs=32, + gpu_memory_utilization=0.8, + enforce_eager=True, + enable_prefix_caching=False, + disable_log_stats=False, + ) + try: + results = evaluate_gsm8k_offline( + spec_llm, num_questions=200, temperature=1.0, use_chat_completions=True + ) + gsm8k_accuracy = results["accuracy"] + + metrics = spec_llm.get_metrics() + acceptance_rate = compute_acceptance_rate(metrics) + acceptance_len = compute_acceptance_len(metrics) + print( + f"Gemma4 DSpark acceptance_rate={acceptance_rate:.2f}, " + f"acceptance_len={acceptance_len:.2f}, " + f"gsm8k_accuracy={gsm8k_accuracy:.3f}" + ) + + assert acceptance_rate >= 0.588 * 0.9 + assert acceptance_len >= 5.116 * 0.9 + assert gsm8k_accuracy >= 0.937 * 0.9 + finally: + del spec_llm + torch.accelerator.empty_cache() + cleanup_dist_env_and_memory() + + @pytest.mark.parametrize( ["model_path", "expected_accuracy_threshold"], [ diff --git a/tests/v1/engine/test_iteration_logging.py b/tests/v1/engine/test_iteration_logging.py new file mode 100644 index 000000000000..5a08308163ec --- /dev/null +++ b/tests/v1/engine/test_iteration_logging.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import time +from types import SimpleNamespace + +from vllm.v1.engine import EngineCoreOutputs +from vllm.v1.engine.core import EngineCore +from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats + + +class FakeEngineCore: + def _make_iteration_details_stats( + self, iteration_details: SchedulerIterationDetails + ) -> SchedulerStats: + return SchedulerStats(iteration_details=iteration_details) + + +def make_iteration_details() -> SchedulerIterationDetails: + return SchedulerIterationDetails( + iteration_index=1, + num_ctx_requests=2, + num_ctx_tokens=3, + num_generation_requests=4, + num_generation_tokens=5, + elapsed_ms=6.7, + ) + + +def make_fake_engine(log_stats: bool = True) -> SimpleNamespace: + return SimpleNamespace( + log_stats=log_stats, + vllm_config=SimpleNamespace( + observability_config=SimpleNamespace( + enable_logging_iteration_details=True, + ) + ), + ) + + +def test_capture_iteration_details_disabled_without_log_stats(): + engine = make_fake_engine(log_stats=False) + + with EngineCore.capture_iteration_details(engine, None) as iteration_details: + assert iteration_details is None + + assert not hasattr(engine, "_iteration_index") + + +def test_capture_iteration_details_fills_elapsed_time(): + engine = make_fake_engine() + + with EngineCore.capture_iteration_details(engine, None) as iteration_details: + assert iteration_details is not None + assert iteration_details.elapsed_ms == 0.0 + assert iteration_details.is_dummy + time.sleep(0.001) + + assert iteration_details is not None + assert iteration_details.elapsed_ms > 0.0 + assert engine._iteration_index == 1 + + +def test_attach_iteration_details_uses_existing_output(): + iteration_details = make_iteration_details() + outputs = { + 2: EngineCoreOutputs(scheduler_stats=SchedulerStats()), + 1: EngineCoreOutputs(scheduler_stats=SchedulerStats()), + } + + EngineCore._attach_iteration_details(FakeEngineCore(), outputs, iteration_details) + + assert 0 not in outputs + assert outputs[2].scheduler_stats is not None + assert outputs[2].scheduler_stats.iteration_details == iteration_details + assert outputs[1].scheduler_stats is not None + assert outputs[1].scheduler_stats.iteration_details is None + + +def test_attach_iteration_details_falls_back_to_client_zero_without_outputs(): + iteration_details = make_iteration_details() + outputs: dict[int, EngineCoreOutputs] = {} + + EngineCore._attach_iteration_details(FakeEngineCore(), outputs, iteration_details) + + assert set(outputs) == {0} + assert outputs[0].scheduler_stats is not None + assert outputs[0].scheduler_stats.iteration_details == iteration_details diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index a357128d3ce4..77235582fe13 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -140,12 +140,13 @@ run_tests_for_model() { # Start prefill instances for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do # Calculate GPU ID - we'll distribute across available GPUs - GPU_ID=$((i % $(get_num_gpus))) - NEXT_GPU=${GPU_ID} + GPU_START=$((i % $(get_num_gpus))) + GPU_ID=$GPU_START + NEXT_GPU=$GPU_START # Reserve TP*PP GPUs for the prefiller (TP shards across PP stages). PREFILLER_WORLD_SIZE=$((PREFILLER_TP_SIZE * PREFILLER_PP_SIZE)) for (( j=1; j < PREFILLER_WORLD_SIZE; j++ )); do - NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus))) + NEXT_GPU=$(((GPU_START + j) % $(get_num_gpus))) GPU_ID="${GPU_ID},${NEXT_GPU}" done @@ -195,10 +196,12 @@ run_tests_for_model() { # Start decode instances for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do # Calculate GPU ID - we'll distribute across available GPUs, starting from after prefill GPUs - GPU_ID=$(((i + NEXT_GPU + 1) % $(get_num_gpus))) + DECODE_START=$(((i + NEXT_GPU + 1) % $(get_num_gpus))) + GPU_ID=$DECODE_START + NEXT_GPU=$DECODE_START # If DECODER_TP_SIZE is more than 1 for (( j=1; j < DECODER_TP_SIZE; j++ )); do - NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus))) + NEXT_GPU=$(((DECODE_START + j) % $(get_num_gpus))) GPU_ID="${GPU_ID},${NEXT_GPU}" done # Calculate port number (base port + instance number) diff --git a/tests/v1/kv_connector/nixl_integration/test_accuracy.py b/tests/v1/kv_connector/nixl_integration/test_accuracy.py index bb68b7a57241..711f27d20945 100644 --- a/tests/v1/kv_connector/nixl_integration/test_accuracy.py +++ b/tests/v1/kv_connector/nixl_integration/test_accuracy.py @@ -26,6 +26,7 @@ "Qwen/Qwen3.5-0.8B": 0.33, "google/gemma-4-E2B-it": 0.485, "ai21labs/AI21-Jamba2-3B": 0.74, + "deepseek-ai/DeepSeek-V4-Flash": 0.95, } SIMPLE_PROMPT = ( diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index beb639724e80..a23fb95b5683 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -7,7 +7,10 @@ from tests.v1.kv_connector.unit.utils import create_vllm_config from vllm.config import KVEventsConfig, KVTransferConfig -from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored +from vllm.distributed.kv_events import MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( + build_offloading_config, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.events import ( OffloadingEventGroupSpec, OffloadingEventsTracker, @@ -23,6 +26,7 @@ KVCacheSpecKind, ) from vllm.v1.kv_offload.base import ( + Locality, OffloadingEvent, OffloadingKVEventsConfig, OffloadKey, @@ -70,15 +74,15 @@ def _group_config( *, group_idx: int = 0, block_size: int = 4, - block_size_factor: int = 1, - sliding_window_size_in_blocks: int | None = None, + blocks_per_chunk: int = 1, + sliding_window_size_in_chunks: int | None = None, ) -> GroupOffloadConfig: return GroupOffloadConfig( group_idx=group_idx, - gpu_block_size=block_size, - offloaded_block_size=block_size * block_size_factor, - hash_block_size_factor=block_size_factor, - sliding_window_size_in_blocks=sliding_window_size_in_blocks, + tokens_per_block=block_size, + tokens_per_chunk=block_size * blocks_per_chunk, + hashes_per_chunk=blocks_per_chunk, + sliding_window_size_in_chunks=sliding_window_size_in_chunks, kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC, ) @@ -90,7 +94,7 @@ def _record_chunks( num_chunks: int, ) -> list[OffloadKey]: keys: list[OffloadKey] = [] - hbf = group_config.hash_block_size_factor + hbf = group_config.hashes_per_chunk for chunk_idx in range(num_chunks): tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1] assert tail_hash is not None @@ -100,12 +104,81 @@ def _record_chunks( return keys -def _stored_event(keys: list[OffloadKey]) -> OffloadingEvent: - return OffloadingEvent(keys=keys, medium=_CPU_MEDIUM, removed=False) +def _stored_event( + keys: list[OffloadKey], + locality: Locality | None = None, + medium: str = _CPU_MEDIUM, +) -> OffloadingEvent: + return OffloadingEvent( + keys=keys, + medium=medium, + removed=False, + locality=locality, + ) + + +def _removed_event( + keys: list[OffloadKey], + locality: Locality | None = None, + medium: str = _CPU_MEDIUM, +) -> OffloadingEvent: + return OffloadingEvent( + keys=keys, + medium=medium, + removed=True, + locality=locality, + ) + + +def test_take_events_forwards_locality_to_rich_store(): + tracker = _tracker() + req = _request(block_hashes=[_hash(0)], token_count=4) + key = _record_chunks(tracker, req, _group_config(), num_chunks=1)[0] + + events = list( + tracker.take_events( + [_stored_event([key], locality=Locality.LOCAL, medium=MEDIUM_FS)] + ) + ) + + assert len(events) == 1 + assert isinstance(events[0], BlockStored) + assert events[0].token_ids == [1, 2, 3, 4] + assert events[0].block_size == 4 + assert events[0].locality == "LOCAL" + + +def test_take_events_forwards_locality_to_placeholder_store(): + tracker = _tracker(self_describing_kv_events=False) + req = _request(block_hashes=[_hash(0)], token_count=4) + key = _record_chunks(tracker, req, _group_config(), num_chunks=1)[0] + + events = list( + tracker.take_events( + [_stored_event([key], locality=Locality.REMOTE, medium=MEDIUM_FS)] + ) + ) + + assert len(events) == 1 + assert isinstance(events[0], BlockStored) + assert events[0].block_size == 0 + assert events[0].locality == "REMOTE" + +def test_take_events_forwards_locality_to_remove(): + tracker = _tracker() + req = _request(block_hashes=[_hash(0)], token_count=4) + key = _record_chunks(tracker, req, _group_config(), num_chunks=1)[0] -def _removed_event(keys: list[OffloadKey]) -> OffloadingEvent: - return OffloadingEvent(keys=keys, medium=_CPU_MEDIUM, removed=True) + events = list( + tracker.take_events( + [_removed_event([key], locality=Locality.LOCAL, medium=MEDIUM_FS)] + ) + ) + + assert len(events) == 1 + assert isinstance(events[0], BlockRemoved) + assert events[0].locality == "LOCAL" def test_take_events_publishes_routable_block_stored(): @@ -149,14 +222,14 @@ def test_take_events_publishes_routable_block_stored(): def test_take_events_factor_gt_1_chunk_store_and_remove(): block_size = 4 - block_size_factor = 3 + blocks_per_chunk = 3 tracker = _tracker() group_config = _group_config( - block_size=block_size, block_size_factor=block_size_factor + block_size=block_size, blocks_per_chunk=blocks_per_chunk ) req = _request( block_hashes=[_hash(i) for i in range(6)], - token_count=block_size * block_size_factor * 2, + token_count=block_size * blocks_per_chunk * 2, ) keys = _record_chunks(tracker, req, group_config, num_chunks=2) @@ -169,17 +242,17 @@ def test_take_events_factor_gt_1_chunk_store_and_remove(): expected_chunk_hashes = [ _wire_hash(_hash(i)) for i in range( - chunk_idx * block_size_factor, - (chunk_idx + 1) * block_size_factor, + chunk_idx * blocks_per_chunk, + (chunk_idx + 1) * blocks_per_chunk, ) ] assert event.block_hashes == expected_chunk_hashes assert event.block_size == block_size - assert len(event.token_ids) == block_size * block_size_factor + assert len(event.token_ids) == block_size * blocks_per_chunk if chunk_idx == 0: assert event.parent_block_hash is None else: - assert event.parent_block_hash == _wire_hash(_hash(block_size_factor - 1)) + assert event.parent_block_hash == _wire_hash(_hash(blocks_per_chunk - 1)) expected_hashes.extend(expected_chunk_hashes) assert len(tracker._pending_event_metadata) == 2 @@ -194,12 +267,12 @@ def test_take_events_factor_gt_1_chunk_store_and_remove(): def test_take_events_factor_gt_1_store_is_order_independent(): - block_size_factor = 3 + blocks_per_chunk = 3 tracker = _tracker() - group_config = _group_config(block_size_factor=block_size_factor) + group_config = _group_config(blocks_per_chunk=blocks_per_chunk) req = _request( block_hashes=[_hash(i) for i in range(6)], - token_count=4 * block_size_factor * 2, + token_count=4 * blocks_per_chunk * 2, ) keys = _record_chunks(tracker, req, group_config, num_chunks=2) unknown_key = make_offload_key(_hash(12345), 0) @@ -244,7 +317,7 @@ def test_take_events_opt_out_keeps_placeholders(): def test_record_store_skips_sliding_window_group(): tracker = _tracker() - group_config = _group_config(sliding_window_size_in_blocks=2) + group_config = _group_config(sliding_window_size_in_chunks=2) req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12) keys = _record_chunks(tracker, req, group_config, num_chunks=3) @@ -258,8 +331,8 @@ def test_record_store_skips_sliding_window_group(): def test_take_events_groups_removed_hashes_by_kv_group(): tracker = _tracker() - group0_config = _group_config(group_idx=0, block_size_factor=2) - group1_config = _group_config(group_idx=1, block_size_factor=2) + group0_config = _group_config(group_idx=0, blocks_per_chunk=2) + group1_config = _group_config(group_idx=1, blocks_per_chunk=2) req0 = _request(block_hashes=[_hash(0), _hash(1)], token_count=8) req1 = _request(block_hashes=[_hash(10), _hash(11)], token_count=8) key0 = _record_chunks(tracker, req0, group0_config, num_chunks=1)[0] @@ -293,7 +366,7 @@ def test_take_events_supports_restore_after_eviction(): assert not tracker._pending_event_metadata req.all_token_ids = [5, 6, 7, 8] - tracker.record_store(req, group_config, offload_block_idx=0, offload_key=key) + tracker.record_store(req, group_config, chunk_idx=0, offload_key=key) second_store = list(tracker.take_events([_stored_event([key])])) assert len(second_store) == 1 @@ -351,4 +424,4 @@ def test_tiering_rejects_self_describing_kv_events(): ) with pytest.raises(ValueError, match="TieringOffloadingSpec"): - TieringOffloadingSpec(vllm_config, kv_cache_config) + TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config)) 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 6c55b91d8da4..bc5190db00f5 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -64,6 +64,47 @@ def test_scheduler_reports_allocation_failure(request_runner): assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1 +@pytest.mark.parametrize("async_scheduling", [True, False]) +@pytest.mark.parametrize("prompt_offset", [-1, -2]) +def test_last_block_offloaded_at_request_finish( + request_runner, async_scheduling: bool, prompt_offset: int +): + """EOS fills the last block at request finish — verify req_status is kept alive. + + prompt = block_size + prompt_offset tokens → not a full block at schedule time, + so _build_store_jobs creates no store job. After EOS, request_finished + keeps req_status alive so _build_store_jobs can process it on the next step. + + prompt_offset=-1: EOS fills the block → store job created on next step. + prompt_offset=-2: block remains partial → no store job, cleanup in + _build_store_jobs deletes req_status. + """ + block_size = 4 + runner = request_runner( + block_size=block_size, + num_gpu_blocks=10, + async_scheduling=async_scheduling, + ) + # prompt = block_size + prompt_offset tokens + runner.new_request(token_ids=[0] * (block_size + prompt_offset)) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(list(keys)) + ) + + # Run with one step (EOS) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + ) + + cs = runner.connector_scheduler + # Verify req_status is kept alive for _build_store_jobs to process + # regardless of whether there are storable blocks + assert "0" in cs._req_status, ( + "req_status was deleted but should be kept alive " + "for _build_store_jobs to process finished_req_ids." + ) + + def test_scheduler_reports_lookup_sync_delay(request_runner): runner = request_runner( block_size=4, @@ -105,20 +146,20 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner): @pytest.mark.parametrize("async_scheduling", [True, False]) def test_offloading_connector(request_runner, async_scheduling: bool): block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 runner = request_runner( block_size=block_size, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) # 3 blocks, store just the middle block (skip first and last) # blocks = [0, 1, 2], [3, 4, 5], [6, 7, 8] - runner.new_request(token_ids=[0] * offloaded_block_size * 3) + runner.new_request(token_ids=[0] * tokens_per_chunk * 3) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(list(keys)[1:2]) ) @@ -126,7 +167,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): # add block missing 1 token -> no offload runner.run( - decoded_tokens=[0] * (offloaded_block_size - 1), + decoded_tokens=[0] * (tokens_per_chunk - 1), expected_stored=(3, 4, 5), ) runner.manager.touch.assert_not_called() @@ -141,7 +182,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) - runner.run(decoded_tokens=[0] * (offloaded_block_size + 1)) + runner.run(decoded_tokens=[0] * (tokens_per_chunk + 1)) # 1 more block (+ token for kicking off offloading) # now check touch was called with all 6 blocks @@ -149,7 +190,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): generate_store_output(keys) ) runner.run( - decoded_tokens=[0] * (offloaded_block_size + 1), + decoded_tokens=[0] * (tokens_per_chunk + 1), expected_stored=(15, 16, 17), ) runner.manager.touch.assert_called() @@ -160,7 +201,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.run(decoded_tokens=[EOS_TOKEN_ID]) # create a new request differing only on the last token - runner.new_request(token_ids=[0] * (offloaded_block_size * 6 - 1) + [1]) + runner.new_request(token_ids=[0] * (tokens_per_chunk * 6 - 1) + [1]) runner.run(decoded_tokens=[0]) runner.manager.touch.assert_called() block_hashes2 = list(runner.manager.touch.call_args.args[0]) @@ -173,12 +214,12 @@ def test_offloading_connector(request_runner, async_scheduling: bool): # terminate request runner.run( decoded_tokens=[EOS_TOKEN_ID], - expected_stored=tuple(range(6 * block_size_factor)), + expected_stored=tuple(range(6 * blocks_per_chunk)), ) - # full_block_tokens - num_computed_tokens < offloaded_block_size + # full_block_tokens - num_computed_tokens < tokens_per_chunk runner.new_request( - token_ids=[0] * block_size + [1] * (offloaded_block_size - block_size) + token_ids=[0] * block_size + [1] * (tokens_per_chunk - block_size) ) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) @@ -187,7 +228,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.lookup.assert_not_called() # single block lookup with no hits - runner.new_request(token_ids=[1] * offloaded_block_size) + runner.new_request(token_ids=[1] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -196,7 +237,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): # single block lookup with a hit runner.scheduler.reset_prefix_cache() - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -204,9 +245,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2)) # single block lookup with a hit in a middle block - runner.new_request( - token_ids=[0] * offloaded_block_size * 2 + [1] * offloaded_block_size - ) + runner.new_request(token_ids=[0] * tokens_per_chunk * 2 + [1] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -217,15 +256,15 @@ 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): block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 runner = request_runner( block_size=block_size, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue @@ -233,7 +272,7 @@ def test_request_preemption(request_runner, async_scheduling: bool): # 2 blocks, store all, without flushing # blocks = [0, 1, 2], [3, 4, 5] - runner.new_request(token_ids=[0] * offloaded_block_size * 2) + runner.new_request(token_ids=[0] * tokens_per_chunk * 2) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -247,7 +286,7 @@ def test_request_preemption(request_runner, async_scheduling: bool): generate_store_output(keys) ) runner.run( - decoded_tokens=[0] * (2 * offloaded_block_size - block_size), + decoded_tokens=[0] * (2 * tokens_per_chunk - block_size), complete_transfers=False, ) @@ -297,14 +336,14 @@ def test_on_request_finished_is_not_deferred_until_store_completion( still arrive afterward for already-submitted transfer jobs. """ block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk runner = request_runner( block_size=block_size, num_gpu_blocks=100, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) # Record the order of per-request connector calls on the (mocked) manager. @@ -324,10 +363,10 @@ def test_on_request_finished_is_not_deferred_until_store_completion( # Decode a couple of blocks, keeping every transfer in flight # (complete_transfers=False) so no store completes while the request runs. - runner.new_request(token_ids=[0] * offloaded_block_size * 2) + runner.new_request(token_ids=[0] * tokens_per_chunk * 2) runner.run(decoded_tokens=[0], complete_transfers=False) runner.run( - decoded_tokens=[0] * (2 * offloaded_block_size), + decoded_tokens=[0] * (2 * tokens_per_chunk), complete_transfers=False, ) @@ -347,7 +386,7 @@ def test_on_request_finished_is_not_deferred_until_store_completion( runner.run( decoded_tokens=[], complete_transfers=True, - expected_stored=tuple(range(4 * block_size_factor)), + expected_stored=tuple(range(4 * blocks_per_chunk)), ) # on_request_finished is issued exactly once. @@ -364,19 +403,19 @@ def test_on_request_finished_is_not_deferred_until_store_completion( @pytest.mark.parametrize("async_scheduling", [True, False]) def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool): block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 runner = request_runner( block_size=block_size, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) # store 1 blocks - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -387,7 +426,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: # start a request to load the first block, but don't complete runner.scheduler.reset_prefix_cache() - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run( decoded_tokens=[], @@ -399,7 +438,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: assert transfer_jobs # start a new request to load the same first block - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run( decoded_tokens=[], @@ -428,19 +467,19 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: @pytest.mark.parametrize("async_scheduling", [True, False]) def test_abort_loading_requests(request_runner, async_scheduling: bool): block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 runner = request_runner( block_size=block_size, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) # store 1 blocks - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -451,7 +490,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool): # start a request to load the first block, but don't complete runner.scheduler.reset_prefix_cache() - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.run( decoded_tokens=[], @@ -483,7 +522,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool): def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bool): block_size = 4 num_gpu_blocks = 100 - # sliding_window=8 -> 2 offloaded blocks (block_size_factor=1) + # sliding_window=8 -> 2 offloaded chunks (blocks_per_chunk=1) sliding_window = 8 kv_cache_groups = [ @@ -518,8 +557,8 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # Verify group configs: group 0 = full attention, group 1 = sliding window kv_group_configs = runner.connector_scheduler.config.kv_group_configs assert len(kv_group_configs) == 2 - assert kv_group_configs[0].sliding_window_size_in_blocks is None - assert kv_group_configs[1].sliding_window_size_in_blocks == 2 + assert kv_group_configs[0].sliding_window_size_in_chunks is None + assert kv_group_configs[1].sliding_window_size_in_chunks == 2 # Blocks [0, 1, 2] miss runner.new_request(token_ids=[0] * block_size * 3) @@ -562,11 +601,19 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # Group 1 (sliding window, window=2): only the last 2 blocks # are within the window → loads blocks 1,2 expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)), + # The deferred store from the previous request's last block + # completes during this step, and its blocks are flushed because + # they were reallocated to the new request. + # Only block 1 (sliding window group) is stored — block 0's + # deferred store is flushed because it was reallocated. + expected_stored=((0, 1),), + expected_flushed=((0, 1),), ) - # one touch in get_num_new_matched_tokens x 2 groups + # 4 touch calls: 2 from get_num_new_matched_tokens (2 groups) + # + 2 from _get_reqs_to_store (2 groups) touch_calls = runner.manager.touch.call_args_list - assert len(touch_calls) == 2 + assert len(touch_calls) == 4 # full attention group touched all 3 blocks assert len(touch_calls[0].args[0]) == 3 # sliding window group touched just the last 2 blocks @@ -587,16 +634,16 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo @pytest.mark.parametrize("async_scheduling", [True, False]) def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool): - hash_block_size = 4 + tokens_per_hash = 4 num_gpu_blocks = 100 - # Group 0: block_size=12 (offloaded_block_size=12) - # Group 1: block_size=16 (offloaded_block_size=16) + # Group 0: block_size=12 (tokens_per_chunk=12) + # Group 1: block_size=16 (tokens_per_chunk=16) kv_cache_groups = [ KVCacheGroupSpec( ["layer0"], FullAttentionSpec( - block_size=hash_block_size * 3, + block_size=tokens_per_hash * 3, num_kv_heads=1, head_size=1, dtype=torch.float32, @@ -605,7 +652,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool KVCacheGroupSpec( ["layer1"], FullAttentionSpec( - block_size=hash_block_size * 4, + block_size=tokens_per_hash * 4, num_kv_heads=1, head_size=1, dtype=torch.float32, @@ -614,7 +661,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool ] runner = request_runner( - block_size=hash_block_size, + block_size=tokens_per_hash, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, @@ -623,10 +670,10 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # Verify group configs kv_group_configs = runner.connector_scheduler.config.kv_group_configs assert len(kv_group_configs) == 2 - assert kv_group_configs[0].gpu_block_size == 12 - assert kv_group_configs[0].offloaded_block_size == 12 - assert kv_group_configs[1].gpu_block_size == 16 - assert kv_group_configs[1].offloaded_block_size == 16 + assert kv_group_configs[0].tokens_per_block == 12 + assert kv_group_configs[0].tokens_per_chunk == 12 + assert kv_group_configs[1].tokens_per_block == 16 + assert kv_group_configs[1].tokens_per_chunk == 16 # Prompt: 25 tokens, unaligned to both block sizes. # Group 0 blocks: [0, 1], ending_token_offset = 24 @@ -932,20 +979,20 @@ def test_hit_pending_does_not_break_streak(self): @pytest.mark.parametrize("async_scheduling", [True, False]) def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool): """With REQUEST_LEVEL policy, all blocks are stored — including prefix hits.""" - gpu_block_size = 4 - block_size_factor = 3 - offloaded_block_size = gpu_block_size * block_size_factor + tokens_per_block = 4 + blocks_per_chunk = 3 + tokens_per_chunk = tokens_per_block * blocks_per_chunk num_gpu_blocks = 100 runner = request_runner( - block_size_factor=block_size_factor, - block_size=gpu_block_size, + blocks_per_chunk=blocks_per_chunk, + block_size=tokens_per_block, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, ) - # Store 1 offloaded block (3 GPU blocks) via a normal request. - runner.new_request(token_ids=[0] * offloaded_block_size) + # Store 1 offloaded chunk (3 GPU blocks) via a normal request. + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -962,14 +1009,14 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling policy=OffloadPolicy.REQUEST_LEVEL ) - # New request with 2 offloaded blocks; first matches what's in CPU. - runner.new_request(token_ids=[0] * offloaded_block_size * 2) + # New request with 2 offloaded chunks; first matches what's in CPU. + runner.new_request(token_ids=[0] * tokens_per_chunk * 2) runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) - # Load the first offloaded block from CPU. + # Load the first offloaded chunk from CPU. runner.run(decoded_tokens=[0], expected_loaded=(0, 1, 2)) # Store must include ALL 6 GPU blocks (both the loaded prefix and @@ -989,7 +1036,7 @@ def test_loads_do_not_populate_fence_index(request_runner): """Loads don't populate _block_id_to_pending_jobs (protected by delay_free_blocks while in flight).""" runner = request_runner( - block_size_factor=3, + blocks_per_chunk=3, block_size=4, num_gpu_blocks=100, async_scheduling=False, @@ -1008,7 +1055,7 @@ def test_fence_at_update_state_after_alloc(request_runner): req1 just freed. """ runner = request_runner( - block_size_factor=1, + blocks_per_chunk=1, block_size=4, num_gpu_blocks=2, async_scheduling=False, @@ -1059,7 +1106,7 @@ def test_fence_at_build_store_jobs(request_runner): reusing a finished request's pending-store block is flushed by _build_store_jobs's fence.""" runner = request_runner( - block_size_factor=1, + blocks_per_chunk=1, block_size=4, num_gpu_blocks=2, async_scheduling=False, @@ -1108,16 +1155,16 @@ def capture_fence(): def test_complete_store_called_per_job(request_runner, async_scheduling: bool): """complete_store fires per-job, not deferred to request finish. Each call carries only that store's keys.""" - gpu_block_size = 4 - block_size_factor = 3 - offloaded_block_size = gpu_block_size * block_size_factor + tokens_per_block = 4 + blocks_per_chunk = 3 + tokens_per_chunk = tokens_per_block * blocks_per_chunk runner = request_runner( - block_size_factor=block_size_factor, - block_size=gpu_block_size, + blocks_per_chunk=blocks_per_chunk, + block_size=tokens_per_block, num_gpu_blocks=100, async_scheduling=async_scheduling, ) - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -1131,7 +1178,7 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool): # Second store: fires when block 1 is fully populated, with different keys. runner.run( - decoded_tokens=[0] * (offloaded_block_size + 1), + decoded_tokens=[0] * (tokens_per_chunk + 1), expected_stored=(3, 4, 5), ) assert runner.manager.complete_store.call_count == 1 @@ -1148,25 +1195,25 @@ def test_complete_store_called_per_job(request_runner, async_scheduling: bool): def test_max_offload_tokens_validation(request_runner, async_scheduling: bool): """Validates max_offload_tokens: type coercion, boundary values, and capping. - Setup: 3 offloaded blocks × 3 GPU blocks each = 9 GPU block offsets (0–8). + Setup: 3 offloaded chunks × 3 GPU blocks each = 9 GPU block offsets (0–8). """ - gpu_block_size = 4 - block_size_factor = 3 - offloaded_block_size = gpu_block_size * block_size_factor # 12 + tokens_per_block = 4 + blocks_per_chunk = 3 + tokens_per_chunk = tokens_per_block * blocks_per_chunk # 12 num_gpu_blocks = 100 all_offsets = (0, 1, 2, 3, 4, 5, 6, 7, 8) def make_runner(): return request_runner( - block_size=gpu_block_size, + block_size=tokens_per_block, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) def setup(r, max_offload_tokens): r.new_request( - token_ids=[0] * offloaded_block_size * 3, + token_ids=[0] * tokens_per_chunk * 3, kv_transfer_params={"max_offload_tokens": max_offload_tokens}, ) r.manager.prepare_store.side_effect = lambda keys, req_context: ( @@ -1228,9 +1275,9 @@ def setup(r, max_offload_tokens): setup(r, 0) r.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=()) - # positive int cap -> limits offload to first 2 offloaded blocks (offsets 0–5) + # positive int cap -> limits offload to first 2 chunks (offsets 0–5) r = make_runner() - setup(r, 24) # 24 tokens = 2 offloaded blocks × 12 tokens each + setup(r, 24) # 24 tokens = 2 offloaded chunks × 12 tokens each r.run( decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0, 1, 2, 3, 4, 5), @@ -1242,8 +1289,8 @@ def setup(r, max_offload_tokens): def test_offload_prompt_only(request_runner, async_scheduling: bool): """offload_prompt_only=True offloads prompt blocks but never decode blocks. - Setup: a 2-offloaded-block prompt followed by enough decode tokens to fill - 4 more offloaded blocks. The flag clamps the offloadable token count to the + Setup: a 2-chunk prompt followed by enough decode tokens to fill + 4 more offloaded chunks. The flag clamps the offloadable token count to the prompt length, so only the prompt's blocks (GPU offsets 0-5) are ever eligible for store; the decode blocks (offsets >= 6) are skipped. @@ -1253,18 +1300,18 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool): subtleties. The decode steps are still enough for the prompt store to complete and show up in expected_stored. """ - gpu_block_size = 4 - block_size_factor = 3 - offloaded_block_size = gpu_block_size * block_size_factor # 12 + tokens_per_block = 4 + blocks_per_chunk = 3 + tokens_per_chunk = tokens_per_block * blocks_per_chunk # 12 num_prompt_blocks = 2 num_decode_blocks = 4 prompt_offsets = (0, 1, 2, 3, 4, 5) runner = request_runner( - block_size=gpu_block_size, + block_size=tokens_per_block, num_gpu_blocks=100, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, extra_config_overrides={"offload_prompt_only": True}, ) @@ -1272,9 +1319,9 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool): generate_store_output(keys) ) - runner.new_request(token_ids=[0] * offloaded_block_size * num_prompt_blocks) + runner.new_request(token_ids=[0] * tokens_per_chunk * num_prompt_blocks) runner.run( - decoded_tokens=[0] * (offloaded_block_size * num_decode_blocks), + decoded_tokens=[0] * (tokens_per_chunk * num_decode_blocks), expected_stored=prompt_offsets, ) @@ -1291,21 +1338,21 @@ def test_offload_prompt_only(request_runner, async_scheduling: bool): @pytest.mark.parametrize("async_scheduling", [True, False]) def test_reset_cache(request_runner, async_scheduling: bool): """reset_cache flushes in-flight loads, calls manager.reset_cache(), resets - next_stored_block_idx for active requests and clears job tracking.""" + next_stored_chunk_idx for active requests and clears job tracking.""" block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 runner = request_runner( block_size=block_size, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) - # Store 1 offloaded block (3 GPU blocks) to CPU. - runner.new_request(token_ids=[0] * offloaded_block_size) + # Store 1 offloaded chunk (3 GPU blocks) to CPU. + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -1317,7 +1364,7 @@ def test_reset_cache(request_runner, async_scheduling: bool): # Reset GPU prefix cache then start a request that loads from CPU. # Leave the load in-flight so that reset_cache must flush it. runner.scheduler.reset_prefix_cache() - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1 runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) @@ -1335,11 +1382,11 @@ def test_reset_cache(request_runner, async_scheduling: bool): # Record job counter to verify the reset counter is set correctly. job_counter_before_reset = runner.connector_scheduler._job_counter - # After update_state_after_alloc, next_stored_block_idx is advanced to + # After update_state_after_alloc, next_stored_chunk_idx is advanced to # skip the loaded prefix; reset_cache must bring it back to 0. for req_status in runner.connector_scheduler._req_status.values(): for group_state in req_status.group_states: - assert group_state.next_stored_block_idx > 0 + assert group_state.next_stored_chunk_idx > 0 # Reset the cache runner.connector_scheduler.reset_cache() @@ -1354,18 +1401,18 @@ def test_reset_cache(request_runner, async_scheduling: bool): # All internal job tracking must be cleared. assert not runner.connector_scheduler._jobs assert not runner.connector_scheduler._block_id_to_pending_jobs - if runner.connector_scheduler._blocks_being_loaded is not None: - assert not runner.connector_scheduler._blocks_being_loaded + if runner.connector_scheduler._chunks_being_loaded is not None: + assert not runner.connector_scheduler._chunks_being_loaded # Job reset counter must equal the job counter so that completions for # pre-reset jobs arriving from workers are silently discarded. assert runner.connector_scheduler._stale_job_threshold == job_counter_before_reset - # next_stored_block_idx must be reset to 0 for every active request so + # next_stored_chunk_idx must be reset to 0 for every active request so # that post-reset stores restart from block 0. for req_status in runner.connector_scheduler._req_status.values(): for group_state in req_status.group_states: - assert group_state.next_stored_block_idx == 0 + assert group_state.next_stored_chunk_idx == 0 @pytest.mark.parametrize("async_scheduling", [True, False]) @@ -1376,14 +1423,14 @@ def test_reset_cache_finalizes_finished_request_with_pending_store( without calling on_request_finished twice. """ block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk runner = request_runner( block_size=block_size, num_gpu_blocks=100, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) finalized: list[str] = [] @@ -1396,10 +1443,10 @@ def test_reset_cache_finalizes_finished_request_with_pending_store( # Decode a couple of blocks and keep every transfer in flight, so the # request has pending store jobs. - runner.new_request(token_ids=[0] * offloaded_block_size * 2) + runner.new_request(token_ids=[0] * tokens_per_chunk * 2) runner.run(decoded_tokens=[0], complete_transfers=False) runner.run( - decoded_tokens=[0] * (2 * offloaded_block_size), + decoded_tokens=[0] * (2 * tokens_per_chunk), complete_transfers=False, ) @@ -1430,7 +1477,7 @@ def test_pending_transfer_defers_prefix_lookup(): With async scheduling, a preempted request's store can be flushed by the worker before the scheduler consumes its completion. If the request is re-admitted in that window, the connector should defer it instead of - looking up offloaded blocks and later asserting when a load is queued while + looking up offloaded chunks and later asserting when a load is queued while the store job is still tracked. """ scheduler = object.__new__(OffloadingConnectorScheduler) @@ -1464,27 +1511,27 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner re-admission path must defer while the scheduler still tracks the store. """ block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk runner = request_runner( block_size=block_size, num_gpu_blocks=100, async_scheduling=True, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue num_free_blocks_empty = free_block_queue.num_free_blocks req_id = "0" - runner.new_request(token_ids=[0] * offloaded_block_size * 2) + runner.new_request(token_ids=[0] * tokens_per_chunk * 2) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) runner.run(decoded_tokens=[0], complete_transfers=False) runner.run( - decoded_tokens=[0] * (2 * offloaded_block_size - block_size), + decoded_tokens=[0] * (2 * tokens_per_chunk - block_size), complete_transfers=False, ) @@ -1530,8 +1577,8 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): - Group 0: full attention (MLA-like), block_size=16 - Group 1: SWA, block_size=4, sliding_window=8 - alignment_block_count = 16 / 4 = 4 SWA blocks per alignment segment. - sliding_window_size_in_blocks = ceil(8 / 4) = 2. + alignment_chunk_count = 16 / 4 = 4 SWA blocks per alignment segment. + sliding_window_size_in_chunks = ceil(8 / 4) = 2. Within each segment of 4 SWA blocks, only the trailing 2 are stored. With 32 tokens (2 full-attn blocks, 8 SWA blocks): @@ -1574,17 +1621,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): kv_cache_groups=kv_cache_groups, ) - # Verify config: alignment_block_count computed correctly + # Verify config: alignment_chunk_count computed correctly kv_group_configs = runner.connector_scheduler.config.kv_group_configs assert len(kv_group_configs) == 2 # Group 0: full attention -> no alignment skip - assert kv_group_configs[0].alignment_block_count is None - assert kv_group_configs[0].sliding_window_size_in_blocks is None - assert kv_group_configs[0].offloaded_block_size == full_attn_block_size - # Group 1: SWA -> alignment_block_count = 16/4 = 4, tail = 2 - assert kv_group_configs[1].alignment_block_count == 4 - assert kv_group_configs[1].sliding_window_size_in_blocks == 2 - assert kv_group_configs[1].offloaded_block_size == swa_block_size + assert kv_group_configs[0].alignment_chunk_count is None + assert kv_group_configs[0].sliding_window_size_in_chunks is None + assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size + # Group 1: SWA -> alignment_chunk_count = 16/4 = 4, tail = 2 + assert kv_group_configs[1].alignment_chunk_count == 4 + assert kv_group_configs[1].sliding_window_size_in_chunks == 2 + assert kv_group_configs[1].tokens_per_chunk == swa_block_size # Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks # (block_size=4). Decode 1 token to kick off processing (stores are @@ -1602,9 +1649,9 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): ) runner.run( decoded_tokens=[EOS_TOKEN_ID], - # Group 0 (full attn, block_size=16): 2 offloaded blocks + # Group 0 (full attn, block_size=16): 2 offloaded chunks # -> GPU blocks (0, 0) and (0, 1) - # Group 1 (SWA, block_size=4): 8 offloaded blocks, skip first 2 + # Group 1 (SWA, block_size=4): 8 offloaded chunks, skip first 2 # per segment of 4: # Segment 0 (blocks 0-3): skip 0,1 -> store (1, 2), (1, 3) # Segment 1 (blocks 4-7): skip 4,5 -> store (1, 6), (1, 7) @@ -1625,7 +1672,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2 runner.run( decoded_tokens=[EOS_TOKEN_ID], - # Group 0: full prefix lookup hits 2 offloaded blocks + # Group 0: full prefix lookup hits 2 offloaded chunks # -> loads GPU blocks (0, 0), (0, 1) # Group 1: sliding window lookup finds trailing 2 from last segment # (blocks 6, 7 which were stored) @@ -1684,7 +1731,7 @@ def test_stale_sliding_window_block_after_prepare_store_failure( runner.new_request(token_ids=[0] * block_size * 3) # First step: prepare_store FAILS -> offloading delayed. - # next_stored_block_idx stays at 0, block_ids[0] still holds the + # next_stored_chunk_idx stays at 0, block_ids[0] still holds the # original block_id for position 0. runner.manager.prepare_store.side_effect = lambda keys, req_context: None runner.run(decoded_tokens=[0]) @@ -1719,19 +1766,19 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): """When skip_reading_prefix_cache=True, the offloading connector must not load any blocks from CPU even if a matching prefix is cached there.""" block_size = 4 - block_size_factor = 3 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 3 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 runner = request_runner( block_size=block_size, num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) # Populate the CPU offload cache with one block. - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -1747,7 +1794,7 @@ def test_skip_reading_prefix_cache(request_runner, async_scheduling: bool): # The offloading connector must not load anything from CPU, but must # still offload the freshly computed blocks (state management intact). runner.new_request( - token_ids=[0] * offloaded_block_size, + token_ids=[0] * tokens_per_chunk, skip_reading_prefix_cache=True, ) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( @@ -2219,7 +2266,7 @@ def test_eagle_verified_survives_eagle_tighten(self, request_runner): def test_full_attn_store_excludes_trailing_decode_block( self, request_runner, async_scheduling: bool ): - """Eagle full-attention group excludes the trailing block only while + """Eagle full-attention group excludes the trailing chunk only while decoding. Setup: 2 groups — group 0 is normal full-attention, group 1 is @@ -2229,8 +2276,8 @@ def test_full_attn_store_excludes_trailing_decode_block( draft-layer KV is volatile until the next block starts). """ block_size = 4 - block_size_factor = 1 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 1 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 kv_cache_groups = [ @@ -2260,7 +2307,7 @@ def test_full_attn_store_excludes_trailing_decode_block( num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) kv_group_configs = runner.connector_scheduler.config.kv_group_configs @@ -2268,7 +2315,7 @@ def test_full_attn_store_excludes_trailing_decode_block( assert not kv_group_configs[0].is_eagle_group assert kv_group_configs[1].is_eagle_group - runner.new_request(token_ids=[0] * offloaded_block_size * 3) + runner.new_request(token_ids=[0] * tokens_per_chunk * 3) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -2292,7 +2339,7 @@ def test_sw_store_excludes_trailing_decode_block( self, request_runner, async_scheduling: bool ): """Eagle sliding-window group stores all prompt blocks but excludes - the trailing block while decoding.""" + the trailing chunk while decoding.""" block_size = 4 sliding_window = 8 num_gpu_blocks = 100 @@ -2321,7 +2368,7 @@ def test_sw_store_excludes_trailing_decode_block( kv_group_configs = runner.connector_scheduler.config.kv_group_configs assert len(kv_group_configs) == 1 assert kv_group_configs[0].is_eagle_group - assert kv_group_configs[0].sliding_window_size_in_blocks == 2 + assert kv_group_configs[0].sliding_window_size_in_chunks == 2 runner.new_request(token_ids=[0] * block_size * 3) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( @@ -2340,8 +2387,8 @@ def test_single_block_stored_at_end_of_prefill( """An eagle group with a single-block prompt stores it at the end of prefill: prompt blocks are stable, so no tail is held back.""" block_size = 4 - block_size_factor = 1 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 1 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 kv_cache_groups = [ @@ -2362,10 +2409,10 @@ def test_single_block_stored_at_end_of_prefill( num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -2378,18 +2425,18 @@ def test_multichunk_store_no_interior_holes( """Eagle store must not drop interior blocks across prefill chunks. Regression: the trailing-block exclusion (num_blocks - 1) was applied - when collecting keys, but next_stored_block_idx advanced by the - non-decremented count, so the trailing block of every chunked-prefill + when collecting keys, but next_stored_chunk_idx advanced by the + non-decremented count, so the trailing chunk of every chunked-prefill chunk was skipped and never re-considered. With the harness chunk budget (1000 tokens) and block_size 4, a prompt longer than one chunk lost the block at the chunk boundary, leaving a permanent gap that caps prefix reuse at the first hole. Only the trailing decode block may be held back; all other blocks must be stored exactly once (no duplicates from - next_stored_block_idx regressing at the prefill->decode transition). + next_stored_chunk_idx regressing at the prefill->decode transition). """ block_size = 4 - block_size_factor = 1 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 1 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 1000 kv_cache_groups = [ @@ -2409,13 +2456,13 @@ def test_multichunk_store_no_interior_holes( num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) assert runner.connector_scheduler.config.kv_group_configs[0].is_eagle_group # Prompt spans more than one prefill chunk (chunk budget 1000 tokens). num_blocks = 256 - runner.new_request(token_ids=[0] * offloaded_block_size * num_blocks) + runner.new_request(token_ids=[0] * tokens_per_chunk * num_blocks) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -2429,7 +2476,7 @@ def test_multichunk_store_no_interior_holes( for b in t.gpu_blocks ) # The stored blocks must be contiguous from 0: no interior block is - # dropped at a chunk boundary. (The bug left a gap at offloaded block + # dropped at a chunk boundary. (The bug left a gap at offloaded chunk # 249, the tail of the first 1000-token chunk.) assert offsets == list(range(len(offsets))), ( f"interior hole in stored blocks: {offsets}" @@ -2439,14 +2486,14 @@ def test_multichunk_store_no_interior_holes( def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool): """Eagle group constrains load: convergence tightens both groups. - Store 3 offloaded blocks per group (all prompt blocks, so the eagle + Store 3 offloaded chunks per group (all prompt chunks, so the eagle group stores all 3 as well). Then a new request loads from CPU. The eagle group pops its trailing hit block on load, tightening the hit to 2 blocks for both groups. """ block_size = 4 - block_size_factor = 1 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 1 + tokens_per_chunk = block_size * blocks_per_chunk num_gpu_blocks = 100 kv_cache_groups = [ @@ -2476,10 +2523,10 @@ def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool) num_gpu_blocks=num_gpu_blocks, async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) - runner.new_request(token_ids=[0] * offloaded_block_size * 3) + runner.new_request(token_ids=[0] * tokens_per_chunk * 3) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -2497,7 +2544,7 @@ def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool) runner.scheduler.reset_prefix_cache() - runner.new_request(token_ids=[0] * offloaded_block_size * 3 + [1]) + runner.new_request(token_ids=[0] * tokens_per_chunk * 3 + [1]) runner.manager.lookup.return_value = LookupResult.HIT runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) @@ -2527,8 +2574,8 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner): GPU blocks before the store completes. """ block_size = 4 - block_size_factor = 1 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 1 + tokens_per_chunk = block_size * blocks_per_chunk # Use 2 GPU blocks so the second run reuses the same blocks, # triggering a fence-based flush of the in-flight job from run 1. @@ -2536,11 +2583,11 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner): block_size=block_size, num_gpu_blocks=2, async_scheduling=False, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) # 4 prompt tokens → 1 GPU block (block 0) - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -2576,7 +2623,7 @@ def capture_fence(): # Run 2: block reuse triggers fence-based flush → cleanup. runner.scheduler.reset_prefix_cache() - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -2603,26 +2650,26 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner): - Run 3: block reuse → both jobs flushed via fence """ block_size = 4 - block_size_factor = 1 - offloaded_block_size = block_size * block_size_factor + blocks_per_chunk = 1 + tokens_per_chunk = block_size * blocks_per_chunk # 4 GPU blocks: block 0 is null, blocks 1-3 are usable. runner = request_runner( block_size=block_size, num_gpu_blocks=4, async_scheduling=False, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, ) # Prompt: 4 tokens → block 1 - runner.new_request(token_ids=[0] * offloaded_block_size) + runner.new_request(token_ids=[0] * tokens_per_chunk) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) # Run 1: 4 decoded tokens → block 2 full → job_0 created for block 1. runner.run( - decoded_tokens=[0] * offloaded_block_size, + decoded_tokens=[0] * tokens_per_chunk, complete_transfers=False, ) assert len(runner.connector_scheduler._jobs) >= 1 @@ -2630,7 +2677,7 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner): # Run 2: 4 more tokens + EOS → block 3 full → more jobs created. # Request finishes → all jobs registered in fence. runner.run( - decoded_tokens=[0] * offloaded_block_size + [EOS_TOKEN_ID], + decoded_tokens=[0] * tokens_per_chunk + [EOS_TOKEN_ID], complete_transfers=False, ) num_jobs = len(runner.connector_scheduler._jobs) @@ -2638,7 +2685,7 @@ def test_multiple_in_flight_stores_all_flushed_by_fence(request_runner): # Run 3: block reuse → fence flushes both jobs. runner.scheduler.reset_prefix_cache() - runner.new_request(token_ids=[0] * offloaded_block_size * 3) + runner.new_request(token_ids=[0] * tokens_per_chunk * 3) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py index 81c00266cfed..25bb664ff402 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_worker.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_worker.py @@ -96,11 +96,12 @@ def _make_worker(kv_cache_config: KVCacheConfig): ) spec = MagicMock(spec=OffloadingSpec) - spec.kv_cache_config = kv_cache_config - spec.vllm_config = MagicMock() spec.get_worker.return_value = MagicMock() - worker = OffloadingConnectorWorker(spec=spec) + worker = OffloadingConnectorWorker( + spec=spec, + kv_cache_config=kv_cache_config, + ) worker.worker = MagicMock() return worker, spec diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 73ea5e2be1d2..b878e294a6ef 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -17,7 +17,6 @@ from vllm.config import ( KVEventsConfig, KVTransferConfig, - VllmConfig, set_current_vllm_config, ) from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole @@ -57,6 +56,7 @@ TransferResult, make_offload_key, ) +from vllm.v1.kv_offload.config import OffloadingConfig from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -123,8 +123,8 @@ def wait(self, job_ids: set[int]) -> None: class MockOffloadingSpec(OffloadingSpec): - def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): - super().__init__(vllm_config, kv_cache_config) + def __init__(self, config: OffloadingConfig): + super().__init__(config) self.manager = MagicMock(spec=OffloadingManager) self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) @@ -175,17 +175,17 @@ def __init__( self, block_size: int, num_gpu_blocks: int, - block_size_factor: int = 1, + blocks_per_chunk: int = 1, async_scheduling: bool = True, kv_cache_groups: list[KVCacheGroupSpec] | None = None, extra_config_overrides: dict[str, Any] | None = None, ): - assert block_size_factor == 1 or kv_cache_groups is None, ( - "block_size_factor > 1 requires all groups to have the same " + assert blocks_per_chunk == 1 or kv_cache_groups is None, ( + "blocks_per_chunk > 1 requires all groups to have the same " "block size, so kv_cache_groups must be None (use default group)" ) - self.block_size_factor: int = block_size_factor + self.blocks_per_chunk: int = blocks_per_chunk self.block_size: int = block_size self.num_gpu_blocks: int = num_gpu_blocks self.async_scheduling: bool = async_scheduling @@ -208,8 +208,8 @@ def __init__( # opt-out tests override this to cover the legacy placeholders. "self_describing_kv_events": True, } - if block_size_factor > 1: - extra_config["block_size"] = block_size * block_size_factor + if blocks_per_chunk > 1: + extra_config["block_size"] = block_size * blocks_per_chunk if extra_config_overrides: extra_config.update(extra_config_overrides) @@ -313,11 +313,9 @@ def __init__( self.connector_scheduler.config.kv_group_configs, kv_cache_config.kv_cache_groups, ): - gpu_block_size = kv_cache_group.kv_cache_spec.block_size - assert group_config.gpu_block_size == gpu_block_size - assert ( - group_config.offloaded_block_size == gpu_block_size * block_size_factor - ) + tokens_per_block = kv_cache_group.kv_cache_spec.block_size + assert group_config.tokens_per_block == tokens_per_block + assert group_config.tokens_per_chunk == tokens_per_block * blocks_per_chunk # extract OffloadingSpec of worker_connector connector_worker = self.worker_connector.connector_worker @@ -389,7 +387,7 @@ def _parse_transfers(self): for block_id in dst_spec.block_ids: self.flushed_gpu_blocks.add(self.gpu_blocks[block_id.item()]) - block_size_factor = self.block_size_factor + blocks_per_chunk = self.blocks_per_chunk for src_spec, dst_spec in self.offloading_spec.get_completed_transfers(): if isinstance(src_spec, GPULoadStoreSpec): @@ -412,7 +410,7 @@ def _parse_transfers(self): # list of (offload_key, sub_block_offset) offload_addresses: list[Any] = [] for offload_key in offload_spec.offload_keys: - for sub_block_idx in range(block_size_factor): + for sub_block_idx in range(blocks_per_chunk): offload_addresses.append((offload_key, sub_block_idx)) assert gpu_spec.block_indices is not None @@ -426,7 +424,7 @@ def _parse_transfers(self): gpu_block_end_offset = gpu_block_offset + group_size assert gpu_block_end_offset <= len(gpu_blocks) - offload_addresses_to_skip = logical_offset % block_size_factor + offload_addresses_to_skip = logical_offset % blocks_per_chunk offload_addresses_end_offset = ( offload_address_offset + offload_addresses_to_skip + group_size ) @@ -651,14 +649,14 @@ def runner_factory( block_size, num_gpu_blocks, async_scheduling, - block_size_factor=1, + blocks_per_chunk=1, kv_cache_groups=None, extra_config_overrides=None, ): runner = RequestRunner( block_size=block_size, num_gpu_blocks=num_gpu_blocks, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, async_scheduling=async_scheduling, kv_cache_groups=kv_cache_groups, extra_config_overrides=extra_config_overrides, diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index 49ca683de732..6155932ea720 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -670,16 +670,25 @@ def test_moriio_wrapper_rejects_invalid_messages(role, payload, match): wrapper._handle_message(payload) -def test_block_id_length_mismatch_raises_value_error(): +def test_local_block_ids_longer_than_remote_raises_value_error(): cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) worker = _worker({"layer": cache}, {"layer": _full_spec()}) - with pytest.raises(ValueError, match="must have the same length"): + with pytest.raises(ValueError, match="longer than remote_block_ids"): moriio_layout.compute_block_transfer_offsets( "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks ) +def test_empty_local_block_ids_is_free_only_noop(): + cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) + worker = _worker({"layer": cache}, {"layer": _full_spec()}) + + assert moriio_layout.compute_block_transfer_offsets( + "layer", cache, worker.layer_to_spec, [], [4, 5], _remote_meta().num_blocks + ) == ([], [], []) + + def test_registration_regions_do_not_split_interleaved_or_mla_cache(): separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16) interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16) 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 9abc9b957c4a..6488135690d8 100644 --- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +++ b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py @@ -253,6 +253,7 @@ def shutdown(self): worker.transfer_id_to_request_id = {"tx-fanin": "req-fanin"} worker._consumer_notification_counts = {} worker._completed_consumer_notifications = set() + worker._pending_unmapped_acks = [] assert worker.get_finished() == (set(), set()) assert worker._consumer_notification_counts == {"tx-fanin": 1} diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index c1088f1c6e88..31d34c69b3fb 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -13,6 +13,7 @@ from unittest.mock import MagicMock, patch import msgspec +import numpy as np import pytest import ray import torch @@ -750,7 +751,10 @@ def test_prefill_tp_size_greater_than_decode_tp_size( worker.block_len_per_layer = [4096 * worker.block_size] worker.num_blocks = 1 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks - worker.src_blocks_data = [(0, worker.block_len_per_layer[0], worker.tp_rank)] + worker.src_blocks_data = np.array( + [(0, worker.block_len_per_layer[0], worker.tp_rank)], + dtype=np.uint64, + ) worker.num_descs = len(worker.src_blocks_data) def check_handshake(remote_tp_size: int): @@ -1124,8 +1128,8 @@ def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices( == worker._mamba_ssm_size[1] ) - assert worker._build_fa_remote(plan, meta, block_size_ratio=1) == [ - (0x1000 + local_block_len, local_block_len, 0) + assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [ + [0x1000 + local_block_len, local_block_len, 0] ] @patch( @@ -1158,10 +1162,13 @@ def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init): worker._region_is_mla = [False, True] worker.num_blocks = 1 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks - worker.src_blocks_data = [ - (0, fa_len, worker.tp_rank), - (0, idx_len, worker.tp_rank), - ] + worker.src_blocks_data = np.array( + [ + (0, fa_len, worker.tp_rank), + (0, idx_len, worker.tp_rank), + ], + dtype=np.uint64, + ) worker.num_descs = len(worker.src_blocks_data) # D_TP=2, P_TP=1 -> tp_ratio=2. SPLIT region scales by tp_ratio; diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 16c47962aefb..4945942ba3a8 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -86,7 +86,9 @@ def test_logical_to_kernel_block_ids_with_hma(): # Test conversion: FA + SW group logical_block_ids = [[0, 1, 2], [3, 4]] - kernel_block_ids = worker._logical_to_kernel_block_ids(logical_block_ids) + kernel_block_ids = worker._logical_to_kernel_block_ids( + logical_block_ids, worker._physical_blocks_per_logical_kv_block + ) expected_kernel_block_ids = [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]] assert kernel_block_ids == expected_kernel_block_ids, ( @@ -179,7 +181,7 @@ def test_read_blocks_for_req_expands_remote_ids( """_read_blocks_for_req must expand remote logical block IDs to kernel block IDs when kernel block size != logical block size. - The hot path always calls _logical_to_remote_kernel_block_ids with + The hot path always calls _logical_to_kernel_block_ids with remote_info.remote_physical_blocks_per_logical (model-agnostic). """ from unittest.mock import MagicMock @@ -750,6 +752,94 @@ def test_nixl_metadata_hybrid_ssm_block_ids(): assert len(req_meta.remote.block_ids[0]) != len(req_meta.remote.block_ids[1]) +class _FakeBlock: + def __init__(self, block_id): + self.block_id = block_id + + +class _FakeSingleTypeManager: + def __init__(self, records, block_size, block_ids): + self.records_new_block_ids = records + self.block_size = block_size + self.req_to_blocks = {"req-1": [_FakeBlock(b) for b in block_ids]} + self.new_block_ids: list[int] = [] + + def take_new_block_ids(self): + ids = self.new_block_ids + self.new_block_ids = [] + return ids + + +def _make_fake_kv_cache_manager(): + from unittest.mock import MagicMock + + from vllm.v1.core.kv_cache_manager import KVCacheManager + + manager = object.__new__(KVCacheManager) + manager.coordinator = MagicMock() + manager.coordinator.single_type_managers = ( + _FakeSingleTypeManager(True, 16, [10, 11, 12, 13, 14, 15]), # attention + _FakeSingleTypeManager(False, 16, [20, 21, 22, 23, 24, 25]), # mamba + ) + return manager + + +@pytest.mark.cpu_test +def test_zeroing_block_ids_cover_only_loaded_attention_blocks(): + """Only zero-recorded (attention) groups contribute, sliced to the + externally-loaded token range; Mamba state blocks are never zeroed.""" + manager = _make_fake_kv_cache_manager() + + # Tokens [0, 16) are locally cached; the load covers tokens [16, 56). + assert manager.get_zeroing_block_ids_in_range("req-1", 16, 56) == [11, 12, 13] + + +@pytest.mark.cpu_test +def test_scheduler_filters_connector_loaded_blocks_from_zeroing(): + """Blocks that will be loaded by the connector must not be zeroed.""" + from vllm.v1.core.sched.scheduler import Scheduler + + class FakeKVCacheManager: + def take_new_block_ids(self): + return [9, 10, 11, 12] + + scheduler = object.__new__(Scheduler) + scheduler.needs_kv_cache_zeroing = True + scheduler.kv_cache_manager = FakeKVCacheManager() + scheduler._skip_zero_block_ids = {10, 12} + + assert scheduler._get_new_block_ids_to_zero() == [9, 11] + assert not scheduler._skip_zero_block_ids + + +@pytest.mark.cpu_test +def test_failed_load_rezeroes_unwritten_skipped_blocks(): + """A failed async load leaves zeroing-skipped blocks unwritten beyond + the valid prefix; they must be zeroed before local recompute.""" + from unittest.mock import MagicMock + + from vllm.v1.core.sched.scheduler import Scheduler + + scheduler = object.__new__(Scheduler) + scheduler.connector = MagicMock() + scheduler.needs_kv_cache_zeroing = True + scheduler.kv_cache_manager = _make_fake_kv_cache_manager() + scheduler.kv_cache_manager.cache_blocks = MagicMock() + scheduler.failed_recving_kv_req_ids = {"req-1"} + scheduler.finished_recving_kv_req_ids = {"req-1"} + + request = MagicMock() + request.request_id = "req-1" + request.num_computed_tokens = 48 # Truncated at the first invalid block. + + scheduler._update_waiting_for_remote_kv(request) + + # Attention blocks covering tokens >= 48 are re-recorded for zeroing + # and flow into the next step's zero list; Mamba blocks are not. + scheduler._skip_zero_block_ids = set() + assert scheduler._get_new_block_ids_to_zero() == [13, 14, 15] + + # ── Mamba N-1 prefill tests ────────────────────────────────────────────── @@ -1140,7 +1230,7 @@ def test_derive_mamba_conv_split( ), ], ) -def test_logical_to_remote_kernel_block_ids( +def test_logical_to_kernel_block_ids_with_remote_ratio( mamba_enabled, swa_enabled, local_physical_per_logical, @@ -1148,7 +1238,7 @@ def test_logical_to_remote_kernel_block_ids( logical_block_ids, expected_kernel_block_ids, ): - """Verify _logical_to_remote_kernel_block_ids uses the remote + """Verify _logical_to_kernel_block_ids uses the remote physical_per_logical for FA expansion, not the local one. This was the root cause of silent accuracy corruption in Qwen3.5 @@ -1169,7 +1259,7 @@ def test_logical_to_remote_kernel_block_ids( swa_enabled=swa_enabled, ) - result = worker._logical_to_remote_kernel_block_ids( + result = worker._logical_to_kernel_block_ids( logical_block_ids, remote_physical_per_logical, ) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 3c9cce1bd686..ad5ccbbbe86c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -487,7 +487,7 @@ def test_start_load_kv_enqueues_to_writer(self): # path here. w._send_heartbeats = lambda metadata: None # Stub logical-to-kernel mapping used by reqs_to_recv. - w._logical_to_kernel_block_ids = lambda x: x + w._logical_to_kernel_block_ids = lambda x, ratio: x meta = NixlConnectorMetadata() meta.push_registrations = { @@ -775,7 +775,7 @@ def test_start_load_kv_with_empty_metadata_is_noop(self): """Empty metadata must not wake the writer or enqueue anything.""" w = _StubWriterWorker.fresh() w._send_heartbeats = lambda metadata: None - w._logical_to_kernel_block_ids = lambda x: x + w._logical_to_kernel_block_ids = lambda x, ratio: x meta = NixlConnectorMetadata() w.start_load_kv(meta) @@ -944,7 +944,7 @@ def _mla_worker_writing_to(d_ranks): rank_offset_factor=0, ) } - w._logical_to_remote_kernel_block_ids = lambda block_ids, ratio: block_ids + w._logical_to_kernel_block_ids = lambda block_ids, ratio: block_ids w.dst_xfer_side_handles = {engine_id: {r: 1000 + r for r in d_ranks}} w.src_xfer_handles_by_block_size = {16: 2000} w._remote_agents = {engine_id: {(0, r): f"agent-{r}" for r in d_ranks}} diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 16420b164b2b..a323ae7ce523 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -45,7 +45,7 @@ # Falcon-H1: parallel hybrid (every layer has both attention and SSM). # The mamba and attention groups end up with different GPU block sizes # after page-size unification, so we leave cpu_block_size=None - # (block_size_factor stays 1). + # (blocks_per_chunk stays 1). ("tiiuae/Falcon-H1-0.5B-Instruct", None, None, True), ] diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 5ab6b68400c9..7c735098230e 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -11,6 +11,7 @@ from types import SimpleNamespace +import numpy as np import pytest from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import ( @@ -102,7 +103,10 @@ def test_build_src_split_handles(self, remote_tp_size): ) worker = _make_mock_worker_for_splits((FullAttentionSpec,)) - src_blocks_data = [(0x2000 + i * 1024, 1024, 0) for i in range(8)] + src_blocks_data = np.array( + [(0x2000 + i * 1024, 1024, 0) for i in range(8)], + dtype=np.uint64, + ) num_descs = len(src_blocks_data) splits = list( worker._build_local_splits_from_plan( @@ -135,11 +139,14 @@ def test_fa_and_ssm_different_split_factors(self): worker = _make_mock_worker_for_splits((FullAttentionSpec, MambaSpec)) # 2 FA descs + 1 SSM desc - src_blocks_data = [ - (1000, 200, 0), # FA desc 0 - (2000, 200, 0), # FA desc 1 - (3000, 400, 0), # SSM desc 0 - ] + src_blocks_data = np.array( + [ + (1000, 200, 0), # FA desc 0 + (2000, 200, 0), # FA desc 1 + (3000, 400, 0), # SSM desc 0 + ], + dtype=np.uint64, + ) splits = list(worker._build_local_splits_from_plan(plan, src_blocks_data, 2)) diff --git a/tests/v1/kv_offload/cpu/test_gpu_worker.py b/tests/v1/kv_offload/cpu/test_gpu_worker.py index 12dbc57fe97c..2f1ce67e9c41 100644 --- a/tests/v1/kv_offload/cpu/test_gpu_worker.py +++ b/tests/v1/kv_offload/cpu/test_gpu_worker.py @@ -23,7 +23,7 @@ NUM_GPU_BLOCKS = [64] NUM_CPU_BLOCKS = [256] GPU_PAGE_SIZES = [512, 1024] -BLOCK_SIZE_FACTORS = [1, 3] +BLOCKS_PER_CHUNK_VALUES = [1, 3] NUM_TENSORS = [4] SEEDS = [0] DEVICE_TYPE = current_platform.device_type @@ -35,7 +35,7 @@ @pytest.mark.parametrize("gpu_to_cpu", [True, False]) @pytest.mark.parametrize("num_mappings", NUM_MAPPINGS) @pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES) -@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS) +@pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES) @pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS) @pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS) @pytest.mark.parametrize("num_tensors", NUM_TENSORS) @@ -48,7 +48,7 @@ def test_transfer( gpu_to_cpu: bool, num_mappings: int, gpu_page_size_bytes: int, - block_size_factor: int, + blocks_per_chunk: int, num_gpu_blocks: int, num_cpu_blocks: int, num_tensors: int, @@ -92,7 +92,7 @@ def test_transfer( mmap_region: SharedOffloadRegion | None = None if use_shared_memory: cpu_page_size = round_up( - gpu_page_size_bytes * num_tensors * block_size_factor, + gpu_page_size_bytes * num_tensors * blocks_per_chunk, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT, ) mmap_region = SharedOffloadRegion( @@ -105,25 +105,25 @@ def test_transfer( worker = CPUOffloadingWorker( kv_caches=kv_caches, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, num_cpu_blocks=num_cpu_blocks, mmap_region=mmap_region, ) # select block mappings - gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * block_size_factor) + gpu_blocks = random.sample(range(num_gpu_blocks), num_mappings * blocks_per_chunk) cpu_blocks = random.sample(range(num_cpu_blocks), num_mappings) # expand cpu blocks to gpu-page granularity for uniform comparison: - # each cpu block maps to block_size_factor consecutive sub-blocks + # each cpu block maps to blocks_per_chunk consecutive sub-blocks cpu_blocks_expanded = [ - cpu_block * block_size_factor + j + cpu_block * blocks_per_chunk + j for cpu_block in cpu_blocks - for j in range(block_size_factor) + for j in range(blocks_per_chunk) ] # maybe skip some GPU blocks to test reading/writing from the middle of a CPU block - blocks_to_skip = block_size_factor - 1 + blocks_to_skip = blocks_per_chunk - 1 if blocks_to_skip > 0: gpu_blocks = gpu_blocks[blocks_to_skip:] cpu_blocks_expanded = cpu_blocks_expanded[blocks_to_skip:] @@ -214,7 +214,7 @@ def test_transfer( @pytest.mark.parametrize("gpu_to_cpu", [True, False]) @pytest.mark.parametrize("num_mappings_per_group", NUM_MAPPINGS_PER_GROUP) @pytest.mark.parametrize("gpu_page_size_bytes", GPU_PAGE_SIZES) -@pytest.mark.parametrize("block_size_factor", BLOCK_SIZE_FACTORS) +@pytest.mark.parametrize("blocks_per_chunk", BLOCKS_PER_CHUNK_VALUES) @pytest.mark.parametrize("num_gpu_blocks", NUM_GPU_BLOCKS) @pytest.mark.parametrize("num_cpu_blocks", NUM_CPU_BLOCKS) @pytest.mark.parametrize("seed", SEEDS) @@ -225,7 +225,7 @@ def test_transfer_multi_group( gpu_to_cpu: bool, num_mappings_per_group: int, gpu_page_size_bytes: int, - block_size_factor: int, + blocks_per_chunk: int, num_gpu_blocks: int, num_cpu_blocks: int, seed: int, @@ -234,7 +234,7 @@ def test_transfer_multi_group( """Test transfers with three KV cache groups: - Group 0: aligned transfer with num_mappings_per_group blocks - Group 1: zero blocks (empty group) - - Group 2: unaligned CPU->GPU transfer (logical_offset=block_size_factor-1, + - Group 2: unaligned CPU->GPU transfer (logical_offset=blocks_per_chunk-1, causing the implementation to skip source sub-blocks) with num_mappings_per_group blocks """ @@ -275,7 +275,7 @@ def test_transfer_multi_group( worker = CPUOffloadingWorker( kv_caches=canonical_kv_caches, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, num_cpu_blocks=num_cpu_blocks, ) @@ -283,7 +283,7 @@ def test_transfer_multi_group( group_sizes_in_cpu_blocks = [num_mappings_per_group, 0, num_mappings_per_group] total_cpu_blocks = sum(group_sizes_in_cpu_blocks) - total_gpu_blocks_needed = total_cpu_blocks * block_size_factor + total_gpu_blocks_needed = total_cpu_blocks * blocks_per_chunk gpu_blocks_all = random.sample(range(num_gpu_blocks), total_gpu_blocks_needed) cpu_blocks_all = random.sample(range(num_cpu_blocks), total_cpu_blocks) @@ -293,7 +293,7 @@ def test_transfer_multi_group( gpu_offset = 0 cpu_offset = 0 for size in group_sizes_in_cpu_blocks: - gpu_count = size * block_size_factor + gpu_count = size * blocks_per_chunk gpu_blocks_per_group.append(gpu_blocks_all[gpu_offset : gpu_offset + gpu_count]) cpu_blocks_per_group.append(cpu_blocks_all[cpu_offset : cpu_offset + size]) gpu_offset += gpu_count @@ -302,15 +302,15 @@ def test_transfer_multi_group( # expand cpu blocks to gpu-page granularity cpu_blocks_expanded_per_group = [ [ - cpu_block * block_size_factor + j + cpu_block * blocks_per_chunk + j for cpu_block in cpu_blocks - for j in range(block_size_factor) + for j in range(blocks_per_chunk) ] for cpu_blocks in cpu_blocks_per_group ] # skip sub-blocks from group 2 to test unaligned transfers. - sub_blocks_to_skip = block_size_factor - 1 # e.g. 2 when block_size_factor=3 + sub_blocks_to_skip = blocks_per_chunk - 1 # e.g. 2 when blocks_per_chunk=3 if sub_blocks_to_skip > 0: gpu_blocks_per_group[2] = gpu_blocks_per_group[2][ sub_blocks_to_skip:-sub_blocks_to_skip @@ -347,7 +347,7 @@ def test_transfer_multi_group( cpu_blocks_expanded_per_group, gpu_blocks_per_group ) ] - num_dst_sub_blocks = num_cpu_blocks * block_size_factor + num_dst_sub_blocks = num_cpu_blocks * blocks_per_chunk else: handler = worker._load_handler src_spec = CPULoadStoreSpec(cpu_blocks) diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index b051f44ecfdc..570924cfccc0 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -11,17 +11,33 @@ 4. Error paths — unregistered specs, missing config, duplicate registration. """ +from typing import cast +from unittest.mock import MagicMock, patch + import pytest import torch -from vllm.config import KVTransferConfig +from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( + build_offloading_config, +) +from vllm.platforms import current_platform from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + MLAAttentionSpec, + SlidingWindowSpec, ) -from vllm.v1.kv_offload.base import OffloadingHistogramMetadata, OffloadingSpec +from vllm.v1.kv_offload.base import ( + CanonicalKVCaches, + OffloadingHistogramMetadata, + OffloadingManager, + OffloadingSpec, + OffloadingWorker, +) +from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec from vllm.v1.kv_offload.factory import OffloadingSpecFactory from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec @@ -39,6 +55,17 @@ def restore_registry(): OffloadingSpecFactory._registry = original +def _get_extra_config(config: VllmConfig) -> dict: + assert config.kv_transfer_config is not None + return config.kv_transfer_config.kv_connector_extra_config + + +def _create_spec(config: VllmConfig, kv_cache_config: KVCacheConfig) -> OffloadingSpec: + return OffloadingSpecFactory.create_spec( + build_offloading_config(config, kv_cache_config) + ) + + def _make_vllm_config( spec_name: str | None = "CPUOffloadingSpec", cpu_bytes_to_use: int | None = None, @@ -95,6 +122,46 @@ def _make_vllm_config( ) +def _make_layout_vllm_config( + spec_name: str = "CPUOffloadingSpec", + cpu_bytes_to_use: int | None = None, + extra_config: dict | None = None, + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + prefill_context_parallel_size: int = 1, + decode_context_parallel_size: int = 1, +) -> VllmConfig: + config = MagicMock() + config.cache_config.block_size = 16 + config.cache_config.enable_prefix_caching = True + config.cache_config.prefix_match_unit = None + config.cache_config.cache_dtype = torch.float16 + config.model_config.model = "test-model" + world_size = ( + tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size + ) + with patch.object(current_platform, "device_count", return_value=world_size): + config.parallel_config = ParallelConfig( + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + prefill_context_parallel_size=prefill_context_parallel_size, + decode_context_parallel_size=decode_context_parallel_size, + ) + config.kv_events_config = None + config.use_v2_model_runner = False + + connector_extra_config = dict(extra_config or {}) + connector_extra_config["spec_name"] = spec_name + if cpu_bytes_to_use is not None: + connector_extra_config["cpu_bytes_to_use"] = cpu_bytes_to_use + config.kv_transfer_config = KVTransferConfig( + kv_connector="OffloadingConnector", + kv_role="kv_both", + kv_connector_extra_config=connector_extra_config, + ) + return cast(VllmConfig, config) + + def _make_kv_cache_config(): """Build a minimal KVCacheConfig with one KV cache tensor.""" num_blocks = 16 @@ -122,6 +189,78 @@ def _make_kv_cache_config(): ) +def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig: + num_blocks = 4 + if packed: + kv_cache_tensors = [ + KVCacheTensor( + size=64, + shared_by=[layer_name], + block_stride=16, + ) + for layer_name in ("layer0", "layer1") + ] + else: + kv_cache_tensors = [ + KVCacheTensor(size=40, shared_by=["layer0"]), + KVCacheTensor(size=24, shared_by=["layer1"]), + ] + + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=kv_cache_tensors, + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0", "layer1"], + FullAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + ], + ) + + +def _make_hybrid_kv_cache_config() -> KVCacheConfig: + return KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor(size=40, shared_by=["full_layer"]), + KVCacheTensor(size=24, shared_by=["mla_layer"]), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full_layer"], + FullAttentionSpec( + block_size=12, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mla_layer"], + MLAAttentionSpec( + block_size=16, + num_kv_heads=1, + head_size=576, + dtype=torch.float32, + ), + ), + ], + ) + + +class SingleArgExternalOffloadingSpec(OffloadingSpec): + def get_manager(self) -> OffloadingManager: + raise NotImplementedError + + def get_worker(self, kv_caches: CanonicalKVCaches) -> OffloadingWorker: + raise NotImplementedError + + # --------------------------------------------------------------------------- # Pre-registration integrity (CI sentinel) # --------------------------------------------------------------------------- @@ -154,7 +293,7 @@ def test_tiering_spec_registered(): def test_get_spec_cls_returns_registered_class(): """Registered spec_name returns correct class.""" config = _make_vllm_config(spec_name="CPUOffloadingSpec") - spec_cls = OffloadingSpecFactory.get_spec_cls(config) + spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) assert spec_cls is CPUOffloadingSpec @@ -162,7 +301,7 @@ def test_get_spec_cls_default_to_cpu(): """Default spec_name (absent from config) resolves to CPUOffloadingSpec.""" config = _make_vllm_config(spec_name=None) config.kv_transfer_config.kv_connector_extra_config.pop("spec_name", None) - spec_cls = OffloadingSpecFactory.get_spec_cls(config) + spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) assert spec_cls is CPUOffloadingSpec @@ -176,16 +315,193 @@ def test_create_cpu_offloading_spec_end_to_end(): Verifies: - cpu_bytes_to_use validation and num_blocks calculation - - block_size % hash_block_size assertion + - block_size % tokens_per_hash assertion - spec instance is CPUOffloadingSpec """ config = _make_vllm_config(cpu_bytes_to_use=65536) kv_cache_config = _make_kv_cache_config() - spec = OffloadingSpecFactory.create_spec(config, kv_cache_config) + spec = _create_spec(config, kv_cache_config) assert isinstance(spec, CPUOffloadingSpec) assert spec.num_blocks > 0 +@pytest.mark.parametrize("packed", [False, True]) +def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool): + cpu_bytes_to_use = 1920 + config = _make_layout_vllm_config( + cpu_bytes_to_use=cpu_bytes_to_use, + extra_config={"block_size": 32}, + tensor_parallel_size=3, + pipeline_parallel_size=2, + ) + + spec = _create_spec(config, _make_sizing_kv_cache_config(packed)) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.cpu_page_size_per_worker == 32 + assert spec.kv_bytes_per_chunk == 192 + assert spec.num_blocks == cpu_bytes_to_use // 192 + + +def test_cpu_spec_rejects_partially_packed_tensor_layout(): + config = _make_layout_vllm_config(cpu_bytes_to_use=65536) + kv_cache_config = _make_sizing_kv_cache_config(packed=False) + kv_cache_config.kv_cache_tensors[0].block_stride = 16 + + with pytest.raises(AssertionError): + _create_spec(config, kv_cache_config) + + +def test_cpu_spec_zero_blocks_skips_tensor_layout_validation(): + config = _make_layout_vllm_config(cpu_bytes_to_use=65536) + kv_cache_config = _make_sizing_kv_cache_config(packed=False) + kv_cache_config.num_blocks = 0 + kv_cache_config.kv_cache_tensors[0].block_stride = 16 + + spec = _create_spec(config, kv_cache_config) + + assert isinstance(spec, CPUOffloadingSpec) + assert spec.cpu_page_size_per_worker == 0 + assert spec.kv_bytes_per_chunk == 0 + assert spec.num_blocks == 0 + + +def test_tiering_spec_aligns_row_size(): + alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT + cpu_bytes_to_use = alignment * 3 + config = _make_layout_vllm_config( + spec_name="TieringOffloadingSpec", + cpu_bytes_to_use=cpu_bytes_to_use, + extra_config={"block_size": 32}, + tensor_parallel_size=3, + pipeline_parallel_size=2, + ) + + spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False)) + + assert isinstance(spec, TieringOffloadingSpec) + assert spec.cpu_page_size_per_worker == 32 + assert spec.kv_bytes_per_chunk == alignment + assert spec.num_blocks == cpu_bytes_to_use // alignment + + +def test_offloading_spec_resolves_prefill_context_parallel_block_sizes(): + config = _make_layout_vllm_config( + cpu_bytes_to_use=65536, + extra_config={"block_size": 64}, + prefill_context_parallel_size=2, + ) + + spec = _create_spec(config, _make_kv_cache_config()) + + assert spec.tokens_per_block == (32,) + assert spec.tokens_per_hash == 32 + assert spec.blocks_per_chunk == 2 + + +def test_offloading_config_preserves_data_parallel_index(): + config = _make_layout_vllm_config() + config.parallel_config.data_parallel_index = 2 + + offloading_config = build_offloading_config(config, _make_kv_cache_config()) + + assert offloading_config.parallel.data_parallel_index == 2 + + +def test_offloading_spec_resolves_heterogeneous_hybrid_block_sizes(): + config = _make_layout_vllm_config(cpu_bytes_to_use=65536) + config.cache_config.block_size = 4 + + spec = _create_spec(config, _make_hybrid_kv_cache_config()) + + assert spec.tokens_per_block == (12, 16) + assert spec.tokens_per_hash == 4 + assert spec.blocks_per_chunk == 1 + + +def _full_attention_spec(block_size: int = 16) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=block_size, num_kv_heads=4, head_size=128, dtype=torch.float32 + ) + + +def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool: + config = _make_layout_vllm_config() + kv_cache_config = KVCacheConfig( + num_blocks=0, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups + ) + offloading_config = build_offloading_config(config, kv_cache_config) + return offloading_config.parallel.is_parallelism_agnostic + + +def test_parallelism_agnostic_for_single_full_attention_group(): + assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())]) + + +@pytest.mark.parametrize( + "kv_cache_groups", + [ + # MLA latent KV is replicated per rank, never head-sharded. + [ + KVCacheGroupSpec( + ["l0"], + MLAAttentionSpec( + block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32 + ), + ) + ], + # Sliding window is not full attention. + [ + KVCacheGroupSpec( + ["l0"], + SlidingWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=128, + dtype=torch.float32, + sliding_window=128, + ), + ) + ], + # Hybrid model: more than one KV cache group. + [ + KVCacheGroupSpec(["l0"], _full_attention_spec()), + KVCacheGroupSpec(["l1"], _full_attention_spec()), + ], + ], +) +def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): + assert not _parallelism_agnostic(kv_cache_groups) + + +def test_parallelism_agnostic_disabled_on_v2_model_runner(): + config = _make_layout_vllm_config() + config.use_v2_model_runner = True + kv_cache_config = KVCacheConfig( + num_blocks=0, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())], + ) + offloading_config = build_offloading_config(config, kv_cache_config) + assert not offloading_config.parallel.is_parallelism_agnostic + + +def test_create_dynamic_spec_receives_translated_config(): + config = _make_layout_vllm_config( + spec_name="SingleArgExternalOffloadingSpec", + extra_config={ + "spec_module_path": "tests.v1.kv_offload.test_factory", + }, + ) + kv_cache_config = _make_kv_cache_config() + offloading_config = build_offloading_config(config, kv_cache_config) + + spec = OffloadingSpecFactory.create_spec(offloading_config) + + assert isinstance(spec, SingleArgExternalOffloadingSpec) + assert spec.config is offloading_config + + # --------------------------------------------------------------------------- # Dynamic import via spec_module_path # --------------------------------------------------------------------------- @@ -205,7 +521,7 @@ def test_dynamic_load_via_spec_module_path(): config.kv_transfer_config.kv_connector_extra_config["spec_module_path"] = ( "vllm.v1.kv_offload.cpu.spec" ) - spec_cls = OffloadingSpecFactory.get_spec_cls(config) + spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) assert spec_cls is CPUOffloadingSpec @@ -218,12 +534,12 @@ def test_unregistered_spec_without_module_path_raises(): """spec_name not in registry + no spec_module_path → ValueError.""" config = _make_vllm_config(spec_name="NonexistentSpec") with pytest.raises(ValueError, match="Unsupported spec type"): - OffloadingSpecFactory.get_spec_cls(config) + OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) # create_spec should also fail (calls get_spec_cls internally) kv_cache_config = _make_kv_cache_config() with pytest.raises(ValueError, match="Unsupported spec type"): - OffloadingSpecFactory.create_spec(config, kv_cache_config) + _create_spec(config, kv_cache_config) def test_cpu_spec_missing_cpu_bytes_to_use_raises(): @@ -232,7 +548,7 @@ def test_cpu_spec_missing_cpu_bytes_to_use_raises(): config.kv_transfer_config.kv_connector_extra_config.pop("cpu_bytes_to_use", None) kv_cache_config = _make_kv_cache_config() with pytest.raises(Exception, match="cpu_bytes_to_use must be specified"): - OffloadingSpecFactory.create_spec(config, kv_cache_config) + _create_spec(config, kv_cache_config) def test_duplicate_registration_raises(): @@ -253,7 +569,7 @@ def test_build_metric_definitions_below_threshold(): from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics config = _make_vllm_config(store_threshold=1) - spec_cls = OffloadingSpecFactory.get_spec_cls(config) + spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) metrics = spec_cls.build_metric_definitions( config.kv_transfer_config.kv_connector_extra_config ) @@ -266,7 +582,7 @@ def test_build_metric_definitions_allocation_size_histogram(): from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics config = _make_vllm_config(store_threshold=0) - spec_cls = OffloadingSpecFactory.get_spec_cls(config) + spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) metrics = spec_cls.build_metric_definitions( config.kv_transfer_config.kv_connector_extra_config ) @@ -291,8 +607,45 @@ def test_build_metric_definitions_returns_counter_at_threshold(): from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics config = _make_vllm_config(store_threshold=2) - spec_cls = OffloadingSpecFactory.get_spec_cls(config) + spec_cls = OffloadingSpecFactory.get_spec_cls(_get_extra_config(config)) metrics = spec_cls.build_metric_definitions( config.kv_transfer_config.kv_connector_extra_config ) assert CPUOffloadingMetrics.STORES_SKIPPED in metrics + + +def test_offloading_spec_accepts_blocks_per_chunk_for_heterogeneous_groups(): + config = _make_layout_vllm_config( + cpu_bytes_to_use=65536, + extra_config={"blocks_per_chunk": 2}, + ) + + spec = _create_spec(config, _make_hybrid_kv_cache_config()) + + assert spec.tokens_per_block == (12, 16) + assert spec.blocks_per_chunk == 2 + + +def test_block_size_and_blocks_per_chunk_are_mutually_exclusive(): + config = _make_layout_vllm_config( + cpu_bytes_to_use=65536, + extra_config={ + "block_size": 64, + "blocks_per_chunk": 2, + }, + ) + + with pytest.raises(ValueError, match="Specify only one"): + _create_spec(config, _make_kv_cache_config()) + + +def test_blocks_per_chunk_must_be_positive(): + config = _make_layout_vllm_config( + cpu_bytes_to_use=65536, + extra_config={ + "blocks_per_chunk": 0, + }, + ) + + with pytest.raises(ValueError, match="greater than 0"): + _create_spec(config, _make_kv_cache_config()) diff --git a/tests/v1/kv_offload/test_file_mapper.py b/tests/v1/kv_offload/test_file_mapper.py index 6f6e0d66196e..6c11f2d465fb 100644 --- a/tests/v1/kv_offload/test_file_mapper.py +++ b/tests/v1/kv_offload/test_file_mapper.py @@ -4,80 +4,60 @@ from unittest.mock import MagicMock -import torch - -from vllm.v1.kv_cache_interface import ( - FullAttentionSpec, - KVCacheGroupSpec, - MLAAttentionSpec, - SlidingWindowSpec, -) -from vllm.v1.kv_offload.base import ( - OffloadingSpec, - make_offload_key, +from vllm.v1.kv_offload.base import OffloadingSpec, make_offload_key +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingGroupConfig, + OffloadingModelConfig, + OffloadingParallelConfig, ) from vllm.v1.kv_offload.file_mapper import FileMapper -# --------------------------------------------------------------------------- -# Shared mocks (mirrors test_fs_tier.py pattern) -# --------------------------------------------------------------------------- - -_MOCK_VLLM_CONFIG = MagicMock() -_MOCK_VLLM_CONFIG.model_config.model = "test-model" -_MOCK_VLLM_CONFIG.cache_config.block_size = 16 -_MOCK_VLLM_CONFIG.cache_config.cache_dtype = "torch.float32" -_MOCK_VLLM_CONFIG.parallel_config.tensor_parallel_size = 1 -_MOCK_VLLM_CONFIG.parallel_config.pipeline_parallel_size = 1 -_MOCK_VLLM_CONFIG.parallel_config.prefill_context_parallel_size = 1 -_MOCK_VLLM_CONFIG.parallel_config.decode_context_parallel_size = 1 -_MOCK_VLLM_CONFIG.parallel_config.rank = 0 - -_MOCK_KV_CACHE_CONFIG = MagicMock() -_MOCK_KV_CACHE_CONFIG.kv_cache_groups = [] - -_MOCK_OFFLOADING_SPEC = MagicMock(spec=OffloadingSpec) -_MOCK_OFFLOADING_SPEC.vllm_config = _MOCK_VLLM_CONFIG -_MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG -_MOCK_OFFLOADING_SPEC.block_size_factor = 1 - - # --------------------------------------------------------------------------- # Helper # --------------------------------------------------------------------------- def make_mapper_from_offloading_spec(**kwargs) -> FileMapper: - """Helper to create FileMapper with customizable mock config.""" - # Create a copy of the mock config to avoid modifying the global one - mock_vllm_config = MagicMock() - mock_vllm_config.model_config.model = kwargs.get("model_name", "test-model") - mock_vllm_config.cache_config.block_size = kwargs.get("hash_block_size", 16) - mock_vllm_config.cache_config.cache_dtype = ( - f"torch.{kwargs.get('dtype', 'float16')}" - ) - mock_vllm_config.parallel_config.tensor_parallel_size = kwargs.get("tp_size", 1) - mock_vllm_config.parallel_config.pipeline_parallel_size = kwargs.get("pp_size", 1) - mock_vllm_config.parallel_config.prefill_context_parallel_size = kwargs.get( - "pcp_size", 1 - ) - mock_vllm_config.parallel_config.decode_context_parallel_size = kwargs.get( - "dcp_size", 1 + """Build a FileMapper from a mocked spec carrying a hand-built config.""" + config = OffloadingConfig( + groups=tuple( + OffloadingGroupConfig( + tokens_per_block=tokens_per_block, + layer_names=(layer_name,), + ) + for tokens_per_block, layer_name in kwargs.get("groups", ()) + ), + worker_kv_bytes_per_block=0, + enable_kv_cache_events=False, + extra_config={}, + engine_id="test-engine", + model=OffloadingModelConfig( + name=kwargs.get("model_name", "test-model"), + dtype=kwargs.get("dtype", "float16"), + ), + cache=OffloadingCacheConfig( + tokens_per_hash=kwargs.get("tokens_per_hash", 16), + blocks_per_chunk=kwargs.get("blocks_per_chunk", 1), + ), + parallel=OffloadingParallelConfig( + rank=kwargs.get("rank", 0), + world_size=kwargs.get("world_size", 1), + tp_size=kwargs.get("tp_size", 1), + pp_size=kwargs.get("pp_size", 1), + pcp_size=kwargs.get("pcp_size", 1), + dcp_size=kwargs.get("dcp_size", 1), + data_parallel_index=0, + is_parallelism_agnostic=kwargs.get("is_parallelism_agnostic", False), + ), ) - mock_vllm_config.parallel_config.rank = kwargs.get("rank", 0) - mock_vllm_config.use_v2_model_runner = kwargs.get("use_v2_model_runner", False) - - mock_kv_cache_config = MagicMock() - mock_kv_cache_config.kv_cache_groups = kwargs.get("kv_cache_groups", []) - - mock_offloading_spec = MagicMock(spec=OffloadingSpec) - mock_offloading_spec.vllm_config = mock_vllm_config - mock_offloading_spec.kv_cache_config = mock_kv_cache_config - mock_offloading_spec.block_size_factor = kwargs.get("block_size_factor", 1) - + spec = MagicMock(spec=OffloadingSpec) + spec.config = config return FileMapper.from_offloading_spec( root_dir=kwargs.get("root_dir", "/tmp/cache"), - offloading_spec=mock_offloading_spec, - gpu_blocks_per_file=mock_offloading_spec.block_size_factor, + offloading_spec=spec, + blocks_per_file=config.cache.blocks_per_chunk, parallel_agnostic=kwargs.get("parallel_agnostic", False), ) @@ -92,7 +72,7 @@ def test_get_file_name_full_structure(): Path must match: _r//_g/.bin Concretely: - - The segment immediately after base_path must end with `_r0` + - The segment immediately after base_path must end with `_r3` - The next segment is the first 3 hex chars of the block hash - The next segment is <2 hex chars>_g - The final segment is .bin @@ -105,7 +85,7 @@ def test_get_file_name_full_structure(): path = fm.get_file_name(key) expected_path = ( - "/tmp/cache/test-model_588656ebcc66_r3/000/10_g2/0001020304050607.bin" + "/tmp/cache/test-model_42b94bdc9933_r3/000/10_g2/0001020304050607.bin" ) assert path == expected_path @@ -114,19 +94,30 @@ def test_get_run_config_fields(): fm = make_mapper_from_offloading_spec( model_name="my-model", dtype="bfloat16", - tp_size=2, + tp_size=4, + pp_size=3, + pcp_size=2, + dcp_size=2, + groups=((64, "layer0"),), + tokens_per_hash=64, + blocks_per_chunk=3, ) cfg = fm.get_run_config() assert cfg == { "model_name": "my-model", - "hash_block_size": 16, - "gpu_blocks_per_file": 1, - "tp_size": 2, - "pp_size": 1, - "pcp_size": 1, - "dcp_size": 1, + "tokens_per_hash": 64, + "blocks_per_file": 3, + "tp_size": 4, + "pp_size": 3, + "pcp_size": 2, + "dcp_size": 2, "dtype": "bfloat16", - "kv_cache_groups": [], + "kv_cache_groups": [ + { + "tokens_per_block": 64, + "layer_names": ["layer0"], + } + ], "inference_engine": "vllm", } @@ -137,90 +128,61 @@ def test_get_config_file_path(): assert config_path == f"{fm.base_path}/config.json" -# --------------------------------------------------------------------------- -# parallel_agnostic: honored only for a single non-MLA full-attention group -# --------------------------------------------------------------------------- - - -def _full_attention_group() -> KVCacheGroupSpec: - return KVCacheGroupSpec( - layer_names=["layer0"], - kv_cache_spec=FullAttentionSpec( - block_size=16, num_kv_heads=4, head_size=128, dtype=torch.float32 - ), +def test_hybrid_file_identity_uses_resolved_tokens_per_hash(): + # For heterogeneous groups the namespace records the resolved hash + # granularity (GCD of the group block sizes), which is the actual + # granularity of the offload block hashes. + fm = make_mapper_from_offloading_spec( + groups=((12, "full_layer"), (16, "mla_layer")), + tokens_per_hash=4, ) + assert fm.fields["tokens_per_hash"] == 4 + assert fm.fields["kv_cache_groups"] == [ + {"tokens_per_block": 12, "layer_names": ["full_layer"]}, + {"tokens_per_block": 16, "layer_names": ["mla_layer"]}, + ] -def _sliding_window_group() -> KVCacheGroupSpec: - return KVCacheGroupSpec( - layer_names=["layer0"], - kv_cache_spec=SlidingWindowSpec( - block_size=16, - num_kv_heads=4, - head_size=128, - dtype=torch.float32, - sliding_window=128, - ), - ) +# --------------------------------------------------------------------------- +# parallel_agnostic: opt-in honored only when the config marks the layout +# parallelism-agnostic (predicate computation is covered in test_factory.py) +# --------------------------------------------------------------------------- -def test_parallel_agnostic_enabled_for_single_full_attention(): - # tp/rank are collapsed out of the namespace so the cache is shared - # across tensor-parallel sizes. +def test_parallel_agnostic_collapses_namespace_when_config_allows(): fm = make_mapper_from_offloading_spec( - tp_size=2, + tp_size=4, + pp_size=3, + pcp_size=2, + dcp_size=2, rank=1, - kv_cache_groups=[_full_attention_group()], + is_parallelism_agnostic=True, parallel_agnostic=True, ) assert fm.fields["tp_size"] == 1 + assert fm.fields["pp_size"] == 1 + assert fm.fields["pcp_size"] == 1 + assert fm.fields["dcp_size"] == 1 assert fm.rank == 0 -def test_parallel_agnostic_disabled_for_multiple_groups(): - # More than one KV-cache group (hybrid model) => keep per-layout namespacing. +def test_parallel_agnostic_ignored_when_config_disallows(): fm = make_mapper_from_offloading_spec( tp_size=2, - kv_cache_groups=[_full_attention_group(), _full_attention_group()], - parallel_agnostic=True, - ) - assert fm.fields["tp_size"] == 2 - - -def test_parallel_agnostic_disabled_for_non_full_attention(): - # Single group but not full attention (sliding window) => keep namespacing. - fm = make_mapper_from_offloading_spec( - tp_size=2, - kv_cache_groups=[_sliding_window_group()], + rank=1, + is_parallelism_agnostic=False, parallel_agnostic=True, ) assert fm.fields["tp_size"] == 2 - - -def test_parallel_agnostic_excludes_mla(): - # MLA latent KV is replicated per rank, so its offloaded blocks are not - # parallelism-invariant: the opt-in must not collapse tp/rank. - group = KVCacheGroupSpec( - layer_names=["layer0"], - kv_cache_spec=MLAAttentionSpec( - block_size=16, num_kv_heads=1, head_size=576, dtype=torch.float32 - ), - ) - fm = make_mapper_from_offloading_spec( - tp_size=2, rank=1, kv_cache_groups=[group], parallel_agnostic=True - ) - assert fm.fields["tp_size"] == 2 assert fm.rank == 1 -def test_parallel_agnostic_disabled_on_v2_model_runner(): - # V2's KV layout is not known to be parallelism-invariant: don't collapse. +def test_namespace_kept_without_parallel_agnostic_opt_in(): fm = make_mapper_from_offloading_spec( tp_size=2, rank=1, - kv_cache_groups=[_full_attention_group()], - use_v2_model_runner=True, - parallel_agnostic=True, + is_parallelism_agnostic=True, + parallel_agnostic=False, ) assert fm.fields["tp_size"] == 2 assert fm.rank == 1 diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py index 7eb812435247..d3bacc8f326f 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -47,7 +47,7 @@ def test_config_fingerprint_empty_when_no_fields(self): def test_config_fingerprint_deterministic(self): """Same config fields → same fingerprint.""" view = self._make_view() - fields = {"model": "llama", "dtype": "float16", "block_size_factor": 1} + fields = {"model": "llama", "dtype": "float16", "blocks_per_chunk": 1} with patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", None): t1 = NixlTransport("test:1", view, config_fields=fields) t2 = NixlTransport("test:2", view, config_fields=fields) diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index 52e2c8363232..01fd295302a8 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -1419,9 +1419,9 @@ def _construct(monkeypatch, dp_index=0, **kwargs) -> P2PSecondaryTierManager: or SimpleNamespace(), ) spec = SimpleNamespace( - block_size_factor=1, - vllm_config=SimpleNamespace( - parallel_config=SimpleNamespace(data_parallel_index=dp_index) + blocks_per_chunk=1, + config=SimpleNamespace( + parallel=SimpleNamespace(data_parallel_index=dp_index) ), ) mgr = P2PSecondaryTierManager(spec, memoryview(b""), **kwargs) diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 4ac734d957f4..dcc92a4fa8bf 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -20,14 +20,23 @@ from vllm.distributed.kv_events import MEDIUM_FS from vllm.v1.kv_offload.base import ( + Locality, LookupResult, OffloadingEvent, + OffloadingKVEventsConfig, OffloadKey, ReqContext, ScheduleEndContext, make_offload_key, ) +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingModelConfig, + OffloadingParallelConfig, +) from vllm.v1.kv_offload.tiering.base import JobMetadata +from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.fs.manager import ( FileSystemTierManager, ) @@ -41,35 +50,40 @@ _DTYPE: torch.dtype = torch.float32 _CTX = ReqContext(req_id="test") -_MOCK_VLLM_CONFIG = MagicMock() -_MOCK_VLLM_CONFIG.model_config.model = "test-model" -_MOCK_VLLM_CONFIG.cache_config.block_size = 16 -_MOCK_VLLM_CONFIG.cache_config.cache_dtype = "torch.float32" -_MOCK_VLLM_CONFIG.parallel_config.tensor_parallel_size = 1 -_MOCK_VLLM_CONFIG.parallel_config.pipeline_parallel_size = 1 -_MOCK_VLLM_CONFIG.parallel_config.prefill_context_parallel_size = 1 -_MOCK_VLLM_CONFIG.parallel_config.decode_context_parallel_size = 1 -_MOCK_VLLM_CONFIG.parallel_config.rank = 0 - -_MOCK_KV_CACHE_CONFIG = MagicMock() -_MOCK_KV_CACHE_CONFIG.kv_cache_groups = [] - -_MOCK_OFFLOADING_SPEC = MagicMock() -_MOCK_OFFLOADING_SPEC.vllm_config = _MOCK_VLLM_CONFIG -_MOCK_OFFLOADING_SPEC.kv_cache_config = _MOCK_KV_CACHE_CONFIG -_MOCK_OFFLOADING_SPEC.block_size_factor = 1 - def _make_offloading_spec(enable_kv_cache_events: bool) -> MagicMock: """Mock spec with an explicit global KV events flag.""" spec = MagicMock() - spec.vllm_config = _MOCK_VLLM_CONFIG - spec.kv_cache_config = _MOCK_KV_CACHE_CONFIG - spec.block_size_factor = 1 - spec.kv_events_config.enable_kv_cache_events = enable_kv_cache_events + spec.config = OffloadingConfig( + groups=(), + worker_kv_bytes_per_block=0, + enable_kv_cache_events=enable_kv_cache_events, + extra_config={}, + engine_id="test-engine", + model=OffloadingModelConfig(name="test-model", dtype="float32"), + cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), + parallel=OffloadingParallelConfig( + rank=0, + world_size=1, + tp_size=1, + pp_size=1, + pcp_size=1, + dcp_size=1, + data_parallel_index=0, + is_parallelism_agnostic=False, + ), + ) + spec.blocks_per_chunk = 1 + spec.kv_events_config = OffloadingKVEventsConfig( + enable_kv_cache_events=enable_kv_cache_events, + self_describing_kv_events=False, + ) return spec +_MOCK_OFFLOADING_SPEC = _make_offloading_spec(enable_kv_cache_events=False) + + def key(n: int) -> OffloadKey: return make_offload_key(n.to_bytes(8, "big"), 0) @@ -174,6 +188,7 @@ def fs_tier_with_events(tmp_path): n_read_threads=4, n_write_threads=4, enable_kv_events=True, + locality="LOCAL", ) yield tier tier.shutdown() @@ -239,6 +254,40 @@ def test_invalid_path_raises_at_construction(): ) +@pytest.mark.parametrize("locality", ["local", ""]) +def test_invalid_locality_raises_at_construction(tmp_path, locality): + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + + with pytest.raises(ValueError, match="Locality"): + FileSystemTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=memoryview(tensor.numpy()), + tier_type="fs", + root_dir=str(tmp_path), + locality=locality, + ) + + +def test_factory_forwards_locality_to_fs_tier(tmp_path): + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + tier = SecondaryTierFactory.create_secondary_tier( + { + "type": "fs", + "root_dir": str(tmp_path), + "n_read_threads": 1, + "n_write_threads": 1, + "locality": "LOCAL", + }, + memoryview(tensor.numpy()), + _MOCK_OFFLOADING_SPEC, + ) + try: + assert isinstance(tier, FileSystemTierManager) + assert tier.locality is Locality.LOCAL + finally: + tier.shutdown() + + def test_failed_load_missing_file(fs_tier): """Test that loading a block whose file does not exist results in a failed job.""" tier, _ = fs_tier @@ -430,11 +479,38 @@ def test_successful_store_emits_stored_event(fs_tier_with_events): assert events[0].keys == keys # Literal medium pins the wire contract, not just the constant choice. assert events[0].medium == "FS" + assert events[0].locality is Locality.LOCAL assert not events[0].removed # take_events drains the buffer. assert list(tier.take_events()) == [] +@pytest.mark.parametrize( + ("locality", "expected"), + [(None, None), ("REMOTE", Locality.REMOTE)], +) +def test_store_event_uses_configured_locality(tmp_path, locality, expected): + tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) + locality_config = {} if locality is None else {"locality": locality} + tier = FileSystemTierManager( + offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), + primary_kv_view=memoryview(tensor.numpy()), + tier_type="fs", + root_dir=str(tmp_path), + enable_kv_events=True, + **locality_config, + ) + try: + tier.submit_store(make_job(1, [key(1)], [0])) + assert all(r.success for r in drain(tier)) + + events = list(tier.take_events()) + assert len(events) == 1 + assert events[0].locality is expected + finally: + tier.shutdown() + + def test_load_job_emits_no_event(fs_tier_with_events): tier = fs_tier_with_events tier.submit_store(make_job(1, [key(1)], [0])) diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 37687adce0cb..7500df8b4b94 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -15,15 +15,24 @@ from unittest.mock import MagicMock, patch import numpy as np +import pytest import torch from vllm.v1.kv_offload.base import ( + Locality, LookupResult, + OffloadingKVEventsConfig, OffloadKey, ReqContext, ScheduleEndContext, make_offload_key, ) +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingModelConfig, + OffloadingParallelConfig, +) from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager @@ -33,24 +42,30 @@ # --------------------------------------------------------------------------- -def _make_vllm_config(): - return SimpleNamespace( - model_config=SimpleNamespace(model="test/model"), - cache_config=SimpleNamespace(block_size=16, cache_dtype="float16"), - parallel_config=SimpleNamespace( - tensor_parallel_size=1, - pipeline_parallel_size=1, - prefill_context_parallel_size=1, - decode_context_parallel_size=1, +def _make_offloading_config(enable_kv_cache_events: bool) -> OffloadingConfig: + return OffloadingConfig( + groups=(), + worker_kv_bytes_per_block=0, + enable_kv_cache_events=enable_kv_cache_events, + extra_config={}, + engine_id="test-engine", + model=OffloadingModelConfig(name="test/model", dtype="float16"), + cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), + parallel=OffloadingParallelConfig( rank=0, + world_size=1, + tp_size=1, + pp_size=1, + pcp_size=1, + dcp_size=1, + data_parallel_index=0, + is_parallelism_agnostic=False, ), - use_v2_model_runner=False, ) _OFFLOADING_SPEC = SimpleNamespace( - vllm_config=_make_vllm_config(), - kv_cache_config=SimpleNamespace(kv_cache_groups=[]), + config=_make_offloading_config(enable_kv_cache_events=False), ) _STORE_CONFIG = { @@ -182,9 +197,11 @@ def _query_memory(self, queries, mem_type, agent_name): def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace: """Offloading spec stub with an explicit global KV events flag.""" return SimpleNamespace( - vllm_config=_make_vllm_config(), - kv_cache_config=SimpleNamespace(kv_cache_groups=[]), - kv_events_config=SimpleNamespace(enable_kv_cache_events=enable_kv_cache_events), + config=_make_offloading_config(enable_kv_cache_events), + kv_events_config=OffloadingKVEventsConfig( + enable_kv_cache_events=enable_kv_cache_events, + self_describing_kv_events=False, + ), ) @@ -250,6 +267,12 @@ def lookup_and_wait( # --------------------------------------------------------------------------- +@pytest.mark.parametrize("locality", ["local", ""]) +def test_invalid_locality_raises_at_construction(locality): + with pytest.raises(ValueError, match="Locality"): + _make_tier(locality=locality) + + class TestMockObjTierBasic: def setup_method(self): self.tier, self.agent = _make_tier(num_blocks=4) @@ -438,6 +461,7 @@ def setup_method(self): self.tier, self.agent = _make_tier( offloading_spec=_make_events_spec(enable_kv_cache_events=True), enable_kv_events=True, + locality="REMOTE", ) def test_successful_store_emits_stored_event(self): @@ -451,10 +475,32 @@ def test_successful_store_emits_stored_event(self): assert events[0].keys == keys # Literal medium pins the wire contract, not just the constant choice. assert events[0].medium == "OBJ" + assert events[0].locality is Locality.REMOTE assert not events[0].removed # take_events drains the buffer. assert list(self.tier.take_events()) == [] + @pytest.mark.parametrize( + ("locality", "expected"), + [(None, None), ("LOCAL", Locality.LOCAL)], + ) + def test_store_event_uses_configured_locality(self, locality, expected): + locality_config = {} if locality is None else {"locality": locality} + tier, _ = _make_tier( + offloading_spec=_make_events_spec(enable_kv_cache_events=True), + enable_kv_events=True, + **locality_config, + ) + try: + tier.submit_store(make_job(1, [key(1)], [0])) + assert all(r.success for r in drain(tier)) + + events = list(tier.take_events()) + assert len(events) == 1 + assert events[0].locality is expected + finally: + tier.shutdown() + def test_mixed_job_results_emit_event_only_for_successful_job(self): """With a failed and a successful store job resolving in the same poll, exactly one event is emitted and its keys belong to the diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 13460f68d7cc..cf546b304cd8 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -35,6 +35,7 @@ JobMetadata, JobResult, SecondaryTierManager, + TieringOffloadingMetrics, ) from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory @@ -233,9 +234,9 @@ def manager_setup(self): secondary_tiers=[self.secondary_tier1, self.secondary_tier2], ) - def _simulate_on_schedule_end(self): + def _simulate_on_schedule_end(self, new_req_ids: list[str] | None = None): """Simulate end of scheduler step: lifecycle flush + drain events.""" - ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) + ctx = ScheduleEndContext(new_req_ids=new_req_ids or [], preempted_req_ids=()) self.manager.on_schedule_end(ctx) list(self.manager.take_events()) @@ -395,6 +396,87 @@ def test_promotion_from_secondary(self, manager_setup): # Next lookup should succeed assert count_hits(self.manager, blocks) == 3 + def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup): + """Resolved lookups report one sync delay sample on allocation.""" + self._start_request() + blocks = to_keys(range(2)) + + # No tier has these blocks: they resolve immediately as misses. + for block in blocks: + assert self.manager.lookup(block, _CTX) is LookupResult.MISS + + stats = self.manager.get_stats() + if stats is not None: + assert f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count" not in ( + stats.reduce() + ) + + self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id]) + + stats = self.manager.get_stats() + assert stats is not None + reduced = stats.reduce() + assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1 + assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in reduced + + def test_lookup_reports_async_delay_across_promotion(self, manager_setup): + """A new request reports async delay at schedule end.""" + self._start_request() + block = to_keys(range(1))[0] + self.secondary_tier1.blocks[block] = True + + # First lookup finds the block in a secondary tier and defers. + assert self.manager.lookup(block, _CTX) is LookupResult.RETRY + + # The first scheduler step reports async delay for the new request. + self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id]) + stats = self.manager.get_stats() + assert stats is not None + reduced = stats.reduce() + assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1 + assert reduced[f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count"] == 1 + + # Promotion completes on the next scheduler step. + self._simulate_on_schedule_end() + + # Next lookup resolves via the now-promoted primary-tier block. + assert self.manager.lookup(block, _CTX) is LookupResult.HIT + + stats = self.manager.get_stats() + if stats is not None: + assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in ( + stats.reduce() + ) + + def test_lookup_reports_async_delay_on_request_finish(self, manager_setup): + """Never-allocated lookup delays flush at teardown.""" + ctx = ReqContext(req_id="req_lookup_finish") + self._start_request(ctx) + block = to_keys(range(1))[0] + self.secondary_tier1.blocks[block] = True + + # Lookup finds the block in a secondary tier and defers. + assert self.manager.lookup(block, ctx) is LookupResult.RETRY + + self._simulate_on_schedule_end() + stats = self.manager.get_stats() + if stats is not None: + assert f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count" not in ( + stats.reduce() + ) + assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in ( + stats.reduce() + ) + + # Request finishes before the deferred lookup is ever resolved. + self.manager.on_request_finished(ctx) + + stats = self.manager.get_stats() + assert stats is not None + reduced = stats.reduce() + assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1 + assert reduced[f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count"] == 1 + def test_partial_lookup(self, manager_setup): """Test lookup with partial hits.""" blocks = to_keys(range(5)) diff --git a/tests/v1/metrics/test_stats.py b/tests/v1/metrics/test_stats.py index 21f496ea4aea..0de74f0faaa4 100644 --- a/tests/v1/metrics/test_stats.py +++ b/tests/v1/metrics/test_stats.py @@ -1,12 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm.v1.engine import FinishReason +from vllm.v1.core.sched.output import ScheduledEncoderInputStats, SchedulerOutput +from vllm.v1.engine import EngineCoreOutputs, FinishReason from vllm.v1.metrics.stats import ( IterationStats, PrefillStats, PromptTokenStats, RequestStateStats, + SchedulerIterationDetails, + SchedulerStats, ) +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder +from vllm.v1.utils import compute_iteration_details def test_iteration_stats_repr(): @@ -14,6 +19,45 @@ def test_iteration_stats_repr(): assert repr(iteration_stats).startswith("IterationStats(") +def test_scheduler_iteration_details_serialization(): + iteration_details = SchedulerIterationDetails( + iteration_index=1, + num_ctx_requests=2, + num_ctx_tokens=3, + num_generation_requests=4, + num_generation_tokens=5, + elapsed_ms=6.7, + num_encoder_inputs=2, + num_encoder_output_tokens=392, + ) + outputs = EngineCoreOutputs( + scheduler_stats=SchedulerStats( + kv_cache_usage=0.5, + iteration_details=iteration_details, + ) + ) + + encoded = MsgpackEncoder().encode(outputs) + decoded = MsgpackDecoder(EngineCoreOutputs).decode(encoded) + + assert decoded.scheduler_stats is not None + assert decoded.scheduler_stats.kv_cache_usage == 0.5 + assert decoded.scheduler_stats.iteration_details == iteration_details + + +def test_compute_iteration_details_includes_encoder_stats(): + scheduler_output = SchedulerOutput.make_empty() + scheduler_output.scheduled_encoder_input_stats = ScheduledEncoderInputStats( + num_inputs=2, + output_tokens=392, + ) + + iteration_details = compute_iteration_details(scheduler_output) + + assert iteration_details.num_encoder_inputs == 2 + assert iteration_details.num_encoder_output_tokens == 392 + + def test_prefill_kv_computed_with_cache(): """Test that prefill KV compute correctly excludes cached tokens.""" iteration_stats = IterationStats() diff --git a/tests/v1/worker/test_encoder_runner.py b/tests/v1/worker/test_encoder_runner.py index 70c0426640d2..2cacc20afb61 100644 --- a/tests/v1/worker/test_encoder_runner.py +++ b/tests/v1/worker/test_encoder_runner.py @@ -78,6 +78,7 @@ def test_draft_lookahead_uses_boundary_feature_when_cached(): # f0 covers positions 0..6 (+1 skew); f1's first embed covers position 7. assert len(mm_embeds) == 2 + assert [e.modality for e in mm_embeds] == ["image", "image"] assert bool(is_mm_embed[7]) assert int(is_mm_embed.sum()) == 8 @@ -94,6 +95,7 @@ def test_draft_lookahead_tolerates_missing_boundary_feature(): # Only f0 is gathered; f1's boundary position falls back silently. assert len(mm_embeds) == 1 + assert [e.modality for e in mm_embeds] == ["image"] assert not bool(is_mm_embed[7]) assert int(is_mm_embed.sum()) == 7 @@ -154,4 +156,28 @@ def test_multi_request_batch_gathers_per_request(draft_lookahead): # Both requests contribute a feature; with the +1 skew each marks 7 of its # 8 positions (the skew drops one), otherwise all 8. assert len(mm_embeds) == 2 + assert [e.modality for e in mm_embeds] == ["image", "image"] assert int(is_mm_embed.sum()) == (14 if draft_lookahead else 16) + + +def test_gather_preserves_mixed_modalities(): + """Modalities must be attached on tensors in gather order.""" + video = MultiModalFeatureSpec( + data=None, + modality="video", + identifier="v0", + mm_position=PlaceholderRange(offset=0, length=4), + ) + audio = MultiModalFeatureSpec( + data=None, + modality="audio", + identifier="a0", + mm_position=PlaceholderRange(offset=4, length=4), + ) + runner = _make_runner([video, audio], cached=[video, audio]) + + mm_embeds, is_mm_embed = _gather(runner, num_scheduled=8, draft_lookahead=0) + + assert len(mm_embeds) == 2 + assert [e.modality for e in mm_embeds] == ["video", "audio"] + assert int(is_mm_embed.sum()) == 8 diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 67a527405e2d..ff43ac4d6871 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -1,85 +1,52 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Build DeepGEMM's `_C` pybind11 extension for . +"""Build DeepGEMM's TORCH_LIBRARY extension and copy vendored artifacts. -Driven from cmake/external_projects/deepgemm.cmake. The driver runs against -the build interpreter's torch; is only consulted for INCLUDEPY -and SOABI, so target venvs don't need torch installed. +DeepGEMM now registers ops via TORCH_LIBRARY into ``deep_gemm._C_extension`` +(abi3) and exposes the legacy API through ``deep_gemm/_C.py``. This driver +delegates to DeepGEMM's ``setup.py build_ext --inplace`` and copies the shim +plus extension into the cmake output directory. -Usage: python build_deepgemm_C.py +Usage: python build_deepgemm_C.py """ -import json import os +import shutil import subprocess import sys from pathlib import Path -import torch -from torch.utils import cpp_extension - -if len(sys.argv) != 4: - sys.exit(f"usage: {sys.argv[0]} ") +if len(sys.argv) != 3: + sys.exit(f"usage: {sys.argv[0]} ") src = Path(sys.argv[1]).resolve() out = Path(sys.argv[2]).resolve() -target_py = sys.argv[3] +_pkg = src / "deep_gemm" out.mkdir(parents=True, exist_ok=True) -info = json.loads( - subprocess.check_output( - [ - target_py, - "-c", - "import sysconfig, json; " - "print(json.dumps({k: sysconfig.get_config_var(k) " - "for k in ('EXT_SUFFIX', 'INCLUDEPY')}))", - ] - ).decode() +if not (_pkg / "_C.py").is_file(): + sys.exit( + f"DeepGEMM source at {src} is missing deep_gemm/_C.py; " + "expected TORCH_LIBRARY migration layout" + ) + +# Avoid DeepGEMM's clean-git assertion when vendoring a local dirty tree. +env = os.environ.copy() +env.pop("DG_SKIP_CUDA_BUILD", None) + +print(f"[build_deepgemm_C] building in {src} with {sys.executable}", flush=True) +subprocess.check_call( + [sys.executable, "setup.py", "build_ext", "--inplace"], + cwd=src, + env=env, ) -cuda_home = cpp_extension.CUDA_HOME -if cuda_home is None: - sys.exit("CUDA_HOME not found; cannot build DeepGEMM _C") -# CCCL lives outside the standard CUDAToolkit search (mirrors DeepGEMM's setup.py). -includes = [ - info["INCLUDEPY"], - f"{cuda_home}/include", - f"{cuda_home}/include/cccl", - str(src / "csrc"), - str(src / "deep_gemm/include"), - str(src / "third-party/cutlass/include"), - str(src / "third-party/cutlass/tools/util/include"), - str(src / "third-party/fmt/include"), - *cpp_extension.include_paths(device_type="cuda"), -] +shim = _pkg / "_C.py" +shutil.copy2(shim, out / shim.name) -cmd = [ - os.environ.get("CXX", "g++"), - "-shared", - "-fPIC", - "-std=c++20", - "-O3", - "-g0", - "-Wno-psabi", - "-Wno-deprecated-declarations", - "-DTORCH_API_INCLUDE_EXTENSION_H", - "-DTORCH_EXTENSION_NAME=_C", - f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}", - *(f"-I{p}" for p in includes), - str(src / "csrc/python_api.cpp"), - *(f"-L{p}" for p in cpp_extension.library_paths(device_type="cuda")), - f"-L{cuda_home}/lib64", - "-ltorch", - "-ltorch_python", - "-ltorch_cpu", - "-ltorch_cuda", - "-lc10", - "-lc10_cuda", - "-lcudart", - "-lnvrtc", - "-o", - str(out / f"_C{info['EXT_SUFFIX']}"), -] -print("[build_deepgemm_C] " + " ".join(cmd), flush=True) -subprocess.check_call(cmd) +so_files = sorted(_pkg.glob("_C_extension*.so")) +if not so_files: + sys.exit(f"DeepGEMM build did not produce deep_gemm/_C_extension*.so under {src}") +for so in so_files: + shutil.copy2(so, out / so.name) + print(f"[build_deepgemm_C] installed {so.name} -> {out}", flush=True) diff --git a/tools/check_wheel_deepgemm.py b/tools/check_wheel_deepgemm.py index 6f8a03ffd3dc..89788a617c1c 100644 --- a/tools/check_wheel_deepgemm.py +++ b/tools/check_wheel_deepgemm.py @@ -1,41 +1,45 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Assert the installed vLLM has a `_C.cpython-X.Y-*.so` for every CPython -covered by `requires-python`. Fails closed if a Python's `.so` is missing -from the wheel — i.e. the regression that surfaced in #41476/#41512. +"""Assert the vendored DeepGEMM package has the TORCH_LIBRARY binding layout. -Run from a CI test job after vLLM is installed, e.g. the H100 deepgemm -kernel tests in .buildkite/test_areas/kernels.yaml. +Expects ``deep_gemm/_C.py`` plus a single ``_C_extension*.so`` (abi3) under +``vllm.third_party.deep_gemm``. Run after vLLM is installed, e.g. the H100 +deepgemm kernel tests in .buildkite/test_areas/kernels.yaml. """ import importlib.util -import os import sys from pathlib import Path -import regex as re -import tomllib -SO_RE = re.compile(r"^_C\.cpython-(\d)(\d+)-") - - -def required_pythons() -> list[str]: - pyproject = Path(__file__).resolve().parent.parent / "pyproject.toml" - spec = tomllib.loads(pyproject.read_text())["project"]["requires-python"] - m = re.match(r">=3\.(\d+),<3\.(\d+)", spec) - if not m: - sys.exit(f"unexpected requires-python format: {spec!r}") - return [f"3.{v}" for v in range(int(m[1]), int(m[2]))] - - -spec = importlib.util.find_spec("vllm.third_party.deep_gemm") -if spec is None or spec.origin is None: - sys.exit("vllm.third_party.deep_gemm not importable; is vllm installed?") -pkg_dir = Path(spec.origin).parent - -found = {f"{m[1]}.{m[2]}" for f in os.listdir(pkg_dir) if (m := SO_RE.match(f))} -required = required_pythons() -missing = [v for v in required if v not in found] -print(f"deepgemm _C: found {sorted(found)}, required {required}, missing {missing}") -sys.exit(1 if missing else 0) +def main() -> int: + spec = importlib.util.find_spec("vllm.third_party.deep_gemm") + if spec is None or spec.origin is None: + print( + "vllm.third_party.deep_gemm not importable; is vllm installed?", + file=sys.stderr, + ) + return 1 + pkg_dir = Path(spec.origin).parent + + shim = pkg_dir / "_C.py" + so_files = sorted(pkg_dir.glob("_C_extension*.so")) + missing = [] + if not shim.is_file(): + missing.append("_C.py") + if not so_files: + missing.append("_C_extension*.so") + + print( + f"deepgemm vendored binding: shim={shim.is_file()}, " + f"extensions={[p.name for p in so_files]}" + ) + if missing: + print(f"missing: {missing}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index 52f73e9050f7..81c4b4d8985d 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -6,9 +6,8 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake -DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" -# NOTE: This is currently targeting nv-dev branch due to sm120 support -DEEPGEMM_GIT_REF="a6b593d2826719dcf4892609af7b84ee23aaf32a" +DEEPGEMM_GIT_REPO="https://github.com/cleonard530/DeepGEMM.git" +DEEPGEMM_GIT_REF="441c417c6cf7184593421273b7e6d79a0999a8f3" WHEEL_DIR="" # Parse command line arguments diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 365b2f5bb771..a2fc173f0357 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -39,6 +39,7 @@ class ForbiddenImport: "vllm/distributed/device_communicators/shm_broadcast.py", "vllm/distributed/device_communicators/shm_object_storage.py", "vllm/distributed/weight_transfer/ipc_engine.py", + "vllm/distributed/weight_transfer/clients.py", "tests/distributed/test_weight_transfer.py", "vllm/utils/hashing.py", "tests/multimodal/media/test_base.py", diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index 4852a8cdf0f3..c52e0e7ce630 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -128,7 +128,6 @@ r"vllm/model_executor/models/[vV]", r"vllm/model_executor/models/[wW]", r"vllm/model_executor/models/[zZ]", - "vllm/model_executor/layers/fla/ops", ] diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 3c018313d06f..8ab06d446bf4 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2458,6 +2458,128 @@ def triton_fp4_gemm_dynamic_quant( gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y) return y + @staticmethod + def fused_qk_norm_rope_and_cache( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + is_neox: bool, + rms_norm_eps: float, + q_out: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + slot_mapping: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + k_out: torch.Tensor | None, + v_out: torch.Tensor | None, + return_kv: bool, + use_shuffle_layout: bool, + block_size: int, + x: int, + rotary_dim: int = 0, + ): + from aiter.ops.fused_qk_norm_rope_cache_quant import ( + fused_qk_norm_rope_cache_pts_quant_shuffle, + ) + + fused_qk_norm_rope_cache_pts_quant_shuffle( + qkv, + q_weight, + k_weight, + cos_sin_cache, + positions, + qkv.size(0), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + is_neox, + rms_norm_eps, + q_out, + k_cache, + v_cache, + slot_mapping, + k_scale, + v_scale, + k_out, + v_out, + return_kv, + use_shuffle_layout, + block_size, + x, + rotary_dim, + ) + + @staticmethod + def do_qk_norm_rope_kvcache_update( + qkv: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + positions: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + head_dim: int, + is_neox: bool, + rms_norm_eps: float, + q_out: torch.Tensor, + k_out: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + slot_mapping: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + kv_cache_dtype: str, + use_shuffle_layout: bool, + ) -> None: + """Run the fused QK-norm+RoPE+KV-cache op on already-split k/v caches. + + Shared by the AITER FA and unified-attention impls. The caller splits + kv_cache, since the unbind dim depends on the layout (e.g. the unified + encoder-decoder path is K/V-first), and passes use_shuffle_layout + (unified reads NHD and must pass False). + """ + if kv_cache_dtype.startswith("fp8"): + key_cache = key_cache.view(current_platform.fp8_dtype()) + value_cache = value_cache.view(current_platform.fp8_dtype()) + # Partial-rotary support (e.g. GLM-4.7 applies rotary to only a prefix + # of each head's channel dim). + rotary_dim = cos_sin_cache.shape[-1] + kernel_rotary_dim = 0 if rotary_dim == head_dim else rotary_dim + rocm_aiter_ops.fused_qk_norm_rope_and_cache( + qkv=qkv, + q_weight=q_weight, + k_weight=k_weight, + cos_sin_cache=cos_sin_cache, + positions=positions, + num_heads_q=num_heads_q, + num_heads_k=num_heads_k, + num_heads_v=num_heads_k, + head_dim=head_dim, + is_neox=is_neox, + rms_norm_eps=rms_norm_eps, + q_out=q_out, + k_cache=key_cache, + v_cache=value_cache, + slot_mapping=slot_mapping, + k_scale=k_scale, + v_scale=v_scale, + k_out=k_out, + v_out=None, + return_kv=True, + use_shuffle_layout=use_shuffle_layout, + block_size=key_cache.shape[1], + x=16 // key_cache.element_size(), + rotary_dim=kernel_rotary_dim, + ) + @staticmethod def triton_rope_and_cache( query: torch.Tensor, diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 9fcc4ae87e8e..599cac0ed6f2 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2759,9 +2759,9 @@ def cp_gather_and_upconvert_fp8_kv_cache( src_cache: torch.Tensor, dst: torch.Tensor, block_table: torch.Tensor, - seq_lens: torch.Tensor, workspace_starts: torch.Tensor, batch_size: int, + seq_starts: torch.Tensor | None = None, ) -> None: """Gather and upconvert FP8 KV cache to BF16 workspace. @@ -2769,12 +2769,12 @@ def cp_gather_and_upconvert_fp8_kv_cache( src_cache: FP8 KV cache [num_blocks, block_size, 656] dst: BF16 output workspace [total_tokens, 576] block_table: Block indices [num_reqs, max_blocks] - seq_lens: Sequence lengths [num_reqs] workspace_starts: Workspace start offsets [num_reqs] batch_size: Number of requests + seq_starts: Optional source sequence offsets [num_reqs] """ torch.ops._C_cache_ops.cp_gather_and_upconvert_fp8_kv_cache( - src_cache, dst, block_table, seq_lens, workspace_starts, batch_size + src_cache, dst, block_table, workspace_starts, batch_size, seq_starts ) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index f80498c8135c..8fbf5b41f747 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -489,6 +489,11 @@ def _decompose_size_nodes(graph: fx.GraphModule) -> None: size_nodes = list(graph.graph.find_nodes(op="call_method", target="size")) for node in size_nodes: + # Only x.size() (no dim) returns a torch.Size tuple that can't cross + # split boundaries. x.size(dim) already returns a scalar SymInt/int, + # which crosses fine, so leave it untouched. + if len(node.args) > 1 or "dim" in node.kwargs: + continue tensor_node = node.args[0] ev = tensor_node.meta.get("example_value") assert ev is not None, ( diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 99b2892a770e..af2719ce3b2d 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -192,7 +192,7 @@ def forward_custom( z: torch.Tensor, weight: torch.Tensor, ) -> torch.Tensor: - from vllm.model_executor.layers.fla.ops.layernorm_guard import ( + from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( rmsnorm_fn, ) diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py index e322e109c443..2c8b9c43bdee 100644 --- a/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py +++ b/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py @@ -25,6 +25,11 @@ FUSED_QK_ROPE_OP = torch.ops._C.fused_qk_norm_rope.default +# Head dimensions supported by csrc/fused_qknorm_rope_kernel.cu's +# launchFusedQKNormRope and launchFusedQKNormRopeNTokenHeads dispatchers. +# Keep in sync with the switch statements in that file. +SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS: tuple[int, ...] = (64, 128, 256) + P = ParamSpec("P") @@ -190,7 +195,12 @@ def replacement( class QKNormRoPEFusionPass(VllmPatternMatcherPass): - """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists.""" + """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists. + + Registers patterns for both standard vLLM ops and ROCm AITER ops + (when AITER is enabled), so the fusion fires regardless of which + RMSNorm/RoPE implementation the graph uses. + """ @enable_fake_mode def __init__(self, config: VllmConfig) -> None: @@ -216,6 +226,17 @@ def __init__(self, config: VllmConfig) -> None: ) return + for layer in attn_layers.values(): + if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS: + logger.warning_once( + "QK Norm+RoPE fusion not enabled: layer head_size=%d is not " + "supported by fused_qk_norm_rope kernel (supported: %s). " + "Falling back to unfused QK norm + RoPE path.", + layer.head_size, + SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS, + ) + return + self._attention_geometries = tuple( sorted( { diff --git a/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py new file mode 100644 index 000000000000..85d82c191b27 --- /dev/null +++ b/vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py @@ -0,0 +1,492 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import inspect +from collections.abc import Callable +from typing import ParamSpec + +import torch +import torch._inductor.pattern_matcher as pm +from torch import fx +from torch._higher_order_ops.auto_functionalize import auto_functionalized +from torch._inductor.pattern_matcher import PatternMatcherPass + +import vllm.ir.ops +from vllm.config import VllmConfig, get_layers_from_vllm_config +from vllm.config.utils import Range +from vllm.logger import init_logger +from vllm.model_executor.layers.attention.attention import ( + Attention, + get_attention_context, +) +from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op + +from ..inductor_pass import enable_fake_mode +from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass +from .matcher_utils import MatcherRotaryEmbedding +from .rms_quant_fusion import empty_bf16, empty_fp32, empty_i64 + +logger = init_logger(__name__) + +P = ParamSpec("P") + +# Head sizes the fused kernel fused_qk_norm_rope_cache_pts_quant_shuffle() supports +# Other sizes hard-abort, so skip those layers. +SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS: tuple[int, ...] = (64, 128, 256) + + +# --------------------------------------------------------------------------- +# Custom op: fused QK-norm + RoPE + KV cache update +# --------------------------------------------------------------------------- + + +def fused_qk_norm_rope_and_unified_kv_cache_update_impl( + q_out: torch.Tensor, + k_out: torch.Tensor, + qkv: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + is_neox: bool, + layer_name: str = "", +) -> torch.Tensor: + _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name) + if layer_slot_mapping is not None: + attn_layer.impl.do_qk_norm_rope_kvcache_update( + attn_layer, + qkv, + q_out, + k_out, + positions, + q_weight, + k_weight, + rms_norm_eps, + cos_sin_cache, + is_neox, + kv_cache, + layer_slot_mapping, + ) + else: + # Profiling/dummy run: define q_out/k_out (consumed by attention). + q_out.zero_() + k_out.zero_() + + return torch.empty(0, device=qkv.device, dtype=qkv.dtype) + + +def fused_qk_norm_rope_and_unified_kv_cache_update_fake( + q_out: torch.Tensor, + k_out: torch.Tensor, + qkv: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + is_neox: bool, + layer_name: str = "", +) -> torch.Tensor: + return torch.empty(0, device=qkv.device, dtype=qkv.dtype) + + +direct_register_custom_op( + op_name="fused_qk_norm_rope_and_unified_kv_cache_update", + op_func=fused_qk_norm_rope_and_unified_kv_cache_update_impl, + mutates_args=["q_out", "k_out"], + fake_impl=fused_qk_norm_rope_and_unified_kv_cache_update_fake, +) + + +# --------------------------------------------------------------------------- +# Pattern: QK-norm + RoPE + unified_kv_cache_update +# --------------------------------------------------------------------------- + + +class QkNormRopeKvCachePattern: + """ + Match the unfused sequence: + q, k, v = split(qkv, ...) + q = rms_norm(q.view(heads), q_weight).view(flat) + k = rms_norm(k.view(heads), k_weight).view(flat) + q, k = rotary_embedding(positions, q, k, cos_sin_cache, is_neox) + q = q.view(num_heads, head_dim) + k = k.view(num_kv_heads, head_dim) + v = v.view(num_kv_heads, head_dim) + dummy = unified_kv_cache_update(k, v, layer_name) + + Replace with: + q_out = empty(...) + k_out = empty(...) + dummy = fused_qk_norm_rope_and_unified_kv_cache_update( + q_out, k_out, qkv, positions, q_weight, k_weight, + eps, cos_sin_cache, is_neox, layer_name) + v = split(qkv, ...)[2].view(num_kv_heads, head_dim) + """ + + FUSED_OP = torch.ops.vllm.fused_qk_norm_rope_and_unified_kv_cache_update.default + + def __init__( + self, + layer: Attention, + eps: float, + is_neox: bool, + quant_query: bool, + ) -> None: + self.layer_name = layer.layer_name + self.num_heads = layer.num_heads + self.num_kv_heads = layer.num_kv_heads + self.head_size = layer.head_size + self.head_size_v = layer.head_size_v + self.eps = eps + self.is_neox = is_neox + self.quant_query = quant_query + + self.q_size = self.num_heads * self.head_size + self.k_size = self.num_kv_heads * self.head_size + self.v_size = self.num_kv_heads * self.head_size_v + + self.rope_matcher = MatcherRotaryEmbedding( + is_neox=is_neox, + head_size=self.head_size, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + ) + + def get_inputs(self) -> list[torch.Tensor]: + T = 5 + L = 4096 + qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size) + positions = empty_i64(T) + q_weight = empty_bf16(1, self.head_size) + k_weight = empty_bf16(1, self.head_size) + cos_sin_cache = empty_bf16(L, self.head_size) + inputs = [qkv, positions, q_weight, k_weight, cos_sin_cache] + if self.quant_query: + q_scale = empty_fp32(1) + inputs += [q_scale] + return inputs + + def pattern_non_fp8_quant_query( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size) + q_normed = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps) + q_flat = q_normed.view(-1, self.q_size) + + k_by_head = k.view(-1, self.k_size // self.head_size, self.head_size) + k_normed = vllm.ir.ops.rms_norm(k_by_head, k_weight, self.eps) + k_flat = k_normed.view(-1, self.k_size) + + q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache) + + q_rope = q_rope.view(-1, self.num_heads, self.head_size) + k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size) + v = v.view(-1, self.num_kv_heads, self.head_size_v) + dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name) + return dummy, q_rope, k_rope, v + + def replacement_non_fp8_quant_query( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q_out = torch.empty( + qkv.shape[0], + self.num_heads, + self.head_size, + device=qkv.device, + dtype=qkv.dtype, + ) + k_out = torch.empty( + qkv.shape[0], + self.num_kv_heads, + self.head_size, + device=qkv.device, + dtype=qkv.dtype, + ) + _, _, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + v = v.view(qkv.shape[0], self.num_kv_heads, self.head_size_v) + results = auto_functionalized( + self.FUSED_OP, + q_out=q_out, + k_out=k_out, + qkv=qkv, + positions=positions, + q_weight=q_weight, + k_weight=k_weight, + rms_norm_eps=self.eps, + cos_sin_cache=cos_sin_cache, + is_neox=self.is_neox, + layer_name=self.layer_name, + ) + return results[0], results[1], results[2], v + + def pattern_fp8_quant_query( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size) + q_normed = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps) + q_flat = q_normed.view(-1, self.q_size) + + k_by_head = k.view(-1, self.k_size // self.head_size, self.head_size) + k_normed = vllm.ir.ops.rms_norm(k_by_head, k_weight, self.eps) + k_flat = k_normed.view(-1, self.k_size) + + q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache) + # Match the quant-query op Attention.forward inserts (fp8 KV + UNIFIED). + # Explicit auto_functionalized (out=[1]) keeps the quant node in the pattern. + q_out = torch.empty_like(q_rope, dtype=current_platform.fp8_dtype()) + q_quant = auto_functionalized( + torch.ops.vllm.rocm_aiter_per_tensor_quant.default, + out=q_out, + x=q_rope, + scale=q_scale, + is_dynamic=False, + ) + # `scale` is mutable: its copy_ write-back to _q_scale bumps the mutation + # region, so keep q flat (a reshape lands past the barrier and won't match). + q_rope_fp8 = q_quant[1] + q_scale_out = q_quant[2] + + k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size) + v = v.view(-1, self.num_kv_heads, self.head_size_v) + dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name) + return dummy, q_rope_fp8, k_rope, v, q_scale_out + + def replacement_fp8_quant_query( + self, + qkv: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q_out = torch.empty( + qkv.shape[0], + self.num_heads, + self.head_size, + device=qkv.device, + dtype=qkv.dtype, + ) + k_out = torch.empty( + qkv.shape[0], + self.num_kv_heads, + self.head_size, + device=qkv.device, + dtype=qkv.dtype, + ) + _, _, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + v = v.view(qkv.shape[0], self.num_kv_heads, self.head_size_v) + results = auto_functionalized( + self.FUSED_OP, + q_out=q_out, + k_out=k_out, + qkv=qkv, + positions=positions, + q_weight=q_weight, + k_weight=k_weight, + rms_norm_eps=self.eps, + cos_sin_cache=cos_sin_cache, + is_neox=self.is_neox, + layer_name=self.layer_name, + ) + # Re-apply the quant on the kernel's bf16 q_out; fused op does not quant q. + # Same explicit auto_functionalized form as the pattern: [1] = quantized + # q, [2] = scale (returned so the buffer-writeback use is preserved). + q_fp8_flat = results[1].view(-1, self.q_size) + q_fp8_out = torch.empty_like(q_fp8_flat, dtype=current_platform.fp8_dtype()) + q_requant = auto_functionalized( + torch.ops.vllm.rocm_aiter_per_tensor_quant.default, + out=q_fp8_out, + x=q_fp8_flat, + scale=q_scale, + is_dynamic=False, + ) + q_fp8 = q_requant[1] # flat to mirror the pattern (see note above) + q_scale_out = q_requant[2] + return results[0], q_fp8, results[2], v, q_scale_out + + @staticmethod + def wrap_trace_fn( + trace_fn: Callable[P, fx.GraphModule], + *process_fx_fns: Callable[[fx.GraphModule], None], + ) -> Callable[P, fx.GraphModule]: + def wrapped(*args: P.args, **kwargs: P.kwargs) -> fx.GraphModule: + gm = trace_fn(*args, **kwargs) + for process_fx in process_fx_fns: + process_fx(gm) + + return gm + + return wrapped + + @staticmethod + def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None: + from torch._inductor.fx_passes.post_grad import view_to_reshape + + view_to_reshape(gm) + + def _register(self, pattern, replacement, pm_pass) -> None: + trace_fn = QkNormRopeKvCachePattern.wrap_trace_fn( + pm.fwd_only, + QkNormRopeKvCachePattern.fx_view_to_reshape, + ) + + # Pre-build the search pattern with `ignore_types=(int, torch.SymInt)` + # and pass it via `search_fn_pattern=` so torch skips both of its + # internal `fx_to_pattern` calls and treats dynamic-shape SymInts as + # wildcards. + inputs = self.get_inputs() + argnames = [*inspect.signature(pattern).parameters.keys()] + search_gm = trace_fn(pattern, inputs) + search_fn_pattern = pm.fx_to_pattern( + search_gm, + ignore_types=(int, torch.SymInt), + argnames=argnames, + ) + + pm.register_replacement( + pattern, + replacement, + inputs, + trace_fn, + pm_pass, + search_fn_pattern=search_fn_pattern, + ) + + def register(self, pm_pass: PatternMatcherPass) -> None: + # make_fx counts `self` in bound-method code params; wrap as plain fns. + # Distinct names per branch so mypy doesn't see one name, two signatures. + if self.quant_query: + + def pattern_q(qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale): + return self.pattern_fp8_quant_query( + qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale + ) + + def replacement_q( + qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale + ): + return self.replacement_fp8_quant_query( + qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale + ) + + self._register(pattern_q, replacement_q, pm_pass) + else: + + def pattern_noq(qkv, positions, q_weight, k_weight, cos_sin_cache): + return self.pattern_non_fp8_quant_query( + qkv, positions, q_weight, k_weight, cos_sin_cache + ) + + def replacement_noq(qkv, positions, q_weight, k_weight, cos_sin_cache): + return self.replacement_non_fp8_quant_query( + qkv, positions, q_weight, k_weight, cos_sin_cache + ) + + self._register(pattern_noq, replacement_noq, pm_pass) + + +# --------------------------------------------------------------------------- +# Pass class +# --------------------------------------------------------------------------- + + +class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass): + """ + Fuse QK-norm + RoPE + KV cache update into a single AITER HIP kernel. + + Supersedes both QKNormRoPEFusionPass and RopeKVCacheFusionPass for + attention layers that support the combined operation, eliminating two + separate kernel launches and the intermediate memory traffic. + """ + + @enable_fake_mode + def __init__(self, config: VllmConfig) -> None: + super().__init__(config) + + self.patterns: PatternMatcherPass = PatternMatcherPass( + pass_name="qk_norm_rope_kvcache_fusion_pass" + ) + + cc = config.compilation_config + self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num + + dtype = config.model_config.dtype + if dtype not in (torch.bfloat16, torch.float16): + logger.warning_once( + "QK Norm+RoPE+KVCache fusion not enabled: unsupported dtype %s", dtype + ) + return + + attn_layers = get_layers_from_vllm_config(config, Attention) + + for _, layer in attn_layers.items(): + if not layer.impl.fused_qk_norm_rope_kvcache_supported(): + continue + if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS: + logger.warning_once( + "QK Norm+RoPE+KVCache fusion not enabled for a layer: " + "head_size=%d is not supported by the " + "fused_qk_norm_rope_cache_pts_quant_shuffle kernel " + "(supported: %s). Falling back to the unfused path.", + layer.head_size, + SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS, + ) + continue + if layer.head_size_v != layer.head_size: + # The fused kernel uses a single head_dim for q/k/v. + logger.warning_once( + "QK Norm+RoPE+KVCache fusion not enabled for a layer: " + "head_size_v=%d differs from head_size=%d, which the fused " + "kernel does not support. Falling back to the unfused path.", + layer.head_size_v, + layer.head_size, + ) + continue + for epsilon in [1e-5, 1e-6]: + for neox in [True, False]: + for quant_q in [False, True]: + QkNormRopeKvCachePattern( + layer=layer, + eps=epsilon, + is_neox=neox, + quant_query=quant_q, + ).register(self.patterns) + + self.dump_patterns(config, self.patterns) + + @VllmInductorPass.time_and_log + def __call__(self, graph: fx.Graph) -> None: + self.matched_count = self.patterns.apply(graph) + logger.info( + "QK-Norm+RoPE+KVCache fusion: replaced %s pattern(s) " + "with AITER fused_qk_norm_rope_cache_pts_quant_shuffle", + self.matched_count, + ) + + def is_applicable_for_range(self, compile_range: Range) -> bool: + return compile_range.end <= self.max_token_num + + def uuid(self) -> str: + return VllmInductorPass.hash_source(self, QkNormRopeKvCachePattern) diff --git a/vllm/compilation/passes/pass_manager.py b/vllm/compilation/passes/pass_manager.py index 4b98ac57745a..67a3e4cbae5c 100644 --- a/vllm/compilation/passes/pass_manager.py +++ b/vllm/compilation/passes/pass_manager.py @@ -38,6 +38,7 @@ from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass + from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass from .fusion.rms_quant_fusion import RMSNormQuantFusionPass from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass from .utility.scatter_split_replace import ScatterSplitReplacementPass @@ -171,6 +172,11 @@ def configure(self, config: VllmConfig) -> None: if rocm_aiter_ops.is_enabled(): self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)] + if self.pass_config.fuse_qk_norm_rope_kvcache: + self.passes += [SplitCoalescingPass(config)] + self.passes += [ScatterSplitReplacementPass(config)] + self.passes += [QkNormRopeKvCacheFusionPass(config)] + if ( self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled() diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 810c40131fc7..810e97c4cadc 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -146,10 +146,16 @@ class PassConfig: """Fuse paired q/kv RMS norms in MLA attention.""" fuse_rope_kvcache: bool = None # type: ignore[assignment] """Fuse the QK rope + KV cache ops.""" + fuse_qk_norm_rope_kvcache: bool = Field(default=None) # type: ignore[assignment] + """Fuse QK RMSNorm + RoPE + KV cache update into a single AITER HIP + kernel. Supersedes both enable_qk_norm_rope_fusion and fuse_rope_kvcache + for layers that support it. Auto-enabled at O1+ on ROCm for models + with QK-norm (e.g. Qwen3-MoE).""" rope_kvcache_fusion_max_token_num: int = 256 """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode. Larger batch sizes e.g. during prefill will use the unfused kernels. + Also applies to the fused QK-Norm+RoPE+KVCache pass. """ fi_allreduce_fusion_max_size_mb: float | None = None @@ -228,6 +234,8 @@ def compute_hash(self) -> str: "fuse_act_padding", "fuse_mla_dual_rms_norm", "fuse_rope_kvcache", + "fuse_qk_norm_rope_kvcache", + "enable_qk_norm_rope_fusion", "fuse_rope_kvcache_cat_mla", mode="wrap", ) @@ -288,6 +296,12 @@ def __post_init__(self) -> None: "The fusion will be disabled." ) self.fuse_rope_kvcache = False + if self.fuse_qk_norm_rope_kvcache and not current_platform.is_rocm(): + logger.warning_once( + "QK-Norm+RoPE+KVCache fusion requires ROCm with AITER. " + "The fusion will be disabled." + ) + self.fuse_qk_norm_rope_kvcache = False if self.fuse_rope_kvcache_cat_mla and not current_platform.is_cuda_alike(): logger.warning_once( "MLA KV cache update with RoPE fusion enabled but the " @@ -302,10 +316,13 @@ def log_enabled_passes(self) -> None: after all defaults are finalized. TODO also log the compile ranges for which this is enabled. """ + fusion_prefixes = ("fuse_", "enable_") enabled_fusions = [ - f.name[len("fuse_") :] + f.name[len(prefix) :] for f in fields(self) # type: ignore[arg-type] - if getattr(self, f.name) and f.name.startswith("fuse_") + if getattr(self, f.name) + for prefix in fusion_prefixes + if f.name.startswith(prefix) ] if enabled_fusions: @@ -947,6 +964,7 @@ def __post_init__(self) -> None: # TODO(zhuhaoran): support rope native forward match and remove this. # Linked issue: https://github.com/vllm-project/vllm/issues/28042 self.custom_ops.append("+rotary_embedding") + if ( self.pass_config.fuse_rope_kvcache and "+rotary_embedding" not in self.custom_ops @@ -955,6 +973,12 @@ def __post_init__(self) -> None: # Linked issue: https://github.com/vllm-project/vllm/issues/28042 self.custom_ops.append("+rotary_embedding") + if ( + self.pass_config.fuse_qk_norm_rope_kvcache + and "+rotary_embedding" not in self.custom_ops + ): + self.custom_ops.append("+rotary_embedding") + if ( is_torch_equal_or_newer("2.9.0.dev") and "combo_kernels" not in self.inductor_compile_config @@ -1137,6 +1161,16 @@ def set_splitting_ops_for_v1( "to enable RoPE+KV cache fusion." ) self.pass_config.fuse_rope_kvcache = False + if self.pass_config.fuse_qk_norm_rope_kvcache: + logger.warning_once( + "fuse_qk_norm_rope_kvcache is enabled, but " + "splitting_ops is None and Inductor graph partition " + "is not enabled. Disabling fuse_qk_norm_rope_kvcache. " + "Please either set splitting_ops to an empty list [] " + "or set use_inductor_graph_partition to True " + "to enable QK-Norm+RoPE+KV cache fusion." + ) + self.pass_config.fuse_qk_norm_rope_kvcache = False self.splitting_ops.append("vllm::unified_kv_cache_update") self.splitting_ops.append("vllm::unified_mla_kv_cache_update") diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index aff94fe03d28..f13562f59c0e 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -178,6 +178,9 @@ class KernelConfig: enable_cutedsl_warmup: bool = True """If True, run CuTeDSL compile warmup during kernel warmup.""" + enable_bf16x3_router_gemm: bool = False + """If True, use the experimental SM100 BF16x3 CuteDSL router GEMM.""" + moe_backend: MoEBackend = "auto" """Backend for MoE expert computation kernels. Available options: diff --git a/vllm/config/lora.py b/vllm/config/lora.py index 94a679941c85..4323da96d806 100644 --- a/vllm/config/lora.py +++ b/vllm/config/lora.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from typing import TYPE_CHECKING, Any, Literal import torch @@ -75,8 +76,14 @@ class LoRAConfig: """If True, force the engine to use the universal 2D MoE LoRA wrapper (`FusedMoEWithLoRA`) regardless of the model's `is_3d_moe_weight` flag, so that 2D-format and 3D-format MoE LoRA adapters can be served in the same - deployment. Only meaningful forMoE models; ignored otherwise. Default False + deployment. Only meaningful for MoE models; ignored otherwise. Default False keeps the existing model-driven behavior.""" + enable_moe_shared_loras: bool = False + """If True, load MoE expert adapters in the "shared-outer" layout, where the + gate/up (`w1`/`w3`) lora_A and the down (`w2`) lora_B are shared across all + experts (stored once with expert-dim 1) instead of per-expert. The shared + factors are broadcast to the expert count at kernel time. Only meaningful for + MoE models whose adapters use this layout; ignored otherwise.""" def compute_hash(self) -> str: """ @@ -97,6 +104,7 @@ def compute_hash(self) -> str: factors.append(self.lora_dtype) factors.append(self.enable_tower_connector_lora) factors.append(self.enable_mixed_moe_lora_format) + factors.append(self.enable_moe_shared_loras) # target_modules affects which modules get LoRA applied factors.append( tuple(sorted(self.target_modules)) if self.target_modules else None @@ -129,3 +137,13 @@ def verify_with_model_config(self, model_config: ModelConfig): self.lora_dtype = model_config.dtype elif isinstance(self.lora_dtype, str): self.lora_dtype = getattr(torch, self.lora_dtype) + + architectures = getattr(model_config, "architectures", None) or [] + is_inkling = any("Inkling" in arch for arch in architectures) + if is_inkling and os.environ.get("INKLING_MULTIMEM_AR", "1") != "0": + raise ValueError( + "Inkling LoRA requires INKLING_MULTIMEM_AR=0: the Lamport " + "fused-collective path bypasses the LoRA-wrapped wo_ud and dense " + "down_proj layers on decode-sized batches, silently dropping " + "their LoRA. Set INKLING_MULTIMEM_AR=0." + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index 0ce01de166c6..e36b672cd828 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -83,7 +83,9 @@ RunnerOption = Literal["auto", RunnerType] ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] -TokenizerMode = Literal["auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4"] +TokenizerMode = Literal[ + "auto", "hf", "slow", "mistral", "deepseek_v32", "deepseek_v4", "inkling" +] ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"] LogprobsMode = Literal[ "raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs" @@ -628,6 +630,8 @@ def __post_init__( self.tokenizer_mode = "deepseek_v32" elif arch == "DeepseekV4ForCausalLM": self.tokenizer_mode = "deepseek_v4" + elif arch in ("InklingForCausalLM", "InklingForConditionalGeneration"): + self.tokenizer_mode = "inkling" if self.tokenizer_mode != "auto": logger.info( diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 909e4418f46e..7c270b0c0eb9 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -652,6 +652,7 @@ def use_sequence_parallel_moe(self) -> bool: ) and self.enable_expert_parallel and self.tensor_parallel_size > 1 + and self.data_parallel_size > 1 ) @property diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index 34d7fa5d6364..370fbaf2731f 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -18,6 +18,7 @@ kInt8StaticChannelSym, kMxfp4Dynamic, kMxfp8Dynamic, + kNvfp4Static, ) # User-facing names addressable from quantization_config. @@ -134,6 +135,11 @@ def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any: "int8_per_channel_weight_only": QuantizationConfigArgs( moe=QuantSpec(weight=kInt8StaticChannelSym), ), + # Online NVFP4 on MoE with per-token dynamic activation scales (Blackwell + + # FlashInfer TRTLLM only); linear stays unquantized (no `linear` field). + "nvfp4_per_token": QuantizationConfigArgs( + moe=QuantSpec(weight=kNvfp4Static), + ), } diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index b9c612621982..8ba55a2ec964 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -10,6 +10,7 @@ from typing_extensions import Self from vllm.config import LoadConfig +from vllm.config.cache import CacheDType from vllm.config.kernel import MoEBackend from vllm.config.model import HfOverrides, ModelConfig from vllm.config.parallel import ParallelConfig @@ -54,6 +55,7 @@ "step3p5_mtp", "hy_v3_mtp", "gemma4_mtp", + "inkling_mtp", ] NgramGPUTypes = Literal["ngram_gpu"] DFlashModelTypes = Literal["dflash"] @@ -118,6 +120,9 @@ class SpeculativeConfig: """Attention backend to use for the draft model. When `None`, the backend is automatically selected. Useful when the drafter requires a different attention backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN).""" + kv_cache_dtype: CacheDType | None = None + """KV cache dtype for the draft model. When `None`, the draft inherits the + target model's `--kv-cache-dtype`.""" max_model_len: int | None = Field(default=None, ge=1) """The maximum model length of the draft model. Used when testing the ability to skip speculation for some sequences.""" @@ -549,6 +554,26 @@ def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig: {"n_predict": n_predict, "architectures": ["HYV3MTPModel"]} ) + if hf_config.model_type in ("inkling_mm_model", "inkling_model"): + mtp_config = getattr(hf_config, "mtp_config", None) or {} + hf_config = getattr(hf_config, "text_config", hf_config) + checkpoint_depths = mtp_config.get("num_nextn_predict_layers", 0) + if checkpoint_depths < 1: + raise ValueError("The Inkling checkpoint does not contain MTP weights") + hf_config.model_type = "inkling_mtp" + hf_config.update( + { + # Inkling currently exposes only the first checkpoint depth. + "n_predict": 1, + "num_nextn_predict_layers": checkpoint_depths, + "chain_hidden_post_norm": mtp_config.get( + "chain_hidden_post_norm", False + ), + "local_layer_ids": mtp_config.get("local_layer_ids", []), + "architectures": ["InklingMTPModel"], + } + ) + if hf_config.model_type in ("gemma4_assistant", "gemma4_unified_assistant"): hf_config.model_type = "gemma4_mtp" text_config = getattr(hf_config, "text_config", hf_config) @@ -626,6 +651,18 @@ def compose_draft_hf_overrides( SpeculativeConfig._apply_composed_hf_override, target_hf_overrides ) + @staticmethod + def _is_custom_proposer_path(model: str | None) -> bool: + """True if ``model`` is a dotted import path (e.g. ``pkg.MyProposer``).""" + if model is None: + return False + if model.startswith(("http://", "https://", "file://")): + return False + if "/" in model: + return False + parts = model.split(".") + return len(parts) >= 2 and all(part.isidentifier() for part in parts) + def __post_init__(self): # Note: "method" is a new parameter that helps to extend the # configuration of non-model-based proposers, and the "model" parameter @@ -636,14 +673,9 @@ def __post_init__(self): # default. # infer method from user args - # Check if the model field contains a custom module path (e.g., 'pkg.Mod') - if ( - self.model is not None - and "." in self.model - and not self.model.startswith(("http://", "https://", "file://")) - and "/" not in self.model # not a HuggingFace repo (org/model) + if self.method is None and SpeculativeConfig._is_custom_proposer_path( + self.model ): - # Treat as a custom class path self.method = "custom_class" elif self.method is None: if self.model in ("ngram", "[ngram]"): @@ -850,6 +882,7 @@ def __post_init__(self): elif ( "dspark" in self.draft_model_config.model.lower() or "Qwen3DSparkModel" in self.draft_model_config.architectures + or "Gemma4DSparkModel" in self.draft_model_config.architectures ): self.method = "dspark" elif self.draft_model_config.hf_config.model_type == "medusa": @@ -863,7 +896,7 @@ def __post_init__(self): if ( self.num_speculative_tokens > 1 and self.draft_model_config.hf_config.model_type - != "step3p5_mtp" + not in ("step3p5_mtp", "inkling_mtp") ): logger.warning( "Enabling num_speculative_tokens > 1 will run " @@ -900,6 +933,7 @@ def __post_init__(self): if self.method == "dspark" and ( "Qwen3DSparkModel" not in self.draft_model_config.architectures + and "Gemma4DSparkModel" not in self.draft_model_config.architectures ): # DeepSeek-V4 DSpark reuses the full DeepSeek-V4 config # and its weights ship in the target checkpoint. @@ -908,6 +942,23 @@ def __post_init__(self): "DSparkDraftModel" ] self.update_arch_() + elif ( + self.method == "dspark" + and "Gemma4DSparkModel" in self.draft_model_config.architectures + ): + # Normalize the self-contained Gemma4 draft's config keys to + # the DSpark conventions. + hf = self.draft_model_config.hf_config + if ( + getattr(hf, "dspark_target_layer_ids", None) is None + and getattr(hf, "target_layer_ids", None) is not None + ): + hf.dspark_target_layer_ids = hf.target_layer_ids + if ( + getattr(hf, "n_predict", None) is None + and getattr(hf, "block_size", None) is not None + ): + hf.n_predict = hf.block_size if self.method in ("dflash", "dspark"): self.parallel_drafting = True @@ -942,6 +993,14 @@ def __post_init__(self): "`num_speculative_tokens` was not provided" ) + if ( + self.draft_model_config.hf_config.model_type == "inkling_mtp" + and self.num_speculative_tokens != 1 + ): + raise ValueError( + "Inkling MTP currently supports exactly one speculative token" + ) + if self.method == "dspark": # DSpark is a semi-autoregressive *block* drafter. A # speculative length smaller than the checkpoint's block diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index dc0180034141..2e57cb1ec8e4 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -68,9 +68,11 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset( { "DeepseekV2ForCausalLM", - "Qwen2MoeForCausalLM", "GraniteMoeForCausalLM", + "InklingForCausalLM", + "InklingForConditionalGeneration", "LongcatFlashNgramForCausalLM", + "Qwen2MoeForCausalLM", } ) @@ -191,6 +193,15 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: return rocm_aiter_ops.is_enabled() and check_aiter_fused_qk_rmsnorm() +def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: + """Enable fused QK-norm + RoPE + KV cache update on ROCm with AITER.""" + from vllm._aiter_ops import rocm_aiter_ops + + if not rocm_aiter_ops.is_enabled(): + return False + return cfg.compilation_config.is_custom_op_enabled("rotary_embedding") + + OPTIMIZATION_LEVEL_00 = { "compilation_config": { "pass_config": { @@ -203,6 +214,8 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: "fuse_act_padding": False, "fuse_mla_dual_rms_norm": False, "fuse_rope_kvcache": False, + "fuse_qk_norm_rope_kvcache": False, + "enable_qk_norm_rope_fusion": False, "fuse_rope_kvcache_cat_mla": False, }, "cudagraph_mode": CUDAGraphMode.NONE, @@ -224,6 +237,8 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": False, + "fuse_qk_norm_rope_kvcache": False, + "enable_qk_norm_rope_fusion": False, "fuse_rope_kvcache_cat_mla": False, }, "cudagraph_mode": CUDAGraphMode.PIECEWISE, @@ -245,6 +260,8 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, + "fuse_qk_norm_rope_kvcache": enable_qk_norm_rope_kvcache, + "enable_qk_norm_rope_fusion": False, "fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion, }, "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE, @@ -266,6 +283,8 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool: "fuse_act_padding": enable_norm_pad_fusion, "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion, "fuse_rope_kvcache": enable_rope_kvcache_fusion, + "fuse_qk_norm_rope_kvcache": enable_qk_norm_rope_kvcache, + "enable_qk_norm_rope_fusion": False, "fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion, }, "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE, @@ -1166,6 +1185,8 @@ def __post_init__(self): in ( "DeepseekV4ForCausalLM", "DeepSeekV4MTPModel", + "InklingForCausalLM", + "InklingForConditionalGeneration", "MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration", ) @@ -1967,6 +1988,21 @@ def _set_compile_ranges(self): compile_range_end, ) + if compilation_config.pass_config.fuse_qk_norm_rope_kvcache: + max_token_num = ( + compilation_config.pass_config.rope_kvcache_fusion_max_token_num + ) + if max_token_num is not None: + if compile_range_end is not None and max_token_num < compile_range_end: + computed_compile_ranges_endpoints.append(max_token_num) + else: + logger.debug( + "Max num batched tokens below qk_norm+rope+kvcache " + "fusion threshold, fusion enabled for " + "num_tokens <= %d.", + compile_range_end, + ) + if compilation_config.compile_ranges_endpoints is not None: for x in compilation_config.compile_ranges_endpoints: assert isinstance(x, int) diff --git a/vllm/cute_utils/__init__.py b/vllm/cute_utils/__init__.py index 1eee51019fda..ca445284bf41 100644 --- a/vllm/cute_utils/__init__.py +++ b/vllm/cute_utils/__init__.py @@ -1,11 +1,33 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from cutlass import BFloat16, Float32, Int64, Uint32, cute +import torch +from cutlass import ( + BFloat16, + Float8E4M3FN, + Float16, + Float32, + Int32, + Int64, + Uint32, + cute, +) from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, vector from cutlass.cute.nvgpu import cpasync from cutlass.cutlass_dsl import T, dsl_user_op +_TORCH_TO_CUTE_DTYPE = { + torch.bfloat16: BFloat16, + torch.float8_e4m3fn: Float8E4M3FN, +} + +_CUTE_TO_PTX_DTYPE = { + BFloat16: "bf16", + Float16: "f16", + Float8E4M3FN: "e4m3", + Float32: "f32", +} + # https://github.com/NVIDIA/cutlass/blob/v4.3.2/include/cute/arch/copy_sm90_desc.hpp#L193-L197 EVICT_NORMAL = Int64(0x1000000000000000) EVICT_FIRST = Int64(0x12F0000000000000) @@ -63,22 +85,35 @@ def fence_before_tma_store(*, loc=None, ip=None): @dsl_user_op -def mma_bf16( - a: cute.TensorSSA, b: cute.TensorSSA, c: cute.TensorSSA, *, loc=None, ip=None -): - if a.element_type == BFloat16: - a = cute.recast_tensor(a, Uint32) - if b.element_type == BFloat16: - b = cute.recast_tensor(b, Uint32) - - mlir_ty = Float32.mlir_type +def mma_sync(a, b, c: cute.Tensor, *, loc=None, ip=None): + a_ty = _CUTE_TO_PTX_DTYPE[a.element_type] + b_ty = _CUTE_TO_PTX_DTYPE[b.element_type] + c_ty = _CUTE_TO_PTX_DTYPE[c.element_type] + mlir_ty = c.element_type.mlir_type + K = 256 // a.element_type.width # 32B + + # Recast expects tensor-backed fragments; materialize SSA fragments here so + # callsites can pass converted FP8 fragments directly. + if isinstance(a, cute.TensorSSA): + a_ = cute.make_rmem_tensor_like(a) + a_.store(a, loc=loc, ip=ip) + a = a_ + if isinstance(b, cute.TensorSSA): + b_ = cute.make_rmem_tensor_like(b) + b_.store(b, loc=loc, ip=ip) + b = b_ + + a = cute.recast_tensor(a, Int32, loc=loc, ip=ip) + b = cute.recast_tensor(b, Int32, loc=loc, ip=ip) out = llvm.inline_asm( llvm.StructType.get_literal([mlir_ty] * 4), [a[i].ir_value(loc=loc, ip=ip) for i in range(4)] + [b[i].ir_value(loc=loc, ip=ip) for i in range(2)] + [c[i].ir_value(loc=loc, ip=ip) for i in range(4)], - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, " + f"mma.sync.aligned.m16n8k{K}.row.col.{c_ty}.{a_ty}.{b_ty}.{c_ty} " + "{$0, $1, $2, $3}, " + "{$4, $5, $6, $7}, " + "{$8, $9}, " "{$10, $11, $12, $13};", "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", has_side_effects=False, @@ -92,7 +127,7 @@ def mma_bf16( loc=loc, ip=ip, ) - return cute.TensorSSA(vec, 4, Float32) + return cute.TensorSSA(vec, 4, c.element_type) def _bf16x2_unary(asm: str, a: Uint32, *, loc=None, ip=None) -> Uint32: diff --git a/vllm/cute_utils/cvt.py b/vllm/cute_utils/cvt.py index f4e8f0b0bc9b..4707f3f7421e 100644 --- a/vllm/cute_utils/cvt.py +++ b/vllm/cute_utils/cvt.py @@ -84,6 +84,32 @@ def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA: return cute.TensorSSA(vec, 2, Uint32) +@dsl_user_op +def fp8x4_to_fp16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA: + out = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * 2), + [x.ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, $2;\n\t" + "cvt.rn.f16x2.e4m3x2 $0, lo;\n\t" + "cvt.rn.f16x2.e4m3x2 $1, hi;\n\t" + "}\n", + "=r,=r,r", + has_side_effects=False, + is_align_stack=False, + loc=loc, + ip=ip, + ) + vec = vector.from_elements( + ir.VectorType.get([2], T.i32(), loc=loc), + [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(2)], + loc=loc, + ip=ip, + ) + return cute.TensorSSA(vec, 2, Uint32) + + @dsl_user_op def fp32x4_to_fp8x4( a0: Float32, diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 81006da401d4..7f540bc4b1f0 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -16,6 +16,7 @@ has_flashinfer_nvlink_one_sided, has_flashinfer_nvlink_two_sided, ) +from vllm.utils.func_utils import supports_kw from vllm.utils.import_utils import has_deep_ep, has_deep_ep_v2, has_mori from .base_device_communicator import All2AllManagerBase, Cache @@ -690,6 +691,7 @@ def __init__(self, cpu_group): self.max_num_tokens = 0 self.top_k = 0 self.num_experts = 0 + self._combine_supports_output = False def initialize( self, @@ -790,6 +792,12 @@ def initialize( workspace_size_per_rank=self.workspace_size, mnnvl_config=ep_config, ) + try: + self._combine_supports_output = supports_kw( + self.moe_alltoall.combine, "output", allow_var_kwargs=False + ) + except (TypeError, ValueError): + self._combine_supports_output = False self.gpus_per_node = gpus_per_node self.initialized = True @@ -804,6 +812,27 @@ def initialize( # different shape sequences, so a world-level barrier would deadlock. dist.barrier(group=self.cpu_group) + def combine_into( + self, + payload: torch.Tensor, + runtime_max_tokens_per_rank: int, + output: torch.Tensor, + ) -> None: + """Combine into ``output``, with a fallback for older FlashInfer.""" + assert self.moe_alltoall is not None + if self._combine_supports_output: + self.moe_alltoall.combine( + payload=payload, + runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, + output=output, + ) + else: + combined_output = self.moe_alltoall.combine( + payload=payload, + runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, + ) + output.copy_(combined_output) + def get_handle(self, kwargs): return self diff --git a/vllm/distributed/kv_events.py b/vllm/distributed/kv_events.py index c2faf34095de..be7e7363fb9d 100644 --- a/vllm/distributed/kv_events.py +++ b/vllm/distributed/kv_events.py @@ -74,6 +74,8 @@ class BlockStored(KVCacheEvent): # filter groups as they are learned. Remove events only need group_idx+hash. kv_cache_spec_kind: str | None = None kv_cache_spec_sliding_window: int | None = None + locality: str | None = None + """LOCAL or REMOTE relative to the publisher; None means unspecified.""" def __hash__(self) -> int: return hash( @@ -88,6 +90,7 @@ def __hash__(self) -> int: self.group_idx, self.kv_cache_spec_kind, self.kv_cache_spec_sliding_window, + self.locality, ) ) @@ -96,6 +99,8 @@ class BlockRemoved(KVCacheEvent): block_hashes: list[ExternalBlockHash] medium: str | None group_idx: int | None = None + locality: str | None = None + """LOCAL or REMOTE relative to the publisher; None means unspecified.""" def __hash__(self) -> int: return hash( @@ -103,6 +108,7 @@ def __hash__(self) -> int: tuple(self.block_hashes), self.medium, self.group_idx, + self.locality, ) ) 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 3d0d216d996f..ea41adf437ce 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 @@ -1030,6 +1030,16 @@ def __init__( use_mla=self.use_mla, ) self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + # READ-mode producer: a decode release-ACK can arrive BEFORE + # start_load_kv populates transfer_id_to_request_id (the notify races + # ahead of the scheduler->worker sync). Buffer such ACKs and retry them + # next get_finished tick instead of dropping them -- dropping loses the + # completion, so the request is never marked done_sending, its KV blocks + # leak, and the prefill KV cache wedges at high concurrency. Buffered + # BEFORE resolve_moriio_transfer_ack, so each ACK is counted exactly once + # (on the tick its mapping exists) -- the heterogeneous-TP ack-counting + # is preserved. + self._pending_unmapped_acks: list = [] # TODO: consider the integration of flashinfer or other backends. self.backend_name = backend.get_name() @@ -1542,16 +1552,22 @@ def get_finished(self) -> tuple[set[str], set[str]]: # pop_finished_req_ids returns release ACKs sent by decode. Keep # duplicate ACKs because heterogeneous TP can fan multiple decode # ranks into one prefill rank for the same transfer_id. - finished_acks = self.moriio_wrapper.pop_finished_req_ids() + # Combine freshly-arrived ACKs with any buffered from prior ticks + # whose transfer_id wasn't mapped yet (notify raced ahead of + # start_load_kv); retry the lookup every tick. Buffered before + # resolve_moriio_transfer_ack so each ACK is counted exactly once. + finished_acks = self._pending_unmapped_acks + list( + self.moriio_wrapper.pop_finished_req_ids() + ) + self._pending_unmapped_acks = [] resolved_transfer_ids: set[TransferId] = set() for ack in finished_acks: transfer_id = ack if isinstance(ack, str) else ack.transfer_id if transfer_id not in self.transfer_id_to_request_id: - logger.warning( - "Could not find %s in transfer_id_to_request_id " - "lookup table. This could lead to a possible hang.", - transfer_id, - ) + # Mapping not populated yet -- buffer and retry next tick, + # do NOT drop (dropping leaks producer KV at high conc and + # wedges the prefill). + self._pending_unmapped_acks.append(ack) continue resolved_transfer_id = resolve_moriio_transfer_ack( ack, @@ -1933,6 +1949,23 @@ def _compute_block_transfer_offsets( ), ) + @staticmethod + def _is_sq_full_status(status) -> bool: + """True if a MoRIIO transfer status is a transient RDMA send-queue-full + rejection (retryable backpressure), not a terminal failure. + + read_remote_data posts the RDMA READ synchronously (the mori executor + joins its worker before returning and marks the status on the calling + thread), so a send-queue-full rejection is a Failed() status the moment + the call returns. mori surfaces it as a generic ERR_RDMA_OP carrying + "SQ full" in the message (no distinct code), so we match the message. + Only meaningful once status.Failed() is True. + """ + try: + return bool(status.Failed()) and "SQ full" in (status.Message() or "") + except Exception: + return False + def _read_blocks( self, local_block_ids: list[int], @@ -1950,6 +1983,8 @@ def _read_blocks( dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0) sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id) + # SQ-full backpressure deadline, shared across this request's layers. + _sq_deadline = time.monotonic() + self.moriio_config.transfer_timeout for layer_name in self.layer_name_to_local_kv_cache_metadata: sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index( layer_name @@ -1962,9 +1997,35 @@ def _read_blocks( remote_tp_size=remote_tp_size, ) # TODO : apply multi-session batch-read when moriio support it - transfer_status = self.moriio_wrapper.read_remote_data( - offs[2], offs[0], offs[1], sessions[sess_idx] - ) + # + # SQ-full backpressure: read_remote_data posts the RDMA READ + # SYNCHRONOUSLY, so a send-queue-full rejection (per-QP HW cap) comes + # back as a Failed() status right here. A SEPARATE CQ-poll thread + # drains completions and frees SQ depth, so back off and RE-POST + # rather than let a transient rejection abort the request. No + # self-deadlock (the drain is off-thread); the reserve is + # all-or-nothing (nothing posted on a rejected attempt). Bounded by + # transfer_timeout; on sustained overload store the failed status and + # let get_finished handle it non-fatally (notify prefill + drop). + _backoff = 0.001 + while True: + transfer_status = self.moriio_wrapper.read_remote_data( + offs[2], offs[0], offs[1], sessions[sess_idx] + ) + if not self._is_sq_full_status(transfer_status): + break + if time.monotonic() > _sq_deadline: + logger.warning( + "MoRIIO READ send queue stayed full past " + "transfer_timeout for req %s layer %s; storing failed " + "status (get_finished notifies prefill and drops the " + "request). Raise qp_per_transfer if frequent.", + request_id, + layer_name, + ) + break + time.sleep(_backoff) + _backoff = min(_backoff * 2, 0.05) with self.moriio_wrapper.lock: self._recving_transfers[request_id].append(transfer_status) self._recving_transfers_callback_addr[request_id] = ( 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 d80541ade637..90aa5f823a81 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 @@ -312,10 +312,16 @@ 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]]: - if len(local_block_ids) != len(remote_block_ids): + # A shorter (or empty) local list is the READ-mode "drop the transfer, just + # free the prefill blocks" case (full-prefix-hit / aborted-before-scheduled): + # decode pulls fewer blocks than the prefill holds. The zip loop below pairs + # local[i]<->remote[i] and sizes by len(local), so a short local transfers + # only what decode allocated and an empty local is a no-op. A longer local + # list is a genuine bug and still fails loudly. + if len(local_block_ids) > len(remote_block_ids): raise ValueError( - "local_block_ids and remote_block_ids must have the same length: " - f"{len(local_block_ids)} != {len(remote_block_ids)}" + "local_block_ids longer than remote_block_ids: " + f"{len(local_block_ids)} > {len(remote_block_ids)}" ) geometry = get_layer_transfer_geometry( layer_name, kv_cache, layer_to_spec, remote_num_blocks diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index d48cb2a5ff37..834c5fccc7ac 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -160,7 +160,7 @@ def _compute_desc_ids( def _build_local_splits_from_plan( self, plan: TPMapping, - src_blocks_data: list[tuple[int, int, int]], + src_blocks_data: np.ndarray, num_fa_descs: int, ) -> Iterator[list[tuple[int, int, int]]]: """Build split handle data for P_TP > D_TP scenario. @@ -188,12 +188,13 @@ def _build_local_splits_from_plan( # Per-FA-descriptor replicate flag, in _build_fa_local emission order. fa_desc_replicated = self._fa_desc_replicated(num_fa_descs) + src_blocks_list = src_blocks_data.tolist() for p_idx, p_rank in enumerate(plan.all_source_ranks): fa_slot = plan.rank_to_attention_slot.get(p_rank, 0) handle: list[tuple[int, int, int]] = [] - for j, (addr, local_len, dev) in enumerate(src_blocks_data): + for j, (addr, local_len, dev) in enumerate(src_blocks_list): if j < num_fa_descs: if fa_desc_replicated[j]: # REPLICATE (MLA): whole block written on every rank. @@ -1277,51 +1278,48 @@ def _build_mamba_local( self, base_addresses: list[int], block_size_ratio: int, - ) -> list[tuple[int, int, int]]: + ) -> np.ndarray: """Build desc regions (conv sub-projections + ssm) per layer for - local mamba blocks with DS conv layout.""" + local mamba blocks with DS conv layout, as an Nx3 uint64 array.""" assert block_size_ratio == 1, ( "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " f"Got block_size_ratio={block_size_ratio}." ) + assert base_addresses, "Local KV cache base addresses must not be empty." assert self._conv_decomp is not None conv_offsets = self._conv_decomp.local_conv_offsets conv_size, ssm_size = self._mamba_ssm_size num_blocks = self._logical_num_blocks * block_size_ratio physical_per_logical = self._physical_blocks_per_logical_kv_block + device_id = self.device_id + block_arange = np.arange(num_blocks, dtype=np.uint64) - result: list[tuple[int, int, int]] = [] + parts: list[np.ndarray] = [] for i, base_addr in enumerate(base_addresses): # Jump one page_size, but ssm page_size may be bigger when kernel # locks block size to a specific value (physical_per_logical scale). page_stride = ( self.block_len_per_layer[i] // block_size_ratio * physical_per_logical ) + blk_addrs = base_addr + block_arange * page_stride for off, sz in conv_offsets: - for blk in range(num_blocks): - result.append( - (base_addr + blk * page_stride + off, sz, self.device_id) - ) + parts.append(self._stack_descs(blk_addrs + off, sz, device_id)) # SSM temporal state follows the conv state. - for blk in range(num_blocks): - result.append( - ( - base_addr + blk * page_stride + conv_size, - ssm_size, - self.device_id, - ) - ) - return result + parts.append(self._stack_descs(blk_addrs + conv_size, ssm_size, device_id)) + return np.concatenate(parts) def _build_mamba_remote( self, nixl_agent_meta: NixlAgentMetadata, tp_ratio: int, transfer_info: EngineTransferInfo, - ) -> list[tuple[int, int, int]]: + ) -> np.ndarray: """Build remote desc regions (conv sub-projections + ssm) per layer. For hetero-TP, each D rank reads only its sub-projection slice from - the P rank.""" + the P rank. Returns an Nx3 uint64 array.""" + assert nixl_agent_meta.kv_caches_base_addr, ( + "Remote KV cache base addresses must not be empty." + ) assert self._conv_decomp is not None effective_ratio = max(tp_ratio, 1) # Mamba conv state is always TP-sharded, even when attention KV @@ -1338,35 +1336,41 @@ def _build_mamba_remote( remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical device_id = nixl_agent_meta.device_id + block_arange = np.arange(num_blocks, dtype=np.uint64) - result: list[tuple[int, int, int]] = [] + parts: list[np.ndarray] = [] # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case # block lengths vary across layers (e.g. MLA). for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical + blk_addrs = base_addr + block_arange * page_stride for off, sz in conv_offsets: - for blk in range(num_blocks): - result.append((base_addr + blk * page_stride + off, sz, device_id)) + parts.append(self._stack_descs(blk_addrs + off, sz, device_id)) # SSM temporal state is also TP-sharded on the heads dimension. - for blk in range(num_blocks): - ssm_addr = ( - base_addr - + blk * page_stride - + conv_size_remote - + local_offset * ssm_read_size - ) - result.append((ssm_addr, ssm_read_size, device_id)) - return result + ssm_addrs = blk_addrs + conv_size_remote + local_offset * ssm_read_size + parts.append(self._stack_descs(ssm_addrs, ssm_read_size, device_id)) + return np.concatenate(parts) + + @staticmethod + def _stack_descs(addrs: np.ndarray, length: int, device_id: int) -> np.ndarray: + out = np.empty((addrs.shape[0], 3), dtype=np.uint64) + out[:, 0] = addrs + out[:, 1] = length + out[:, 2] = device_id + return out def _build_fa_local( self, base_addresses: list[int], block_size_ratio: int, - ) -> list[tuple[int, int, int]]: - """Build local FA descriptors for all layers.""" + ) -> np.ndarray: + """Build local FA descriptors for all layers as an Nx3 uint64 array.""" assert self.transfer_topo is not None + assert base_addresses, "Local KV cache base addresses must not be empty." num_blocks = self.num_blocks * block_size_ratio - result: list[tuple[int, int, int]] = [] + device_id = self.device_id + block_arange = np.arange(num_blocks, dtype=np.uint64) + parts: list[np.ndarray] = [] for i, base_addr in enumerate(base_addresses): kv_block_len = ( self.get_backend_aware_kv_block_len( @@ -1375,20 +1379,21 @@ def _build_fa_local( // block_size_ratio ) page_stride = self.block_len_per_layer[i] // block_size_ratio - for block_id in range(num_blocks): - block_offset = block_id * page_stride - addr = base_addr + block_offset - result.append((addr, kv_block_len, self.device_id)) - return result + addrs = base_addr + block_arange * page_stride + parts.append(self._stack_descs(addrs, kv_block_len, device_id)) + return np.concatenate(parts) def _build_fa_remote( self, plan: TPMapping, nixl_agent_meta: NixlAgentMetadata, block_size_ratio: int, - ) -> list[tuple[int, int, int]]: - """Build remote FA descriptors for all layers.""" + ) -> np.ndarray: + """Build remote FA descriptors for all layers as an Nx3 uint64 array.""" assert self.transfer_topo is not None + assert nixl_agent_meta.kv_caches_base_addr, ( + "Remote KV cache base addresses must not be empty." + ) fa_group_idx = next( i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) ) @@ -1396,7 +1401,9 @@ def _build_fa_remote( # per-rank offset; REPLICATE regions read the whole block once. split_reads = len(plan.source_ranks_per_group[fa_group_idx]) num_blocks = nixl_agent_meta.num_blocks - result: list[tuple[int, int, int]] = [] + device_id = nixl_agent_meta.device_id + block_arange = np.arange(num_blocks, dtype=np.uint64) + parts: list[np.ndarray] = [] for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): replicated = self._is_region_replicated(i) # Read our whole local region size from remote.. @@ -1417,18 +1424,14 @@ def _build_fa_remote( local_block_len = local_block_len // num_reads page_size = nixl_agent_meta.block_lens[i] - for block_id in range(num_blocks): - block_offset = block_id * page_size - # For each block, grab the kv heads chunk belonging to current local - # tp rank of size local_block_len. - addr = base_addr + block_offset + rank_offset - result.append((addr, local_block_len, nixl_agent_meta.device_id)) - return result + addrs = base_addr + rank_offset + block_arange * page_size + parts.append(self._stack_descs(addrs, local_block_len, device_id)) + return np.concatenate(parts) def register_local_xfer_handler( self, block_size: int, - ) -> tuple[int, list[tuple[int, int, int]]]: + ) -> tuple[int, np.ndarray]: """ Function used for register local xfer handler with local block_size or Remote block_size. @@ -1461,9 +1464,8 @@ def register_local_xfer_handler( # remote has been seen. Currently we always register 4 regions # because local descs are created before knowing the remote TP. logger.debug("Registering local Mamba descriptors (4 regions/layer)") - blocks_data.extend( - self._build_mamba_local(local_base_addresses, block_size_ratio) - ) + mamba = self._build_mamba_local(local_base_addresses, block_size_ratio) + blocks_data = np.concatenate([blocks_data, mamba]) descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) # NIXL_INIT_AGENT to be used for preparations of local descs. @@ -1650,13 +1652,8 @@ def add_remote_agent( engine_id, remote_tp_rank, ) - blocks_data.extend( - self._build_mamba_remote( - nixl_agent_meta, - tp_ratio, - transfer_info, - ) - ) + mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info) + blocks_data = np.concatenate([blocks_data, mamba]) # Register with NIXL. descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) @@ -1855,7 +1852,7 @@ def save_kv_to_host(self, metadata: NixlConnectorMetadata): for req_id, meta in metadata.reqs_to_save.items(): meta.local_physical_block_ids = self._logical_to_kernel_block_ids( - meta.local_block_ids + meta.local_block_ids, self._physical_blocks_per_logical_kv_block ) if logger.isEnabledFor(logging.DEBUG): logger.debug( @@ -2210,30 +2207,36 @@ def get_mapped_blocks( return mapped_2d.flatten().astype(np.int64) - def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds: + def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds: """ - Convert logical block ids to kernel physical block ids. + Convert block ids to kernel physical block ids. This is required when the logical block size (the one set by the user) does not match the one required by the attn backend. + `ratio` is the number of physical blocks per logical block. + We always receive logical blocks from the engine, so we expand them here eg: + logical block ids: [(SW-clipped) [1], (FA) [2, 3]], ratio=2 + physical block ids: [(SW-clipped) [2, 3], (FA) [4, 5, 6, 7]] """ - if self._physical_blocks_per_logical_kv_block == 1: + if ratio == 1: # Noop when physical and logical block sizes are the same return block_ids - block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape( - 1, -1 - ) - # Mamba blocks have no logical<>physical discrepancy + block_arange = np.arange(0, ratio).reshape(1, -1) + # Mamba blocks have no logical<>physical discrepancy (block-size=1) group_specs = self.kv_cache_config.kv_cache_groups - return [ - BlockTable.map_to_kernel_blocks( - np.array(group), - self._physical_blocks_per_logical_kv_block, - block_arange, - ).tolist() - if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) - else group - for i, group in enumerate(block_ids) - ] + physical_block_ids = [] + for i, group in enumerate(block_ids): + spec = group_specs[i].kv_cache_spec + if isinstance(spec, MambaSpec): + physical_block_ids.append(group) + else: + physical_block_ids.append( + BlockTable.map_to_kernel_blocks( + np.array(group), + ratio, + block_arange, + ).tolist() + ) + return physical_block_ids def _apply_prefix_caching( self, @@ -2311,39 +2314,6 @@ def _apply_prefix_caching( remote_block_ids[i] = remote_group[:num_blocks] return local_block_ids, remote_block_ids - def _logical_to_remote_kernel_block_ids( - self, block_ids: BlockIds, remote_physical_per_logical: int - ) -> BlockIds: - """Map logical block IDs to physical kernel block IDs on the remote. - - Args: - block_ids: per-group lists of logical block IDs. - remote_physical_per_logical: remote engine's physical blocks - per logical block. - - Returns: - Same structure with FA groups expanded (each logical block L - becomes kernel blocks [L*remote_physical_per_logical, .. - L*remote_physical_per_logical + - remote_physical_per_logical - 1]). - Mamba groups are passed through unchanged. - """ - if remote_physical_per_logical == 1: - return block_ids - remote_arange = np.arange(remote_physical_per_logical).reshape(1, -1) - group_specs = self.kv_cache_config.kv_cache_groups - result = [ - BlockTable.map_to_kernel_blocks( - np.array(group), - remote_physical_per_logical, - remote_arange, - ).tolist() - if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) - else group - for i, group in enumerate(block_ids) - ] - return result - def get_backend_aware_kv_block_len( self, layer_idx: int, first_split: bool = True, mamba_view: bool = False ) -> int: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 4414bf60704f..63969382dff8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -48,7 +48,7 @@ def start_load_kv(self, metadata: NixlConnectorMetadata): """ for req_id, meta in metadata.reqs_to_recv.items(): meta.local_physical_block_ids = self._logical_to_kernel_block_ids( - meta.local_block_ids + meta.local_block_ids, self._physical_blocks_per_logical_kv_block ) assert meta.remote is not None # Remote block IDs are kept logical here; expanded in @@ -134,7 +134,7 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): remote_info = self.transfer_topo.get_engine_info(engine_id) tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size) - meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( + meta.remote.block_ids = self._logical_to_kernel_block_ids( meta.remote.block_ids, remote_info.remote_physical_blocks_per_logical, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 090c2074ddd3..859dbae00b6f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -152,7 +152,7 @@ def start_load_kv(self, metadata: NixlConnectorMetadata): # D-side: track reqs waiting for P to push. for req_id, meta in metadata.reqs_to_recv.items(): meta.local_physical_block_ids = self._logical_to_kernel_block_ids( - meta.local_block_ids + meta.local_block_ids, self._physical_blocks_per_logical_kv_block ) assert meta.remote is not None remote_engine_id = meta.remote.engine_id @@ -414,7 +414,9 @@ def _do_start_push_kv( # expands each side using the appropriate ratio. logical_local = self._as_grouped_block_ids(local_block_ids) logical_remote = self._as_grouped_block_ids(remote_block_ids) - physical_local = self._logical_to_kernel_block_ids(logical_local) + physical_local = self._logical_to_kernel_block_ids( + logical_local, self._physical_blocks_per_logical_kv_block + ) push_meta = ReqMeta( local_block_ids=logical_local, @@ -501,7 +503,7 @@ def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta): # Expand D's logical IDs using the ratio learned during the # NIXL handshake. ``meta`` is freshly built by # ``_do_start_push_kv`` so mutating it here is safe. - meta.remote.block_ids = self._logical_to_remote_kernel_block_ids( + meta.remote.block_ids = self._logical_to_kernel_block_ids( meta.remote.block_ids, remote_info.remote_physical_blocks_per_logical, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py new file mode 100644 index 000000000000..5cf4dffd1af6 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Translate vLLM KV cache metadata for native offloading backends.""" + +from typing import TYPE_CHECKING + +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes +from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingGroupConfig, + OffloadingModelConfig, + OffloadingParallelConfig, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheTensor + + +def is_kv_cache_tensor_packed(kv_cache_tensor: "KVCacheTensor") -> bool: + """Return whether a KV cache tensor uses a packed block stride.""" + return bool(kv_cache_tensor.block_stride) + + +def build_offloading_config( + vllm_config: "VllmConfig", + kv_cache_config: "KVCacheConfig", +) -> OffloadingConfig: + """Translate vLLM configuration into the native offloading boundary.""" + kv_transfer_config = vllm_config.kv_transfer_config + assert kv_transfer_config is not None + extra_config = kv_transfer_config.kv_connector_extra_config + assert kv_transfer_config.engine_id is not None + engine_id = kv_transfer_config.engine_id + + parallel_config = vllm_config.parallel_config + context_parallel_factor = ( + parallel_config.decode_context_parallel_size + * parallel_config.prefill_context_parallel_size + ) + groups = tuple( + OffloadingGroupConfig( + tokens_per_block=(group.kv_cache_spec.block_size * context_parallel_factor), + layer_names=tuple(group.layer_names), + ) + for group in kv_cache_config.kv_cache_groups + ) + + _, tokens_per_hash = resolve_kv_cache_block_sizes(kv_cache_config, vllm_config) + for group in groups: + assert group.tokens_per_block % tokens_per_hash == 0, ( + f"tokens_per_block={group.tokens_per_block} not divisible by " + f"tokens_per_hash={tokens_per_hash}. " + f"Hybrid models (e.g. Mamba+Attention) need " + f"--enable-prefix-caching to align block sizes." + ) + + blocks_per_chunk = 1 + blocks_per_chunk_config = extra_config.get("blocks_per_chunk") + tokens_per_chunk = extra_config.get("block_size") + + if blocks_per_chunk_config is not None and tokens_per_chunk is not None: + raise ValueError( + "Specify only one of 'block_size' or 'blocks_per_chunk' " + "in kv_connector_extra_config." + ) + + if blocks_per_chunk_config is not None: + blocks_per_chunk = int(blocks_per_chunk_config) + + if blocks_per_chunk <= 0: + raise ValueError("'blocks_per_chunk' must be greater than 0.") + + elif tokens_per_chunk is not None: + tokens_per_chunk_int = int(tokens_per_chunk) + + unique_tokens_per_block = {group.tokens_per_block for group in groups} + + assert len(unique_tokens_per_block) == 1, ( + "If 'block_size' is specified in kv_connector_extra_config, " + "there must be at least one KV cache group, " + "and all groups must have the same block size." + ) + + tokens_per_block = unique_tokens_per_block.pop() + assert tokens_per_chunk_int % tokens_per_block == 0 + blocks_per_chunk = tokens_per_chunk_int // tokens_per_block + + worker_kv_bytes_per_block = 0 + if kv_cache_config.num_blocks > 0: + packed_tensors = tuple( + is_kv_cache_tensor_packed(tensor) + for tensor in kv_cache_config.kv_cache_tensors + ) + is_packed = any(packed_tensors) + assert not is_packed or all(packed_tensors) + total_gpu_kv_bytes = ( + kv_cache_config.kv_cache_tensors[0].size + if is_packed + else sum(tensor.size for tensor in kv_cache_config.kv_cache_tensors) + ) + worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks + + # Only a single non-MLA full-attention group is parallelism-invariant: + # MLA latent KV is replicated per rank (never head-sharded), and the V2 + # model runner's KV layout is not known to be parallelism-invariant. + single_group = ( + kv_cache_config.kv_cache_groups[0].kv_cache_spec + if len(kv_cache_config.kv_cache_groups) == 1 + else None + ) + is_parallelism_agnostic = ( + not vllm_config.use_v2_model_runner + and single_group is not None + and isinstance(single_group, FullAttentionSpec) + and not isinstance(single_group, MLAAttentionSpec) + ) + + kv_events_config = vllm_config.kv_events_config + return OffloadingConfig( + groups=groups, + worker_kv_bytes_per_block=worker_kv_bytes_per_block, + enable_kv_cache_events=( + kv_events_config is not None and kv_events_config.enable_kv_cache_events + ), + extra_config=extra_config, + engine_id=engine_id, + model=OffloadingModelConfig( + name=vllm_config.model_config.model, + dtype=str(vllm_config.cache_config.cache_dtype).replace("torch.", ""), + ), + cache=OffloadingCacheConfig( + tokens_per_hash=tokens_per_hash, + blocks_per_chunk=blocks_per_chunk, + ), + parallel=OffloadingParallelConfig( + rank=parallel_config.rank, + world_size=parallel_config.world_size, + tp_size=parallel_config.tensor_parallel_size, + pp_size=parallel_config.pipeline_parallel_size, + pcp_size=parallel_config.prefill_context_parallel_size, + dcp_size=parallel_config.decode_context_parallel_size, + data_parallel_index=parallel_config.data_parallel_index, + is_parallelism_agnostic=is_parallelism_agnostic, + ), + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py index 410f84c50ddf..4543cd6dd367 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py @@ -101,7 +101,7 @@ def record_store( self, req: Request, group_config: "GroupOffloadConfig", - offload_block_idx: int, + chunk_idx: int, offload_key: OffloadKey, ) -> None: """Snapshot the KV cache event payload for one offloaded chunk. @@ -111,9 +111,9 @@ def record_store( """ if not self.self_describing_enabled: return - if group_config.sliding_window_size_in_blocks is not None: + if group_config.sliding_window_size_in_chunks is not None: return - meta = self._build_event_metadata(req, group_config, offload_block_idx) + meta = self._build_event_metadata(req, group_config, chunk_idx) self._pending_event_metadata[offload_key] = meta def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]: @@ -142,19 +142,19 @@ def _build_event_metadata( self, req: Request, group_config: "GroupOffloadConfig", - offload_block_idx: int, + chunk_idx: int, ) -> _OffloadEventMetadata: """Build the payload snapshot for one offloaded chunk: its constituent per-block hashes, the whole chunk's tokens, and the per-block ``block_size``.""" - hbf = group_config.hash_block_size_factor + hbf = group_config.hashes_per_chunk assert hbf > 0 - assert offload_block_idx >= 0 + assert chunk_idx >= 0 # per-block token count (= the GPU/hash block size) - sub_block_size = group_config.offloaded_block_size // hbf + tokens_per_hash = group_config.tokens_per_chunk // hbf # chunk c covers hash-blocks [c*hbf, (c+1)*hbf); its tail block's hash # is the chunk's OffloadKey. - first_hash_idx = offload_block_idx * hbf + first_hash_idx = chunk_idx * hbf last_hash_idx = first_hash_idx + hbf assert first_hash_idx >= 0 assert last_hash_idx <= len(req.block_hashes) @@ -164,7 +164,7 @@ def _build_event_metadata( chunk_hashes.append(block_hash) assert len(chunk_hashes) == hbf - if group_config.sliding_window_size_in_blocks is not None: + if group_config.sliding_window_size_in_chunks is not None: # record_store filters these out before calling this helper. raise AssertionError("self-describing events only support full attention") @@ -175,8 +175,8 @@ def _build_event_metadata( parent_block_hash = req.block_hashes[first_hash_idx - 1] assert parent_block_hash is not None - tok_start = offload_block_idx * group_config.offloaded_block_size - tok_end = tok_start + group_config.offloaded_block_size + tok_start = chunk_idx * group_config.tokens_per_chunk + tok_end = tok_start + group_config.tokens_per_chunk assert tok_end <= len(req.all_token_ids) token_ids = tuple(req.all_token_ids[tok_start:tok_end]) @@ -190,7 +190,7 @@ def _build_event_metadata( block_hashes=tuple(chunk_hashes), parent_block_hash=parent_block_hash, token_ids=token_ids, - block_size=sub_block_size, + block_size=tokens_per_hash, lora_id=lora_id, lora_name=lora_name, extra_keys=None, @@ -198,7 +198,12 @@ def _build_event_metadata( kv_cache_spec=group_config.kv_event_group_spec, ) - def _placeholder_stored(self, key: OffloadKey, medium: str) -> BlockStored: + def _placeholder_stored( + self, + key: OffloadKey, + medium: str, + locality: str | None, + ) -> BlockStored: return BlockStored( block_hashes=[ maybe_convert_block_hash(BlockHash(get_offload_block_hash(key))) @@ -210,12 +215,14 @@ def _placeholder_stored(self, key: OffloadKey, medium: str) -> BlockStored: medium=medium, lora_name=None, group_idx=get_offload_group_idx(key), + locality=locality, ) def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: # Metadata is read, NOT popped: the entry must survive until the # eviction event so BlockRemoved can fan out to the same hashes. # Events are self-contained (own parent), so key order is free. + locality = event.locality.value if event.locality is not None else None for key in event.keys: meta = self._pending_event_metadata.get(key) if meta is None: @@ -227,7 +234,7 @@ def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: "placeholder payload. Expected for non-full-attention " "groups; otherwise indicates a missing populate path." ) - yield self._placeholder_stored(key, event.medium) + yield self._placeholder_stored(key, event.medium, locality) continue yield BlockStored( @@ -252,10 +259,12 @@ def _take_stored_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: kv_cache_spec_sliding_window=( meta.kv_cache_spec.kv_cache_spec_sliding_window ), + locality=locality, ) def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: # Keep group_idx unambiguous if a manager batch spans groups. + locality = event.locality.value if event.locality is not None else None by_group: dict[int, list] = {} for key in event.keys: meta = self._pending_event_metadata.pop(key, None) @@ -283,4 +292,5 @@ def _take_removed_event(self, event: OffloadingEvent) -> Iterable[KVCacheEvent]: block_hashes=hashes, medium=event.medium, group_idx=group_idx, + locality=locality, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py index baa168e1708b..58023c393561 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py @@ -321,10 +321,10 @@ def __init__( self.histogram_transfer_size: dict[tuple[int, str], PromMetricT] = {} self.counter_kv_bytes: dict[tuple[int, str], PromMetricT] = {} self.counter_kv_transfer_time: dict[tuple[int, str], PromMetricT] = {} - spec_cls = OffloadingSpecFactory.get_spec_cls(vllm_config) kv_transfer_config = vllm_config.kv_transfer_config assert kv_transfer_config is not None extra_config = kv_transfer_config.kv_connector_extra_config + spec_cls = OffloadingSpecFactory.get_spec_cls(extra_config) self._offloading_metric_metadata: dict[str, OffloadingMetricMetadata] = { **spec_cls.build_metric_definitions(extra_config), **get_connector_metric_definitions(), 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 f896c9cc4923..5e98c1266e20 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -3,9 +3,10 @@ import time from collections.abc import Iterable, Sequence from dataclasses import dataclass, field -from itertools import islice +from itertools import chain, islice from typing import Any, NamedTuple +from vllm.config import VllmConfig from vllm.distributed.kv_events import KVCacheEvent from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata @@ -31,6 +32,7 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( FullAttentionSpec, + KVCacheConfig, KVCacheSpec, MambaSpec, SlidingWindowSpec, @@ -74,32 +76,32 @@ class TransferJobStatus: class GroupOffloadConfig(NamedTuple): group_idx: int - gpu_block_size: int - offloaded_block_size: int - hash_block_size_factor: int + tokens_per_block: int + tokens_per_chunk: int + hashes_per_chunk: int # KV cache spec metadata propagated onto emitted BlockStored events so # KV-aware consumers can classify and filter the group. kv_event_group_spec: OffloadingEventGroupSpec # None below means full attention - sliding_window_size_in_blocks: int | None - # Number of this group's offloaded blocks per full-attention alignment - # segment. Used to skip storing SWA blocks that can never serve a load + sliding_window_size_in_chunks: int | None + # Number of this group's offloaded chunks per full-attention alignment + # segment. Used to skip storing SWA chunks that can never serve a load # hit (e.g. DeepSeek V4 where SWA groups have much smaller block sizes # than the MLA full-attention group). # None for full-attention groups or when the optimization doesn't apply. - alignment_block_count: int | None = None - # True for EAGLE/MTP draft-model attention groups. The trailing block + alignment_chunk_count: int | None = None + # True for EAGLE/MTP draft-model attention groups. The trailing chunk # of these groups is volatile and lacks a stable hash, so it must # be excluded from store and load scheduling. is_eagle_group: bool = False -def get_sliding_window_size_in_blocks( - kv_cache_spec: KVCacheSpec, offloaded_block_size: int +def get_sliding_window_size_in_chunks( + kv_cache_spec: KVCacheSpec, tokens_per_chunk: int ) -> int | None: if isinstance(kv_cache_spec, SlidingWindowSpec): assert kv_cache_spec.sliding_window > 0 - return cdiv(kv_cache_spec.sliding_window, offloaded_block_size) + return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk) if isinstance(kv_cache_spec, MambaSpec): # Mamba depends on a single state @@ -109,116 +111,121 @@ def get_sliding_window_size_in_blocks( return None -def resolve_mamba_align_size(spec: "OffloadingSpec") -> int | None: +def resolve_mamba_align_size( + spec: "OffloadingSpec", kv_cache_config: KVCacheConfig +) -> int | None: """Scan all KV cache groups in *spec* and return the single mamba alignment size, or None if no group requires mamba alignment. For MambaSpec groups in "align" cache mode the hit window must be rounded - down to a multiple of the offloaded block size. Asserts that all such + down to a multiple of the offloaded chunk size. Asserts that all such groups agree on the same value. """ mamba_align_size: int | None = None - for idx, gpu_block_size in enumerate(spec.gpu_block_size): - kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec + for idx, tokens_per_block in enumerate(spec.tokens_per_block): + kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align": - offload_block_size = gpu_block_size * spec.block_size_factor - assert mamba_align_size is None or mamba_align_size == offload_block_size - mamba_align_size = offload_block_size + tokens_per_chunk = tokens_per_block * spec.blocks_per_chunk + assert mamba_align_size is None or mamba_align_size == tokens_per_chunk + mamba_align_size = tokens_per_chunk return mamba_align_size class SchedulerOffloadConfig(NamedTuple): kv_group_configs: tuple[GroupOffloadConfig, ...] - block_size_factor: int + blocks_per_chunk: int num_workers: int offload_prompt_only: bool @classmethod - def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig": + def from_spec( + cls, + spec: OffloadingSpec, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + ) -> "SchedulerOffloadConfig": # Determine the alignment token count from the full-attention group(s). - # This is the offloaded_block_size of the full-attention group; load + # This is the tokens_per_chunk of the full-attention group; load # hits are always aligned to this boundary, so SWA blocks earlier in # each segment can never serve a load hit. Relevant for hybrid # architectures like DeepSeek V4 (MLA + SWA groups). - full_attn_offloaded_block_sizes: set[int] = set() - for idx, gpu_block_size in enumerate(spec.gpu_block_size): - kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec - sw = get_sliding_window_size_in_blocks( - kv_spec, gpu_block_size * spec.block_size_factor + full_attn_tokens_per_chunk: set[int] = set() + for idx, tokens_per_block in enumerate(spec.tokens_per_block): + kv_spec = kv_cache_config.kv_cache_groups[idx].kv_cache_spec + sw = get_sliding_window_size_in_chunks( + kv_spec, tokens_per_block * spec.blocks_per_chunk ) if sw is None: - full_attn_offloaded_block_sizes.add( - gpu_block_size * spec.block_size_factor - ) + full_attn_tokens_per_chunk.add(tokens_per_block * spec.blocks_per_chunk) # Only apply the optimization if there's a single consistent # full-attention alignment size. alignment_tokens: int | None = None - if len(full_attn_offloaded_block_sizes) == 1: - alignment_tokens = full_attn_offloaded_block_sizes.pop() + if len(full_attn_tokens_per_chunk) == 1: + alignment_tokens = full_attn_tokens_per_chunk.pop() - def _alignment_block_count( - offloaded_block_size: int, - sliding_window_size_in_blocks: int | None, + def _alignment_chunk_count( + tokens_per_chunk: int, + sliding_window_size_in_chunks: int | None, ) -> int | None: - if alignment_tokens is None or sliding_window_size_in_blocks is None: + if alignment_tokens is None or sliding_window_size_in_chunks is None: return None - if alignment_tokens <= offloaded_block_size: + if alignment_tokens <= tokens_per_chunk: return None - per_segment = alignment_tokens // offloaded_block_size - if sliding_window_size_in_blocks >= per_segment: + per_segment = alignment_tokens // tokens_per_chunk + if sliding_window_size_in_chunks >= per_segment: return None return per_segment eagle_groups = { idx - for idx, g in enumerate(spec.kv_cache_config.kv_cache_groups) + for idx, g in enumerate(kv_cache_config.kv_cache_groups) if g.is_eagle_group } use_eagle = ( - spec.vllm_config.speculative_config is not None - and spec.vllm_config.speculative_config.use_eagle() + vllm_config.speculative_config is not None + and vllm_config.speculative_config.use_eagle() ) if use_eagle and not eagle_groups: - eagle_groups = set(range(len(spec.kv_cache_config.kv_cache_groups))) + eagle_groups = set(range(len(kv_cache_config.kv_cache_groups))) if eagle_groups: logger.info( "KV offloading: EAGLE/MTP draft attention groups %s " - "detected. The trailing block of these groups will be " + "detected. The trailing chunk of these groups will be " "excluded from offloading due to volatility.", sorted(eagle_groups), ) return cls( - num_workers=spec.vllm_config.parallel_config.world_size, + num_workers=vllm_config.parallel_config.world_size, kv_group_configs=tuple( GroupOffloadConfig( group_idx=idx, - gpu_block_size=gpu_block_size, - offloaded_block_size=gpu_block_size * spec.block_size_factor, - hash_block_size_factor=( - (gpu_block_size * spec.block_size_factor) - // spec.hash_block_size + tokens_per_block=tokens_per_block, + tokens_per_chunk=tokens_per_block * spec.blocks_per_chunk, + hashes_per_chunk=( + (tokens_per_block * spec.blocks_per_chunk) + // spec.tokens_per_hash ), - sliding_window_size_in_blocks=( - sw := get_sliding_window_size_in_blocks( - spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec, - gpu_block_size * spec.block_size_factor, + sliding_window_size_in_chunks=( + sw := get_sliding_window_size_in_chunks( + kv_cache_config.kv_cache_groups[idx].kv_cache_spec, + tokens_per_block * spec.blocks_per_chunk, ) ), - alignment_block_count=_alignment_block_count( - gpu_block_size * spec.block_size_factor, sw + alignment_chunk_count=_alignment_chunk_count( + tokens_per_block * spec.blocks_per_chunk, sw ), kv_event_group_spec=get_offloading_event_group_spec( - spec.kv_cache_config.kv_cache_groups[idx] + kv_cache_config.kv_cache_groups[idx] ), is_eagle_group=idx in eagle_groups, ) - for idx, gpu_block_size in enumerate(spec.gpu_block_size) + for idx, tokens_per_block in enumerate(spec.tokens_per_block) ), - block_size_factor=spec.block_size_factor, + blocks_per_chunk=spec.blocks_per_chunk, offload_prompt_only=spec.offload_prompt_only, ) @@ -227,11 +234,11 @@ def _alignment_block_count( class RequestGroupState: offload_keys: list[OffloadKey] = field(default_factory=list) block_ids: list[int] = field(default_factory=list) - # index of next block (of size offloaded_block_size) to offload - next_stored_block_idx: int = 0 - # number of offloaded blocks hit (including GPU prefix cache) + # Index of the next chunk to offload. + next_stored_chunk_idx: int = 0 + # Number of offloaded chunks hit (including GPU prefix cache) # when the request first started - num_hit_blocks: int = 0 + num_hit_chunks: int = 0 @dataclass(slots=True) @@ -278,11 +285,11 @@ def update_offload_keys(self) -> None: ): for req_block_hash in islice( self.req.block_hashes, - group_config.hash_block_size_factor * len(group_state.offload_keys) - + group_config.hash_block_size_factor + group_config.hashes_per_chunk * len(group_state.offload_keys) + + group_config.hashes_per_chunk - 1, None, - group_config.hash_block_size_factor, + group_config.hashes_per_chunk, ): group_state.offload_keys.append( make_offload_key(req_block_hash, group_config.group_idx) @@ -298,46 +305,46 @@ def update_block_id_groups( for group_state, new_blocks in zip(self.group_states, new_block_id_groups): group_state.block_ids.extend(new_blocks) - def storable_blocks( + def storable_chunks( self, group_config: "GroupOffloadConfig", num_offloadable_tokens: int ) -> int: - """Number of leading offloaded blocks eligible for store. + """Number of leading offloaded chunks eligible for store. - For eagle/MTP groups the volatile trailing block of the offloadable + For eagle/MTP groups the volatile trailing chunk of the offloadable range is excluded while decoding: the draft-layer KV of the last accepted position may be rewritten after spec-token rejection. During - prefill the trailing block is stable (the draft input for a chunk's + prefill the trailing chunk is stable (the draft input for a chunk's last position is the next prompt token), so it is stored immediately. The exclusion must be applied consistently everywhere - ``next_stored_block_idx`` is derived: otherwise the trailing block of + ``next_stored_chunk_idx`` is derived: otherwise the trailing chunk of each step is skipped on collection but jumped over by - ``next_stored_block_idx``, so it is never re-considered and a + ``next_stored_chunk_idx``, so it is never re-considered and a permanent hole breaks prefix-reuse lookup. """ - num_blocks = num_offloadable_tokens // group_config.offloaded_block_size + num_chunks = num_offloadable_tokens // group_config.tokens_per_chunk is_decoding = num_offloadable_tokens > self.req.num_prompt_tokens if group_config.is_eagle_group and is_decoding: - num_blocks = max(0, num_blocks - 1) - return num_blocks + num_chunks = max(0, num_chunks - 1) + return num_chunks def advance_stored_idx(self, num_offloadable_tokens: int) -> None: - # max(): at the prefill->decode transition of a block-aligned prompt, - # storable_blocks drops by one (the eagle exclusion kicks in), and the - # index must not move backwards past already-stored blocks. + # max(): at the prefill->decode transition of a chunk-aligned prompt, + # storable_chunks drops by one (the eagle exclusion kicks in), and the + # index must not move backwards past already-stored chunks. for group_config, group_state in zip( self.config.kv_group_configs, self.group_states ): - group_state.next_stored_block_idx = max( - group_state.next_stored_block_idx, - self.storable_blocks(group_config, num_offloadable_tokens), + group_state.next_stored_chunk_idx = max( + group_state.next_stored_chunk_idx, + self.storable_chunks(group_config, num_offloadable_tokens), ) - def update_num_hit_blocks(self, num_cached_tokens: int) -> None: + def update_num_hit_chunks(self, num_cached_tokens: int) -> None: for group_config, group_state in zip( self.config.kv_group_configs, self.group_states ): - group_state.num_hit_blocks = ( - num_cached_tokens // group_config.offloaded_block_size + group_state.num_hit_chunks = ( + num_cached_tokens // group_config.tokens_per_chunk ) @@ -354,22 +361,26 @@ class OffloadingConnectorScheduler: def __init__( self, spec: OffloadingSpec, + vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, ): - self.config = SchedulerOffloadConfig.from_spec(spec) + self.config = SchedulerOffloadConfig.from_spec( + spec, vllm_config, kv_cache_config + ) self.manager: OffloadingManager = spec.get_manager() self._connector_stats = OffloadingConnectorStats() full_attention_groups: list[int] = [] sliding_window_groups: list[int] = [] for group_config in self.config.kv_group_configs: - if group_config.sliding_window_size_in_blocks is None: + if group_config.sliding_window_size_in_chunks is None: full_attention_groups.append(group_config.group_idx) else: sliding_window_groups.append(group_config.group_idx) # sort sliding window groups by window size in decreasing order def _sliding_window_sort_key(i: int) -> int: - val = self.config.kv_group_configs[i].sliding_window_size_in_blocks + val = self.config.kv_group_configs[i].sliding_window_size_in_chunks assert val is not None return val @@ -378,7 +389,9 @@ def _sliding_window_sort_key(i: int) -> int: # used by _lookup self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups) self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups - self._mamba_align_size: int | None = resolve_mamba_align_size(spec) + self._mamba_align_size: int | None = resolve_mamba_align_size( + spec, kv_cache_config + ) self._req_status: dict[ReqId, RequestOffloadState] = {} self._current_batch_load_jobs: dict[int, TransferJob] = {} @@ -386,9 +399,9 @@ def _sliding_window_sort_key(i: int) -> int: # GPU block IDs allocated in the current engine step self._current_batch_allocated_block_ids: set[int] = set() # if GPU prefix caching is enabled, - # track loaded blocks to avoid redundant loads - self._blocks_being_loaded: set[OffloadKey] | None = ( - set() if spec.vllm_config.cache_config.enable_prefix_caching else None + # Track loaded chunks to avoid redundant loads. + self._chunks_being_loaded: set[OffloadKey] | None = ( + set() if vllm_config.cache_config.enable_prefix_caching else None ) # Job ID counter shared by loads and stores. @@ -431,10 +444,28 @@ def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None: if not pending: del self._block_id_to_pending_jobs[bid] + def _calc_num_offloadable_tokens( + self, req_status: RequestOffloadState, num_computed_tokens: int + ) -> int: + num = min(num_computed_tokens, req_status.req.num_tokens) + max_offload_tokens = req_status.max_offload_tokens + if max_offload_tokens is not None: + num = min(num, max_offload_tokens) + if self.config.offload_prompt_only: + num = min(num, req_status.req.num_prompt_tokens) + return num + + def _maybe_cleanup_finished_req( + self, req_id: str, req_status: RequestOffloadState + ) -> None: + """Clean up req_status if finished and no in-flight jobs.""" + if req_status.req.is_finished() and not req_status.transfer_jobs: + del self._req_status[req_id] + def _maximal_prefix_lookup( self, keys: Iterable[OffloadKey], req_context: ReqContext ) -> int | None: - """Return the number of consecutive offloaded blocks from the start, + """Return the number of consecutive offloaded chunks from the start, or None if the backend deferred a lookup.""" hit_count = 0 defer_lookup = False @@ -490,18 +521,18 @@ def _touch(self, req_status: RequestOffloadState): for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - if group_config.sliding_window_size_in_blocks is None: + if group_config.sliding_window_size_in_chunks is None: self.manager.touch(group_state.offload_keys, req_status.req_context) else: - # we aim to keep just blocks that are necessary to hit - # the original request (+ decoded blocks) - blocks_to_skip = max( + # Keep only chunks needed to hit the original request, plus + # decoded chunks. + chunks_to_skip = max( 0, - group_state.num_hit_blocks - - group_config.sliding_window_size_in_blocks, + group_state.num_hit_chunks + - group_config.sliding_window_size_in_chunks, ) self.manager.touch( - group_state.offload_keys[blocks_to_skip:], + group_state.offload_keys[chunks_to_skip:], req_status.req_context, ) @@ -531,7 +562,7 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: defer_lookup = False lookup_groups = self._lookup_groups - # Tracks which eagle groups have already popped their volatile trailing block + # Tracks which eagle groups have already popped their volatile trailing chunk # in the current convergence iteration. Reset when a non-eagle group # tightens the hit boundary, requiring a fresh pop. eagle_verified: set[int] = set() @@ -544,79 +575,76 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: group_idx ] group_state: RequestGroupState = req_status.group_states[group_idx] - offloaded_block_size = group_config.offloaded_block_size + tokens_per_chunk = group_config.tokens_per_chunk offload_keys = group_state.offload_keys assert ( - len(offload_keys) - >= req_status.req.num_tokens // offloaded_block_size + len(offload_keys) >= req_status.req.num_tokens // tokens_per_chunk ) is_eagle_unverified = ( group_config.is_eagle_group and group_idx not in eagle_verified ) - # Constrain to block-aligned boundary for this group + # Constrain to a chunk-aligned boundary for this group. max_hit_size_tokens = min( - max_hit_size_tokens, len(offload_keys) * offloaded_block_size + max_hit_size_tokens, len(offload_keys) * tokens_per_chunk ) - if max_hit_size_tokens - num_computed_tokens < offloaded_block_size: - # we can only load less than a block, better skip + if max_hit_size_tokens - num_computed_tokens < tokens_per_chunk: + # We can only load less than a chunk, so skip. return 0 - sliding_window_size_in_blocks = ( - group_config.sliding_window_size_in_blocks + sliding_window_size_in_chunks = ( + group_config.sliding_window_size_in_chunks ) - # For eagle groups, query one extra block that will be popped. + # For eagle groups, query one extra chunk that will be popped. # We only need to increase the query size for sliding window groups. query_max = max_hit_size_tokens - if is_eagle_unverified and sliding_window_size_in_blocks is not None: + if is_eagle_unverified and sliding_window_size_in_chunks is not None: query_max = min( - max_hit_size_tokens + offloaded_block_size, - len(offload_keys) * offloaded_block_size, + max_hit_size_tokens + tokens_per_chunk, + len(offload_keys) * tokens_per_chunk, ) - num_blocks = min( - cdiv(query_max, offloaded_block_size), len(offload_keys) - ) - start_block_idx = num_computed_tokens // offloaded_block_size - offload_keys = offload_keys[start_block_idx:num_blocks] + num_chunks = min(cdiv(query_max, tokens_per_chunk), len(offload_keys)) + start_chunk_idx = num_computed_tokens // tokens_per_chunk + offload_keys = offload_keys[start_chunk_idx:num_chunks] # end index (in the sliced offload_keys) up to which we # have backend-confirmed hits - num_hit_blocks: int | None - if sliding_window_size_in_blocks is None: - num_hit_blocks = self._maximal_prefix_lookup( + num_hit_chunks: int | None + if sliding_window_size_in_chunks is None: + num_hit_chunks = self._maximal_prefix_lookup( offload_keys, req_status.req_context ) else: - required_window = sliding_window_size_in_blocks + required_window = sliding_window_size_in_chunks if is_eagle_unverified: required_window += 1 - num_hit_blocks = self._sliding_window_lookup( + num_hit_chunks = self._sliding_window_lookup( offload_keys, required_window, req_status.req_context, ) - if num_hit_blocks == 0: + if num_hit_chunks == 0: return 0 - if num_hit_blocks is None: + if num_hit_chunks is None: defer_lookup = True else: if is_eagle_unverified: - num_hit_blocks -= 1 + num_hit_chunks -= 1 eagle_verified.add(group_idx) max_hit_size_tokens = min( max_hit_size_tokens, - offloaded_block_size * (start_block_idx + num_hit_blocks), + tokens_per_chunk * (start_chunk_idx + num_hit_chunks), ) new_num_hit_tokens = max_hit_size_tokens - num_computed_tokens - if new_num_hit_tokens < offloaded_block_size: - # we can only load less than a block, better skip + if new_num_hit_tokens < tokens_per_chunk: + # We can only load less than a chunk, so skip. return 0 if new_num_hit_tokens < num_hit_tokens: @@ -632,7 +660,7 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: # sliding window works with the new_num_hit_tokens lookup_groups = self._sliding_window_groups - looked_up_sliding_window |= sliding_window_size_in_blocks is not None + looked_up_sliding_window |= sliding_window_size_in_chunks is not None num_hit_tokens = new_num_hit_tokens if defer_lookup: @@ -642,28 +670,28 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: ) return None - # possibly delay request if any of the hit blocks is already being loaded - if self._blocks_being_loaded: + # Possibly delay the request if any hit chunk is already being loaded. + if self._chunks_being_loaded: for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - offloaded_block_size = group_config.offloaded_block_size - sliding_window_size_in_blocks = ( - group_config.sliding_window_size_in_blocks + tokens_per_chunk = group_config.tokens_per_chunk + sliding_window_size_in_chunks = ( + group_config.sliding_window_size_in_chunks ) offload_keys = group_state.offload_keys - num_blocks = cdiv( - num_computed_tokens + num_hit_tokens, offloaded_block_size + num_chunks = cdiv( + num_computed_tokens + num_hit_tokens, tokens_per_chunk ) - start_block_idx = num_computed_tokens // offloaded_block_size - offload_keys = offload_keys[start_block_idx:num_blocks] - if sliding_window_size_in_blocks is not None: - offload_keys = offload_keys[-sliding_window_size_in_blocks:] - if any(key in self._blocks_being_loaded for key in offload_keys): - # hit blocks are being loaded, delay request + start_chunk_idx = num_computed_tokens // tokens_per_chunk + offload_keys = offload_keys[start_chunk_idx:num_chunks] + if sliding_window_size_in_chunks is not None: + offload_keys = offload_keys[-sliding_window_size_in_chunks:] + if any(key in self._chunks_being_loaded for key in offload_keys): + # Hit chunks are being loaded, so delay the request. logger.debug( "Delaying request %s since some of its" - " blocks are already being loaded", + " chunks are already being loaded", req_status.req.request_id, ) return None @@ -740,7 +768,7 @@ def get_num_new_matched_tokens( req_status.deferred_lookup_start_time = lookup_start else: self._maybe_observe_lookup_async_delay(req_status) - req_status.update_num_hit_blocks(num_computed_tokens + (num_hit_tokens or 0)) + req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0)) self._touch(req_status) @@ -771,10 +799,10 @@ def update_state_after_alloc( block.block_id for block in group_blocks if block.block_id != 0 ) - gpu_block_size = group_config.gpu_block_size - offloaded_block_size = group_config.offloaded_block_size + tokens_per_block = group_config.tokens_per_block + tokens_per_chunk = group_config.tokens_per_chunk offload_keys = group_state.offload_keys - num_gpu_blocks = cdiv(num_cached_tokens, gpu_block_size) + num_gpu_blocks = cdiv(num_cached_tokens, tokens_per_block) assert len(group_blocks) >= num_gpu_blocks num_locally_computed_gpu_blocks = num_gpu_blocks @@ -786,24 +814,24 @@ def update_state_after_alloc( assert ( num_locally_computed_tokens - <= num_locally_computed_gpu_blocks * gpu_block_size + <= num_locally_computed_gpu_blocks * tokens_per_block ) num_pending_gpu_blocks = num_gpu_blocks - num_locally_computed_gpu_blocks - if group_config.sliding_window_size_in_blocks is not None: + if group_config.sliding_window_size_in_chunks is not None: assert ( num_pending_gpu_blocks - <= group_config.sliding_window_size_in_blocks - * self.config.block_size_factor + <= group_config.sliding_window_size_in_chunks + * self.config.blocks_per_chunk ) - num_blocks = cdiv(num_cached_tokens, offloaded_block_size) - assert len(offload_keys) >= num_blocks + num_chunks = cdiv(num_cached_tokens, tokens_per_chunk) + assert len(offload_keys) >= num_chunks if num_pending_gpu_blocks: - start_block_idx = ( - num_locally_computed_gpu_blocks // self.config.block_size_factor + start_chunk_idx = ( + num_locally_computed_gpu_blocks // self.config.blocks_per_chunk ) - keys_to_load.extend(offload_keys[start_block_idx:num_blocks]) + keys_to_load.extend(offload_keys[start_chunk_idx:num_chunks]) dst_block_ids.extend( block.block_id @@ -814,11 +842,11 @@ def update_state_after_alloc( group_sizes.append(num_pending_gpu_blocks) block_indices.append(num_locally_computed_gpu_blocks) - # Skip prefix-hit blocks for block-level policy; for - # request-level, next_stored_block_idx stays at 0 so all - # blocks (including hits) are offloaded. + # Skip prefix-hit chunks for block-level policy; for + # request-level, next_stored_chunk_idx stays at 0 so all + # chunks (including hits) are offloaded. if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL: - group_state.next_stored_block_idx = num_blocks + group_state.next_stored_chunk_idx = num_chunks src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context) dst_spec = GPULoadStoreSpec( @@ -841,8 +869,8 @@ def update_state_after_alloc( is_store=False, ) - if self._blocks_being_loaded is not None: - self._blocks_being_loaded.update(keys_to_load) + if self._chunks_being_loaded is not None: + self._chunks_being_loaded.update(keys_to_load) def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: """ @@ -877,16 +905,16 @@ def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: # Zero out stale block_ids in sliding window groups' pending-store # positions. Only sliding window groups can have stale entries (blocks # freed by remove_skipped_blocks then reallocated). Only positions in - # [next_stored_block_idx * bsf, end) need checking where end is the + # [next_stored_chunk_idx * bsf, end) need checking where end is the # pre-extend length: earlier positions were already offloaded, later # ones are fresh allocations from this step. if self._sliding_window_groups and self._current_batch_allocated_block_ids: - block_size_factor = self.config.block_size_factor + blocks_per_chunk = self.config.blocks_per_chunk for req_id, req_status in self._req_status.items(): ends = new_block_ids_end.get(req_id) for i, grp_idx in enumerate(self._sliding_window_groups): group_state = req_status.group_states[grp_idx] - start = group_state.next_stored_block_idx * block_size_factor + start = group_state.next_stored_chunk_idx * blocks_per_chunk end = ends[i] if ends is not None else len(group_state.block_ids) for j in range(start, end): if ( @@ -899,80 +927,77 @@ def _build_store_jobs( self, scheduler_output: SchedulerOutput, ) -> dict[int, TransferJob]: - block_size_factor = self.config.block_size_factor + blocks_per_chunk = self.config.blocks_per_chunk store_jobs: dict[int, TransferJob] = {} - for req_id in scheduler_output.num_scheduled_tokens: + for req_id in chain( + scheduler_output.num_scheduled_tokens, + scheduler_output.finished_req_ids or (), + ): req_status = self._req_status.get(req_id) if req_status is None: continue req = req_status.req - num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] - num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens - # with async scheduling, some tokens may be missing - num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens) - max_offload_tokens = req_status.max_offload_tokens - if max_offload_tokens is not None: - num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens) - - # Skip decode-phase blocks: clamp to the prompt length so only - # prefill (prompt) blocks become eligible for store. next_stored_idx - # never advances past this boundary, so decode blocks are never - # queued in this or any later step. - if self.config.offload_prompt_only: - num_offloadable_tokens = min( - num_offloadable_tokens, req.num_prompt_tokens - ) + if req.is_finished(): + num_tokens_after_batch = req.num_tokens + else: + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens - # Filter out blocks skipped due to sliding window attention / SSM + num_offloadable_tokens = self._calc_num_offloadable_tokens( + req_status, num_tokens_after_batch + ) + + # Filter out chunks skipped due to sliding window attention / SSM # or unreachable by the load path's alignment constraints. new_offload_keys: list[OffloadKey] = [] for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - num_blocks = req_status.storable_blocks( + num_chunks = req_status.storable_chunks( group_config, num_offloadable_tokens ) - start_block_idx = group_state.next_stored_block_idx - if num_blocks <= start_block_idx: + start_chunk_idx = group_state.next_stored_chunk_idx + if num_chunks <= start_chunk_idx: continue - offload_keys = group_state.offload_keys[start_block_idx:num_blocks] - # For each block to offload, take the last corresponding GPU block. - # e.g. if block size factor is 3 and GPU block IDs are - # 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8. + offload_keys = group_state.offload_keys[start_chunk_idx:num_chunks] + # For each chunk, take the last corresponding GPU block. For + # blocks_per_chunk=3 and GPU block IDs 1 5 6 7 2 4 9 3 8, + # this selects GPU blocks 6 4 8. # A block_id of 0 means either a sliding window / SSM skip # or a stale entry that was zeroed out — skip it either way. offload_block_ids = group_state.block_ids[ - start_block_idx * block_size_factor - + block_size_factor - - 1 : num_blocks * block_size_factor : block_size_factor + start_chunk_idx * blocks_per_chunk + + blocks_per_chunk + - 1 : num_chunks * blocks_per_chunk : blocks_per_chunk ] assert len(offload_keys) == len(offload_block_ids) - alignment_block_count = group_config.alignment_block_count - tail = group_config.sliding_window_size_in_blocks + alignment_chunk_count = group_config.alignment_chunk_count + tail = group_config.sliding_window_size_in_chunks for key_idx, (offload_key, block_id) in enumerate( zip(offload_keys, offload_block_ids) ): if block_id == 0: continue - # Skip SWA blocks that can never serve a load hit: + # Skip SWA chunks that can never serve a load hit: # within each full-attention alignment segment, only the - # trailing `tail` blocks are reachable by + # trailing `tail` chunks are reachable by # _sliding_window_lookup. For DeepSeek V4 with 100K # tokens this reduces SWA stores by ~78%. - if alignment_block_count is not None: + if alignment_chunk_count is not None: assert tail is not None - abs_block_idx = start_block_idx + key_idx - pos_in_segment = abs_block_idx % alignment_block_count - if pos_in_segment < alignment_block_count - tail: + abs_chunk_idx = start_chunk_idx + key_idx + pos_in_segment = abs_chunk_idx % alignment_chunk_count + if pos_in_segment < alignment_chunk_count - tail: continue new_offload_keys.append(offload_key) if not new_offload_keys: req_status.advance_stored_idx(num_offloadable_tokens) + self._maybe_cleanup_finished_req(req_id, req_status) continue store_output = self.manager.prepare_store( @@ -982,11 +1007,13 @@ def _build_store_jobs( self._connector_stats.increase_counter( _ConnectorMetricName.ALLOCATION_FAILURE ) - logger.warning("Request %s: cannot store blocks", req_id) + logger.warning("Request %s: cannot store chunks", req_id) + self._maybe_cleanup_finished_req(req_id, req_status) continue if not store_output.keys_to_store: req_status.advance_stored_idx(num_offloadable_tokens) + self._maybe_cleanup_finished_req(req_id, req_status) continue self._touch(req_status) @@ -1002,29 +1029,29 @@ def _build_store_jobs( self.config.kv_group_configs, req_status.group_states ): is_sliding_window = ( - group_config.sliding_window_size_in_blocks is not None + group_config.sliding_window_size_in_chunks is not None ) - num_blocks = req_status.storable_blocks( + num_chunks = req_status.storable_chunks( group_config, num_offloadable_tokens ) - start_block_idx = group_state.next_stored_block_idx + start_chunk_idx = group_state.next_stored_chunk_idx block_ids = group_state.block_ids num_group_blocks = 0 start_gpu_block_idx: int | None = None for idx, offload_key in enumerate( - group_state.offload_keys[start_block_idx:num_blocks] + group_state.offload_keys[start_chunk_idx:num_chunks] ): if offload_key not in keys_to_store: continue - offloaded_block_idx = start_block_idx + idx + chunk_idx = start_chunk_idx + idx self._events_tracker.record_store( - req, group_config, offloaded_block_idx, offload_key + req, group_config, chunk_idx, offload_key ) - gpu_block_idx = offloaded_block_idx * block_size_factor - for i in range(block_size_factor): + gpu_block_idx = chunk_idx * blocks_per_chunk + for i in range(blocks_per_chunk): block_id = block_ids[gpu_block_idx + i] if block_id == 0: continue @@ -1039,8 +1066,8 @@ def _build_store_jobs( group_sizes.append(num_group_blocks) block_indices.append(start_gpu_block_idx or 0) - group_state.next_stored_block_idx = max( - group_state.next_stored_block_idx, num_blocks + group_state.next_stored_chunk_idx = max( + group_state.next_stored_chunk_idx, num_chunks ) src_spec = GPULoadStoreSpec( @@ -1076,13 +1103,20 @@ def _build_store_jobs( ) logger.debug( - "Request %s offloading %s blocks upto %d tokens (job %d)", + "Request %s offloading %s chunks upto %d tokens (job %d)", req_id, len(keys_to_store), num_offloadable_tokens, job_id, ) + if req.is_finished(): + # Register non-sliding-window blocks for flush detection. + for bid in non_sliding_window_block_ids: + self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id) + if bid in self._current_batch_allocated_block_ids: + self._current_batch_jobs_to_flush.add(job_id) + return store_jobs def build_connector_meta( @@ -1198,8 +1232,8 @@ def update_connector_output(self, connector_output: KVConnectorOutput): self.manager.complete_store(job_status.keys, req_status.req_context) else: self.manager.complete_load(job_status.keys, req_status.req_context) - if self._blocks_being_loaded: - self._blocks_being_loaded.difference_update(job_status.keys) + if self._chunks_being_loaded: + self._chunks_being_loaded.difference_update(job_status.keys) if self._block_id_to_pending_jobs: # Sliding window blocks are tracked from store creation # and must be cleaned up unconditionally. @@ -1245,8 +1279,6 @@ def request_finished( Optional KVTransferParams to be included in the request outputs returned by the engine. """ - # TODO(orozery): possibly kickoff offload for last block - # which may have been deferred due to async scheduling req_status = self._req_status.get(request.request_id) if req_status is None: @@ -1259,20 +1291,20 @@ def request_finished( self.manager.on_request_finished(req_status.req_context) self._maybe_observe_lookup_async_delay(req_status) - if not req_status.transfer_jobs: - # No in-flight jobs: no later complete_store()/complete_load() calls - # need this request's state. - del self._req_status[request.request_id] - return False, None - # In-flight jobs remain after the request stopped. Their completion may - # still call manager.complete_store()/complete_load(), so keep req_status. - # Pending stores outlive the request's block ownership; register them so - # future reuse of those blocks triggers a flush. + # Update offload keys with final block hash so _build_store_jobs can + # create store jobs for the last block(s) on the next schedule step. + req_status.update_offload_keys() + + # Keep req_status alive: _build_store_jobs will process finished_req_ids + # on the next step and handle cleanup after creating store jobs. + # Register non_sliding_window_block_ids so future block reuse triggers + # a flush via _block_id_to_pending_jobs. for job_id in req_status.transfer_jobs: job_status = self._jobs[job_id] for bid in job_status.non_sliding_window_block_ids or (): self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id) + return False, None def take_events(self) -> Iterable[KVCacheEvent]: @@ -1289,7 +1321,7 @@ def take_events(self) -> Iterable[KVCacheEvent]: yield from self._events_tracker.take_events(self.manager.take_events()) def reset_cache(self) -> None: - """Reset the offloading manager cache, evicting all stored blocks.""" + """Reset the offloading manager cache, evicting all stored chunks.""" # reset_cache cannot be called in the middle of a schedule step assert not self._current_batch_load_jobs @@ -1306,10 +1338,10 @@ def reset_cache(self) -> None: # Reset offloading manager cache self.manager.reset_cache() - # Reset store progress so active requests re-offload from block 0 + # Reset store progress so active requests re-offload from chunk 0. for status in self._req_status.values(): for group_state in status.group_states: - group_state.next_stored_block_idx = 0 + group_state.next_stored_chunk_idx = 0 status.transfer_jobs.clear() # Discard jobs and save job_counter to be able to discard worker responses @@ -1323,8 +1355,8 @@ def reset_cache(self) -> None: # Note: _current_batch_jobs_to_flush is intentionally NOT cleared. # The load flush IDs collected above must be delivered to workers. - if self._blocks_being_loaded is not None: - self._blocks_being_loaded.clear() + if self._chunks_being_loaded is not None: + self._chunks_being_loaded.clear() def shutdown(self) -> None: self.manager.shutdown() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 60e0beb6845c..045e513c3f49 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -10,10 +10,14 @@ OffloadingWorkerMetadata, ReqId, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( + is_kv_cache_tensor_packed, +) from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionBackend from vllm.v1.kv_cache_interface import ( AttentionSpec, + KVCacheConfig, MambaSpec, UniformTypeKVCacheSpecs, ) @@ -33,8 +37,13 @@ class OffloadingConnectorWorker: """Implementation of Worker side methods""" - def __init__(self, spec: OffloadingSpec): + def __init__( + self, + spec: OffloadingSpec, + kv_cache_config: KVCacheConfig, + ): self.spec = spec + self.kv_cache_config = kv_cache_config self.worker: OffloadingWorker | None = None # job_id -> req_id for in-flight loads. @@ -50,7 +59,7 @@ def _init_worker(self, kv_caches: CanonicalKVCaches) -> None: def register_kv_caches( self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]] ): - kv_cache_config = self.spec.kv_cache_config + kv_cache_config = self.kv_cache_config num_blocks = kv_cache_config.num_blocks # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use @@ -58,7 +67,7 @@ def register_kv_caches( # General (non-packed) layouts size the tensor at page_size_bytes per # manager block, so page_size_bytes is the correct offloading stride. layer_is_packed: dict[str, bool] = { - ln: bool(kv_tensor.block_stride) + ln: is_kv_cache_tensor_packed(kv_tensor) for kv_tensor in kv_cache_config.kv_cache_tensors for ln in kv_tensor.shared_by } @@ -92,16 +101,17 @@ def register_kv_caches( if layer_is_packed[layer_name] else page ) + raw = torch.empty( + 0, + dtype=torch.int8, + device=layer_kv_cache.device, + ).set_(layer_kv_cache.untyped_storage()) tensors_per_block[layer_name] = ( - torch.tensor( - [], - dtype=torch.int8, - device=layer_kv_cache.device, - ).set_( - layer_kv_cache.untyped_storage(), - byte_offset, + torch.as_strided( + raw, (num_blocks, page), (block_stride_bytes, 1), + byte_offset, ), ) page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes @@ -141,7 +151,7 @@ def register_kv_caches( ( t for t in kv_cache_config.kv_cache_tensors - if t.block_stride and t.shared_by + if is_kv_cache_tensor_packed(t) and t.shared_by ), None, ) @@ -236,7 +246,7 @@ def register_cross_layers_kv_cache( num_blocks_physical_dim = physical_to_logical.index(num_blocks_logical_dim) assert num_blocks_physical_dim == 0 - kv_cache_groups = self.spec.kv_cache_config.kv_cache_groups + kv_cache_groups = self.kv_cache_config.kv_cache_groups assert len(kv_cache_groups) == 1 kv_cache_spec = kv_cache_groups[0].kv_cache_spec num_layers = len(kv_cache_groups[0].layer_names) @@ -270,7 +280,21 @@ def register_cross_layers_kv_cache( def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata): assert self.worker is not None + + # Pop jobs_to_flush from store_jobs into _unsubmitted_store_jobs + # so the existing submission loop below submits them before wait(). + if kv_connector_metadata.jobs_to_flush: + for job_id in kv_connector_metadata.jobs_to_flush: + entry = kv_connector_metadata.store_jobs.pop(job_id, None) + if entry is not None: + assert isinstance(entry.src_spec, GPULoadStoreSpec) + self._unsubmitted_store_jobs.append( + (job_id, entry.src_spec, entry.dst_spec) + ) + + # Submit deferred stores from previous step (and jobs_to_flush above). for job_id, src_spec, dst_spec in self._unsubmitted_store_jobs: + assert isinstance(src_spec, GPULoadStoreSpec) success = self.worker.submit_store(job_id, src_spec, dst_spec) assert success self._unsubmitted_store_jobs.clear() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 197beca9aece..2fe4bf6a5a7c 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -23,6 +23,9 @@ OffloadingConnectorMetadata, OffloadingWorkerMetadata, ) +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import ( + build_offloading_config, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, OffloadPromMetrics, @@ -56,14 +59,17 @@ def __init__( ): super().__init__(vllm_config, role, kv_cache_config) - spec = OffloadingSpecFactory.create_spec(vllm_config, kv_cache_config) + offloading_config = build_offloading_config(vllm_config, kv_cache_config) + spec = OffloadingSpecFactory.create_spec(offloading_config) self.connector_scheduler: OffloadingConnectorScheduler | None = None self.connector_worker: OffloadingConnectorWorker | None = None if role == KVConnectorRole.SCHEDULER: - self.connector_scheduler = OffloadingConnectorScheduler(spec) + self.connector_scheduler = OffloadingConnectorScheduler( + spec, vllm_config, kv_cache_config + ) elif role == KVConnectorRole.WORKER: - self.connector_worker = OffloadingConnectorWorker(spec) + self.connector_worker = OffloadingConnectorWorker(spec, kv_cache_config) def shutdown(self) -> None: if self.connector_worker is not None: diff --git a/vllm/distributed/weight_transfer/__init__.py b/vllm/distributed/weight_transfer/__init__.py index af3322e0cbb7..c78fd1f3cc5f 100644 --- a/vllm/distributed/weight_transfer/__init__.py +++ b/vllm/distributed/weight_transfer/__init__.py @@ -5,10 +5,32 @@ to inference workers. """ -from vllm.distributed.weight_transfer.base import WeightTransferEngine -from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory +from vllm.distributed.weight_transfer.base import ( + ModuleSource, + ParamMeta, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightSource, + WeightTransferEngine, +) +from vllm.distributed.weight_transfer.clients import ( + HTTPVLLMWeightSyncClient, + RayVLLMWeightSyncClient, +) +from vllm.distributed.weight_transfer.factory import ( + WeightTransferEngineFactory, + WeightTransferTrainerFactory, +) __all__ = [ "WeightTransferEngine", "WeightTransferEngineFactory", + "TrainerWeightTransferEngine", + "WeightTransferTrainerFactory", + "VLLMWeightSyncClient", + "HTTPVLLMWeightSyncClient", + "RayVLLMWeightSyncClient", + "ParamMeta", + "WeightSource", + "ModuleSource", ] diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 6dbd768d253b..2e377e292536 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/vllm/distributed/weight_transfer/base.py @@ -5,9 +5,10 @@ from abc import ABC, abstractmethod from collections.abc import Iterator from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Generic, TypeVar +from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, runtime_checkable import torch +from typing_extensions import Self if TYPE_CHECKING: from vllm.config import VllmConfig @@ -17,6 +18,85 @@ TInitInfo = TypeVar("TInitInfo", bound="WeightTransferInitInfo") TUpdateInfo = TypeVar("TUpdateInfo", bound="WeightTransferUpdateInfo") +TConfig = TypeVar("TConfig", bound="WeightTransferConfig") + +# A trainer supplies its parameters as a `WeightSource` (defined below): a +# re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()` +# channel. The built-in `ModuleSource` uses `materialize_full_tensor`. + + +def materialize_full_tensor(tensor: torch.Tensor) -> torch.Tensor: + """Return a full, locally-materialized tensor ready to send. + + FSDP shards (DTensors) expose `full_tensor()`, a collective all-gather; + regular tensors do not and are returned unchanged. Trainer engines call + this at send time so the (potentially expensive) gather happens exactly + once — reading `.shape`/`.dtype` for metadata does not trigger it. + """ + full_tensor = getattr(tensor, "full_tensor", None) + return full_tensor() if callable(full_tensor) else tensor + + +@dataclass(frozen=True) +class ParamMeta: + """Name / wire dtype / full (HF) shape for one output parameter.""" + + name: str + dtype: torch.dtype + shape: tuple[int, ...] + + +class WeightSource(ABC): + """A re-iterable source of the trainer's weights, handed to a trainer engine. + + Two channels: + + * `metadata()` — `(name, wire dtype, full shape)` for every parameter, + *without* transferring. Cheap when shapes are known locally (FSDP + `DTensor` global shape); may be expensive on first call for backends that + must materialize to learn shapes (e.g. a Megatron-Bridge export), in which + case it should cache. + * iteration — yields fully-materialized `(name, tensor)` pairs, one at a + time. Materializing is typically a collective (FSDP `full_tensor()`, a + Megatron export), so every trainer rank must iterate the same source in the + same order in lockstep, or ranks deadlock. Under pipeline parallelism a + rank may not own a parameter at all — iterating still drives the collective + and the yielded tensor is only meaningful on the sender. + + `iter(source)` must yield a *fresh* pass each round. Backends with custom + producer logic (Megatron export, RDT plans, MoE re-fusing) subclass this. + """ + + @abstractmethod + def metadata(self) -> list[ParamMeta]: + raise NotImplementedError + + @abstractmethod + def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: + raise NotImplementedError + + +class ModuleSource(WeightSource): + """`WeightSource` over `module.named_parameters()` — the common case. + + Handles both plain dense modules and FSDP-sharded ones with no special + casing: iteration all-gathers each `DTensor` via `full_tensor()` (a + collective) and passes regular tensors through. `metadata()` reads the + *global* `.shape` / `.dtype`, so it never triggers a gather. + """ + + def __init__(self, module: torch.nn.Module) -> None: + self._module = module + + def metadata(self) -> list[ParamMeta]: + return [ + ParamMeta(name, p.dtype, tuple(p.shape)) + for name, p in self._module.named_parameters() + ] + + def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: + for name, param in self._module.named_parameters(): + yield name, materialize_full_tensor(param) # Base protocols for backend-specific dataclasses @@ -27,6 +107,26 @@ class WeightTransferInitInfo(ABC): # noqa: B024 pass +@dataclass +class TrainerInitInfo(WeightTransferInitInfo): + """Base trainer-side init info: which trainer rank drives the transfer. + + `rank` is this trainer process's rank, provided **explicitly** by the + caller — the engine does not read it from a global process group, which is + ambiguous once several groups (FSDP / TP / PP / EP) exist. Rank 0 is always + the sender: only it opens the endpoint and drives the inference-side RPCs, + while every rank still runs the trainer-side collectives. Backend subclasses + add their own (positional) fields; `rank` is keyword-only so that ordering + never conflicts. + """ + + rank: int = field(kw_only=True) + + @property + def is_sender(self) -> bool: + return self.rank == 0 + + @dataclass class WeightTransferUpdateInfo(ABC): # noqa: B024 """Base class for backend-specific weight update info.""" @@ -243,3 +343,104 @@ def trainer_send_weights( >>> engine.trainer_send_weights(param_iter, trainer_args) """ raise NotImplementedError + + +@runtime_checkable +class VLLMWeightSyncClient(Protocol): + """Trainer-side stub for the inference engine's weight-sync control plane. + + Mirrors the weight-sync methods that the inference engine exposes + (`EngineClient` / the HTTP RLHF routes / Ray actors). A + `TrainerWeightTransferEngine` drives the full handshake through this + protocol so trainer code never has to know the transport. + + All methods are synchronous and accept plain dicts (matching what the + inference side already accepts). Concurrency that some backends need + (e.g. NCCL must run `update_weights` concurrently with the trainer-side + broadcast) is the engine's responsibility, not the client's, so the + protocol stays a flat four-method surface that any wrapper can implement. + + The protocol is structural (PEP 544), so user implementations need only + define these four methods — no import or subclassing required. + """ + + def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ... + + def start_weight_update(self) -> None: ... + + def update_weights(self, update_info: dict[str, Any]) -> None: ... + + def finish_weight_update(self) -> None: ... + + +class TrainerWeightTransferEngine(ABC, Generic[TConfig, TInitInfo]): + """Trainer-side weight transfer engine. + + Symmetric to `WeightTransferEngine` but lives in the training process. + Constructed via the `trainer_init` factory classmethod; carries any + backend-specific state (NCCL communicators, IPC device info, transfer + plans) on `self`. The `WeightSource` is required at `trainer_init`, + then replayed each round by the no-argument `send_weights()`. + + Multi-rank trainers: `trainer_init` and `send_weights` are + called on *every* trainer rank. Rank 0 is the sender, resolved once at + `trainer_init` into `is_sender`. Non-sender ranks still run every + collective (iterating the source, metadata export, IPC handle all-gather) so + the group stays aligned, but each engine explicitly guards the control-plane + RPCs and the transmit on `self.is_sender`, so only the sender touches the + client. + + Subclasses should define: + init_info_cls: Type of backend-specific trainer init info + config_cls: Type of backend-specific config + """ + + # Subclasses should override these class attributes + init_info_cls: type[TInitInfo] + config_cls: type[TConfig] + + def __init__( + self, + config: TConfig, + *, + client: "VLLMWeightSyncClient", + source: "WeightSource", + is_sender: bool = True, + ) -> None: + self.config = config + self.is_sender = is_sender + # The real client is held on every rank; each engine only *calls* it when + # `is_sender`, so non-sender ranks never touch the wire. + self.client = client + self.source = source + + @classmethod + @abstractmethod + def trainer_init( + cls, + config: TConfig, + init_info: TInitInfo, + *, + client: "VLLMWeightSyncClient", + source: "WeightSource", + ) -> Self: + """Rendezvous with the inference side and return a ready instance. + + Called on every trainer rank. The sender drives the full handshake via + `client` (build the worker-side init info, call + `client.init_weight_transfer_engine`, open the trainer-side endpoint); + non-sender ranks skip the rendezvous and the RPC. `source` is stored on + `self.source`; after return, `send_weights()` is callable. + """ + raise NotImplementedError + + @abstractmethod + def send_weights(self) -> None: + """Push `self.source`'s weights to inference workers and drive the full + update round trip: `start_weight_update`, `update_weights` (run + concurrently with the trainer-side broadcast when the backend requires + it), then `finish_weight_update`. Called on every trainer rank.""" + raise NotImplementedError + + def shutdown(self) -> None: + """Tear down communicators / process groups. Default no-op.""" diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py new file mode 100644 index 000000000000..4f54a6e291e3 --- /dev/null +++ b/vllm/distributed/weight_transfer/clients.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Built-in `VLLMWeightSyncClient` implementations. + +These adapt the inference engine's weight-sync control plane to concrete +transports. A `TrainerWeightTransferEngine` takes one of these (or any object +with the same four methods — the protocol is structural) and drives the full +handshake through it. + +Imports of `ray` / `requests` are deferred to call time so this module is +importable without those packages installed. +""" + +from typing import TYPE_CHECKING, Any + +from vllm.distributed.weight_transfer.base import ( + WeightTransferInitRequest, + WeightTransferUpdateRequest, +) + +if TYPE_CHECKING: + from ray.actor import ActorHandle + + +def _json_safe_update_info(update_info: dict[str, Any]) -> dict[str, Any]: + """Make an update_info dict JSON-serializable for HTTP transport. + + CUDA IPC handles (`ipc_handles`) are tuples of non-JSON-native objects, so + over HTTP they are pickled+base64-encoded into `ipc_handles_pickled` (which + the worker auto-deserializes when `VLLM_ALLOW_INSECURE_SERIALIZATION=1`). + Other backends (NCCL) carry only JSON-native metadata and pass through + unchanged. Mirrors the old IPC `_do_send` HTTP branch. + """ + ipc_handles = update_info.get("ipc_handles") + if ipc_handles is None: + return update_info + + import pickle + + import pybase64 as base64 + + out = {k: v for k, v in update_info.items() if k != "ipc_handles"} + out["ipc_handles_pickled"] = base64.b64encode(pickle.dumps(ipc_handles)).decode( + "utf-8" + ) + return out + + +class HTTPVLLMWeightSyncClient: + """Talks to a vLLM server over the RLHF HTTP routes. + + Mirrors `vllm/entrypoints/serve/dev/rlhf/api_router.py`: + `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`, + `/finish_weight_update`. + """ + + def __init__(self, base_url: str, timeout: float = 300) -> None: + self.base_url = base_url.rstrip("/") + self.timeout = timeout + + def _post(self, path: str, json: dict[str, Any] | None = None) -> None: + import requests + + response = requests.post( + f"{self.base_url}/{path}", json=json, timeout=self.timeout + ) + response.raise_for_status() + + def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: + self._post("init_weight_transfer_engine", {"init_info": init_info}) + + def start_weight_update(self) -> None: + self._post("start_weight_update") + + def update_weights(self, update_info: dict[str, Any]) -> None: + self._post( + "update_weights", {"update_info": _json_safe_update_info(update_info)} + ) + + def finish_weight_update(self) -> None: + self._post("finish_weight_update") + + +class RayVLLMWeightSyncClient: + """Talks to one or more vLLM `AsyncLLM`/`LLM` Ray actors. + + Each call fans out to every handle and blocks on all of them, so a + multi-actor (e.g. multi-DP) deployment is driven as one unit. + """ + + def __init__(self, handle: "ActorHandle | list[ActorHandle]") -> None: + self.handles = handle if isinstance(handle, list) else [handle] + + def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: + import ray + + request = WeightTransferInitRequest(init_info=init_info) + ray.get([h.init_weight_transfer_engine.remote(request) for h in self.handles]) + + def start_weight_update(self) -> None: + import ray + + ray.get([h.start_weight_update.remote() for h in self.handles]) + + def update_weights(self, update_info: dict[str, Any]) -> None: + import ray + + request = WeightTransferUpdateRequest(update_info=update_info) + ray.get([h.update_weights.remote(request) for h in self.handles]) + + def finish_weight_update(self) -> None: + import ray + + ray.get([h.finish_weight_update.remote() for h in self.handles]) diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index a253363d7369..4ea27c5ef584 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/vllm/distributed/weight_transfer/factory.py @@ -6,7 +6,10 @@ from collections.abc import Callable from typing import TYPE_CHECKING -from vllm.distributed.weight_transfer.base import WeightTransferEngine +from vllm.distributed.weight_transfer.base import ( + TrainerWeightTransferEngine, + WeightTransferEngine, +) from vllm.logger import init_logger if TYPE_CHECKING: @@ -14,6 +17,11 @@ from vllm.config import VllmConfig from vllm.config.weight_transfer import WeightTransferConfig + from vllm.distributed.weight_transfer.base import ( + VLLMWeightSyncClient, + WeightSource, + WeightTransferInitInfo, + ) logger = init_logger(__name__) @@ -111,6 +119,94 @@ def create_engine( return engine_cls(config, vllm_config, device, model) +class WeightTransferTrainerFactory: + """Factory for creating trainer-side weight transfer engines. + + Parallel to `WeightTransferEngineFactory`, with its own lazy-import + registry. The trainer-side and worker-side registries are kept separate: + they share backend names by convention, but the trainer process never + instantiates a worker engine and vice versa, so unifying them would only + couple the import graphs. + """ + + _registry: dict[str, Callable[[], type[TrainerWeightTransferEngine]]] = {} + + @classmethod + def register_engine( + cls, + name: str, + module_path_or_cls: "str | type[TrainerWeightTransferEngine]", + class_name: str | None = None, + ) -> None: + """Register a trainer engine. Same conventions as + `WeightTransferEngineFactory.register_engine`.""" + if name in cls._registry: + raise ValueError( + f"Weight transfer trainer engine '{name}' is already registered." + ) + + if isinstance(module_path_or_cls, str): + module_path = module_path_or_cls + if class_name is None: + raise ValueError( + "class_name is required when registering with module path" + ) + + def loader() -> type[TrainerWeightTransferEngine]: + module = importlib.import_module(module_path) + return getattr(module, class_name) + + cls._registry[name] = loader + else: + engine_cls = module_path_or_cls + cls._registry[name] = lambda: engine_cls + + @classmethod + def trainer_init( + cls, + backend: str, + config: "WeightTransferConfig", + init_info: "WeightTransferInitInfo", + *, + client: "VLLMWeightSyncClient", + source: "WeightSource", + ) -> TrainerWeightTransferEngine: + """Build and rendezvous a ready-to-send trainer engine. + + Called on every trainer rank (multi-rank trainers construct on all + ranks; the sender is resolved inside the engine's ``trainer_init``). + + Args: + backend: Backend name (must be registered). + config: Backend-specific weight transfer config. + init_info: Backend-specific trainer init info. + client: Inference-side control-plane client. + source: `WeightSource` of `(name, tensor)` pairs to send each round. + + Raises: + ValueError: If the backend is not registered. + """ + if backend not in cls._registry: + available = list(cls._registry.keys()) + raise ValueError( + f"Invalid weight transfer backend: {backend}. " + f"Available trainer engines: {available}" + ) + engine_cls = cls._registry[backend]() + + logger.info( + "Creating weight transfer trainer engine: %s", + engine_cls.__name__, + ) + + return engine_cls.trainer_init( + config=config, + init_info=init_info, + client=client, + source=source, + ) + + # Register built-in weight transfer engines here. # Registration should be centralized to ensure lazy loading - # engine modules are only imported when actually used. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 742d62ac3698..72dfc28688d6 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -602,6 +602,7 @@ class EngineArgs: enable_tower_connector_lora: bool = LoRAConfig.enable_tower_connector_lora specialize_active_lora: bool = LoRAConfig.specialize_active_lora enable_mixed_moe_lora_format: bool = LoRAConfig.enable_mixed_moe_lora_format + enable_moe_shared_loras: bool = LoRAConfig.enable_moe_shared_loras ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override @@ -665,6 +666,7 @@ class EngineArgs: enable_flashinfer_autotune: bool = get_field( KernelConfig, "enable_flashinfer_autotune" ) + enable_bf16x3_router_gemm: bool | None = None worker_cls: str = ParallelConfig.worker_cls worker_extension_cls: str = ParallelConfig.worker_extension_cls @@ -1350,6 +1352,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--enable-mixed-moe-lora-format", **lora_kwargs["enable_mixed_moe_lora_format"], ) + lora_group.add_argument( + "--enable-moe-shared-loras", + **lora_kwargs["enable_moe_shared_loras"], + ) # Observability arguments observability_kwargs = get_kwargs(ObservabilityConfig) @@ -1508,6 +1514,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--enable-flashinfer-autotune", **kernel_kwargs["enable_flashinfer_autotune"], ) + kernel_group.add_argument( + "--enable-bf16x3-router-gemm", + **kernel_kwargs["enable_bf16x3_router_gemm"], + ) moe_backend_kwargs = kernel_kwargs["moe_backend"] moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") kernel_group.add_argument("--moe-backend", **moe_backend_kwargs) @@ -1754,6 +1764,10 @@ def create_speculative_config( if self.speculative_config is None: return None + self.speculative_config = { + k.replace("-", "_"): v for k, v in self.speculative_config.items() + } + # Note(Shangming): These parameters are not obtained from the cli arg # '--speculative-config' and must be passed in when creating the engine # config. @@ -2204,6 +2218,7 @@ def create_engine_config( enable_tower_connector_lora=self.enable_tower_connector_lora, specialize_active_lora=self.specialize_active_lora, enable_mixed_moe_lora_format=self.enable_mixed_moe_lora_format, + enable_moe_shared_loras=self.enable_moe_shared_loras, max_cpu_loras=self.max_cpu_loras if self.max_cpu_loras and self.max_cpu_loras > 0 else None, @@ -2282,6 +2297,8 @@ def create_engine_config( "are mutually exclusive" ) kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune + if self.enable_bf16x3_router_gemm is not None: + kernel_config.enable_bf16x3_router_gemm = self.enable_bf16x3_router_gemm if self.moe_backend != "auto": kernel_config.moe_backend = self.moe_backend if self.linear_backend != "auto": diff --git a/vllm/entrypoints/chat_utils.py b/vllm/entrypoints/chat_utils.py index b7d96c894aed..c89e9fa79d0e 100644 --- a/vllm/entrypoints/chat_utils.py +++ b/vllm/entrypoints/chat_utils.py @@ -3,6 +3,7 @@ import asyncio import json +import types from abc import ABC, abstractmethod from collections import Counter, defaultdict from collections.abc import Awaitable, Callable, Iterable @@ -10,7 +11,19 @@ from functools import cached_property, lru_cache, partial from itertools import accumulate from pathlib import Path -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, TypeAlias, TypeVar, cast +from typing import ( + TYPE_CHECKING, + Any, + Final, + Generic, + Literal, + TypeAlias, + TypeVar, + Union, + cast, + get_args, + get_origin, +) from openai.types.chat import ( ChatCompletionAssistantMessageParam, @@ -1460,6 +1473,25 @@ def _get_full_multimodal_text_prompt( } +def _collect_known_content_part_fields() -> frozenset[str]: + fields: set[str] = set() + stack: list[Any] = [ChatCompletionContentPartParam] + while stack: + node = stack.pop() + if get_origin(node) in (Union, types.UnionType): + stack.extend(get_args(node)) + elif hasattr(node, "__required_keys__"): + fields |= node.__required_keys__ | node.__optional_keys__ + return frozenset(fields) + + +_KNOWN_CONTENT_PART_FIELDS = _collect_known_content_part_fields() + + +def _collect_extra_fields(part: dict[str, Any]) -> dict[str, Any]: + return {k: v for k, v in part.items() if k not in _KNOWN_CONTENT_PART_FIELDS} + + def _parse_chat_message_content_mm_part( part: ChatCompletionContentPartParam, ) -> tuple[str, _ContentPart]: @@ -1669,7 +1701,9 @@ def _parse_chat_message_content_part( str_content = cast(str, content) _reject_reserved_placeholder_in_text(str_content, mm_parser.model_config) if wrap_dicts: - return {"type": "text", "text": str_content} + result: dict[str, Any] = {"type": "text", "text": str_content} + result.update(_collect_extra_fields(cast(dict[str, Any], part))) + return result else: return str_content @@ -1734,7 +1768,9 @@ def _parse_chat_message_content_part( # emit the single sentinel token as text so the template renders # it inline. The renderer later expands it to N tokens post-tokenize. return {"type": "text", "text": PROMPT_EMBEDS_PLACEHOLDER_TOKEN} - return {"type": modality} + result = {"type": modality} + result.update(_collect_extra_fields(cast(dict[str, Any], part))) + return result if modality == "prompt_embeds": # Emit the renderer token inline regardless of `interleave_strings`, # prompt_embeds are spliced at the token offset so position matters. diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 0280c386f250..b3205728e496 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -356,7 +356,7 @@ def _make_config(value: Any, cls: type[_R]) -> _R: # path; the synchronous `LLM` entrypoint runs multimodal # preprocessing serially. Warn so the setting is not a silent # no-op. See vllm-project/vllm#42901. - if self.model_config.renderer_num_workers > 1: + if self.model_config.renderer_num_workers > 1 and self.runner_type != "pooling": logger.warning_once( "`renderer_num_workers=%d` was set, but the offline `LLM` " "entrypoint uses the synchronous renderer path and runs " diff --git a/vllm/entrypoints/pooling/base/io_processor.py b/vllm/entrypoints/pooling/base/io_processor.py index 3d672a13ec34..cdc4c16cacb1 100644 --- a/vllm/entrypoints/pooling/base/io_processor.py +++ b/vllm/entrypoints/pooling/base/io_processor.py @@ -2,9 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Sequence -from typing import Any, Final +from typing import Any, Final, cast -from vllm import PoolingParams, PoolingRequestOutput, PromptType +from vllm import ( + PoolingParams, + PoolingRequestOutput, +) from vllm.config import VllmConfig from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, @@ -14,18 +17,25 @@ ) from vllm.entrypoints.serve.engine.typing import RendererChatRequest, RendererRequest from vllm.inputs import EngineInput, SingletonPrompt -from vllm.renderers import BaseRenderer, TokenizeParams, merge_kwargs +from vllm.lora.request import LoRARequest +from vllm.renderers import BaseRenderer, merge_kwargs from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq from vllm.tool_parsers import ToolParser from vllm.utils.mistral import is_mistral_tokenizer -from ..scoring.typing import ScoringData from ..typing import ( - OfflineInputsContext, + ALLOfflineInputsContext, + EncodeChatRenderParams, + EncodeCMPLRenderParams, + OfflineEncodeInputsContext, OfflineOutputsContext, PoolingChatLikeRequest, PoolingCompletionLikeRequest, + PoolingEngineInput, PoolingServeContext, + RequestFactory, + RequestGenerator, + ScoringRenderParams, ) @@ -98,16 +108,62 @@ def post_process_online( ####################################### # offline APIs - def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: - assert not isinstance(ctx.prompts, ScoringData) and not ( - isinstance(ctx.prompts, dict) and "data" in ctx.prompts - ) + def get_request_factory_offline( + self, ctx: ALLOfflineInputsContext + ) -> tuple[RequestFactory, int]: + assert isinstance(ctx, OfflineEncodeInputsContext) prompts_seq = prompt_to_seq(ctx.prompts) + num_requests = len(prompts_seq) + pooling_task = ctx.pooling_task + + parsed_prompts = [ + ( + prompt + if isinstance(prompt, bytes) + else parse_model_prompt(self.model_config, prompt) + ) + for prompt in prompts_seq + ] tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( **(ctx.tokenization_kwargs or {}) ) - return self._preprocess_cmpl_offline(prompts=prompts_seq, tok_params=tok_params) + + pooling_params: PoolingParams | Sequence[PoolingParams] + if ctx.pooling_params is None: + pooling_params = PoolingParams() + else: + pooling_params = ctx.pooling_params + + params_seq = self._params_to_seq(pooling_params, num_requests) + + for param in params_seq: + if param.task is None: + param.task = pooling_task + elif pooling_task == "plugin": + # `plugin` task uses io_processor.parse_request to verify inputs. + # We actually allow plugin to overwrite pooling_task. + pass + elif param.task != pooling_task: + msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" + raise ValueError(msg) + + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) + seq_priority = self._priority_to_seq(ctx.priorities, num_requests) + + def request_factory() -> RequestGenerator: + for i in range(num_requests): + yield EncodeCMPLRenderParams( + prompts=parsed_prompts[i], + tok_params=tok_params, + prompt_extras=None, + skip_mm_cache=False, + params=params_seq[i], + lora_requests=seq_lora_requests[i], + priorities=seq_priority[i], + ) + + return request_factory, num_requests def post_process_offline( self, @@ -118,6 +174,41 @@ def post_process_offline( ####################################### # helpers + def render( + self, + render_params: EncodeCMPLRenderParams + | EncodeChatRenderParams + | ScoringRenderParams, + ) -> PoolingEngineInput: + if "conversations" in render_params: + render_params = cast(EncodeChatRenderParams, render_params) + (_,), engine_input = self.renderer.render_chat( + conversations=[render_params["conversations"]], + chat_params=render_params["chat_params"], + tok_params=render_params["tok_params"], + prompt_extras=render_params["prompt_extras"], + skip_mm_cache=render_params["skip_mm_cache"], + ) + elif "prompts" in render_params: + render_params = cast(EncodeCMPLRenderParams, render_params) + engine_input = self.renderer.render_cmpl( + prompts=[render_params["prompts"]], + tok_params=render_params["tok_params"], + prompt_extras=render_params["prompt_extras"], + skip_mm_cache=render_params["skip_mm_cache"], + ) + else: + raise ValueError( + f"Unsupported render_params type {render_params.__class__.__name__}" + ) + + return PoolingEngineInput( + prompts=engine_input[0], + params=render_params["params"], + lora_requests=render_params["lora_requests"], + priorities=render_params["priorities"], + ) + def _preprocess_cmpl_online( self, request: RendererRequest, @@ -196,26 +287,6 @@ def _preprocess_chat_online( return conversation, [engine_input] - def _preprocess_cmpl_offline( - self, - prompts: PromptType | Sequence[PromptType], - tok_params: TokenizeParams, - prompt_extras: dict[str, Any] | None = None, - ) -> Sequence[EngineInput]: - prompts = prompt_to_seq(prompts) - parsed_prompts = [ - ( - prompt - if isinstance(prompt, bytes) - else parse_model_prompt(self.model_config, prompt) - ) - for prompt in prompts - ] - - return self.renderer.render_cmpl( - parsed_prompts, tok_params, prompt_extras=prompt_extras - ) - def _validate_chat_template( self, request_chat_template: str | None, @@ -251,3 +322,35 @@ def _params_to_seq( return params return [params] * num_requests + + def _lora_request_to_seq( + self, + lora_request: LoRARequest | None | Sequence[LoRARequest | None], + num_requests: int, + ) -> Sequence[LoRARequest | None]: + if isinstance(lora_request, Sequence): + if len(lora_request) != num_requests: + raise ValueError( + f"The lengths of prompts ({num_requests}) " + f"and lora_request ({len(lora_request)}) must be the same." + ) + + return lora_request + + return [lora_request] * num_requests + + def _priority_to_seq( + self, + priority: Sequence[int] | None, + num_requests: int, + ) -> Sequence[int]: + if priority is not None: + if len(priority) != num_requests: + raise ValueError( + f"The lengths of prompts ({num_requests}) " + f"and priority ({len(priority)}) must be the same." + ) + + return priority + + return [0] * num_requests diff --git a/vllm/entrypoints/pooling/embed/io_processor.py b/vllm/entrypoints/pooling/embed/io_processor.py index ec52e2efd68d..5c2f825a9675 100644 --- a/vllm/entrypoints/pooling/embed/io_processor.py +++ b/vllm/entrypoints/pooling/embed/io_processor.py @@ -28,11 +28,13 @@ from ..base.io_processor import PoolingIOProcessor from ..scoring.io_processor import JinaRankingIOProcessorMixin from ..typing import ( + ALLOfflineInputsContext, ChunkedEmbeddingMetadata, - OfflineInputsContext, + OfflineEncodeInputsContext, PoolingChatLikeRequest, PoolingCompletionLikeRequest, PoolingServeContext, + RequestFactory, ) from .protocol import ( CohereEmbedContent, @@ -665,7 +667,10 @@ def pre_process_online(self, ctx: PoolingServeContext): ctx.engine_inputs = engine_inputs - def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: + def get_request_factory_offline( + self, ctx: ALLOfflineInputsContext + ) -> tuple[RequestFactory, int]: + assert isinstance(ctx, OfflineEncodeInputsContext) if not isinstance(ctx.prompts, Sequence) or len(ctx.prompts) < 2: raise ValueError("The JinaForRanking model requires at least 2 inputs.") @@ -677,4 +682,4 @@ def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput query=text_prompts[-1], docs=text_prompts[:-1] ) - return super().pre_process_offline(ctx) + return super().get_request_factory_offline(ctx) diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index a005bb92b483..bb3be812c6b0 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -20,10 +20,18 @@ from vllm.pooling_params import PoolingParams from vllm.tasks import SCORE_TYPE_MAP, PoolingTask, SupportedTask +from .base.io_processor import PoolingIOProcessor from .factories import init_pooling_io_processors from .scoring.io_processor import ScoringIOProcessor from .scoring.typing import ScoreInput -from .typing import OfflineInputsContext, OfflineOutputsContext +from .typing import ( + ALLOfflineInputsContext, + OfflineEncodeInputsContext, + OfflineOutputsContext, + OfflinePluginInputsContext, + OfflineScoringInputsContext, + RequestFactory, +) logger = init_logger(__name__) @@ -48,6 +56,9 @@ def __init__(self): chat_template_config=self.chat_template_config, ) + # Use thread pool executor to accelerate preprocessing. + self._executor = self.renderer._executor + def encode( self, prompts: PromptType | Sequence[PromptType] | DataPrompt, @@ -93,43 +104,30 @@ def encode( io_processor = self.pooling_io_processors[pooling_task] - if pooling_params is None: - pooling_params = PoolingParams() - - ctx = OfflineInputsContext( - prompts=prompts, - pooling_params=pooling_params, - tokenization_kwargs=tokenization_kwargs, - ) + ctx: ALLOfflineInputsContext + if isinstance(prompts, dict) and "data" in prompts: + ctx = OfflinePluginInputsContext( + pooling_task=pooling_task, + prompts=prompts, # type: ignore[arg-type] + tokenization_kwargs=tokenization_kwargs, + pooling_params=pooling_params, + lora_request=lora_request, + priorities=None, + ) + else: + ctx = OfflineEncodeInputsContext( + pooling_task=pooling_task, + prompts=prompts, + tokenization_kwargs=tokenization_kwargs, + pooling_params=pooling_params, + lora_request=lora_request, + priorities=None, + ) - engine_inputs = io_processor.pre_process_offline(ctx) - n_inputs = len(engine_inputs) - assert ctx.pooling_params is not None - - params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) - - for param in params_seq: - if param.task is None: - param.task = pooling_task - elif pooling_task == "plugin": - # `plugin` task uses io_processor.parse_request to verify inputs. - # We actually allow plugin to overwrite pooling_task. - pass - elif param.task != pooling_task: - msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" - raise ValueError(msg) - - seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs) - seq_priority = self._priority_to_seq(None, n_inputs) - - self._render_and_add_requests( - prompts=engine_inputs, - params=params_seq, - lora_requests=seq_lora_requests, - priorities=seq_priority, + request_factory, num_requests = io_processor.get_request_factory_offline(ctx) + outputs = self._run_tiling_engine( + io_processor, request_factory, num_requests, use_tqdm=use_tqdm ) - - outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput) outputs = io_processor.post_process_offline( ctx=OfflineOutputsContext(outputs=outputs) ) @@ -357,46 +355,121 @@ def score( io_processor = self.pooling_io_processors[score_type] assert isinstance(io_processor, ScoringIOProcessor) - pooling_task = io_processor.pooling_task scoring_data = io_processor.valid_inputs(data_1, data_2) n_queries = len(scoring_data.data_1) if pooling_params is None: pooling_params = PoolingParams() - ctx = OfflineInputsContext( - prompts=scoring_data, + assert isinstance(pooling_params, PoolingParams) + pooling_task = io_processor.pooling_task + if pooling_params.task is None: + pooling_params.task = pooling_task + elif pooling_params.task != pooling_task: + msg = ( + f"You cannot overwrite {pooling_params.task=!r} with {pooling_task=!r}!" + ) + raise ValueError(msg) + + ctx = OfflineScoringInputsContext( + pooling_task=pooling_task, + scoring_data=scoring_data, pooling_params=pooling_params, tokenization_kwargs=tokenization_kwargs, + lora_request=lora_request, chat_template=chat_template, - n_queries=n_queries, + priorities=None, ) - engine_inputs = io_processor.pre_process_offline(ctx) - n_inputs = len(engine_inputs) - - seq_lora_requests = self._lora_request_to_seq(lora_request, n_inputs) - params_seq = self._params_to_seq(ctx.pooling_params, n_inputs) - - for param in params_seq: - if param.task is None: - param.task = pooling_task - elif param.task != pooling_task: - msg = f"You cannot overwrite {param.task=!r} with {pooling_task=!r}!" - raise ValueError(msg) - - seq_priority = self._priority_to_seq(None, n_inputs) - - self._render_and_add_requests( - prompts=engine_inputs, - params=params_seq, - lora_requests=seq_lora_requests, - priorities=seq_priority, + request_factory, num_requests = io_processor.get_request_factory_offline(ctx) + outputs = self._run_tiling_engine( + io_processor, request_factory, num_requests, use_tqdm=use_tqdm ) - outputs = self._run_engine(use_tqdm=use_tqdm, output_type=PoolingRequestOutput) outputs = io_processor.post_process_offline( ctx=OfflineOutputsContext(outputs=outputs, n_queries=n_queries), ) return [ScoringRequestOutput.from_base(item) for item in outputs] + + def _run_tiling_engine( + self, + io_processor: PoolingIOProcessor, + request_factory: RequestFactory, + num_requests: int, + use_tqdm: bool | Callable[..., tqdm] = True, + ): + # Keeping max_num_seqs * 2 requests in the core can already saturate the core. + # Therefore, keep most requests waiting outside the core. + max_requests_in_core = ( + self.llm_engine.vllm_config.scheduler_config.max_num_seqs * 2 + ) + num_requests_in_core = 0 + num_waited_requests = num_requests + + if use_tqdm: + tqdm_func = use_tqdm if callable(use_tqdm) else tqdm + pbar = tqdm_func( + total=num_requests, + desc="Processed prompts", + dynamic_ncols=True, + postfix=f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s", + ) + + outputs: list[PoolingRequestOutput] = [] + added_request_ids: set[str] = set() + + it = self._executor.map(io_processor.render, request_factory()) + + try: + while num_waited_requests or self.llm_engine.has_unfinished_requests(): + requests = [] + for _ in range(max_requests_in_core - num_requests_in_core): + if num_waited_requests == 0: + break + try: + request = next(it) + requests.append(request) + except StopIteration: + num_waited_requests = 0 + break + + num_waited_requests -= 1 + num_requests_in_core += 1 + + if requests: + request_ids = self._render_and_add_requests( + prompts=[x["prompts"] for x in requests], + params=[x["params"] for x in requests], + lora_requests=[x["lora_requests"] for x in requests], + priorities=[x["priorities"] for x in requests], + ) + + for request_id in request_ids: + # undo assign_request_id + request_id = request_id.split("-", 1)[0] + added_request_ids.add(request_id) + + step_outputs = self.llm_engine.step() + for output in step_outputs: + assert isinstance(output, PoolingRequestOutput) + assert output.finished + outputs.append(output) + added_request_ids.discard(output.request_id) + num_requests_in_core -= 1 + + if use_tqdm: + pbar.update(1) + if pbar.n == num_requests: + pbar.refresh() + + except Exception: + if added_request_ids: + self.llm_engine.abort_request(list(added_request_ids)) + raise + + finally: + if use_tqdm: + pbar.close() + + return sorted(outputs, key=lambda x: int(x.request_id)) diff --git a/vllm/entrypoints/pooling/pooling/io_processor.py b/vllm/entrypoints/pooling/pooling/io_processor.py index b07ceeede32b..b0c38b63a32c 100644 --- a/vllm/entrypoints/pooling/pooling/io_processor.py +++ b/vllm/entrypoints/pooling/pooling/io_processor.py @@ -4,13 +4,19 @@ from typing import Any from vllm import PoolingParams, PoolingRequestOutput -from vllm.inputs import EngineInput from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq from ..base.io_processor import PoolingIOProcessor -from ..typing import OfflineInputsContext, OfflineOutputsContext, PoolingServeContext +from ..typing import ( + ALLOfflineInputsContext, + OfflineEncodeInputsContext, + OfflineOutputsContext, + OfflinePluginInputsContext, + PoolingServeContext, + RequestFactory, +) from .protocol import IOProcessorRequest, IOProcessorResponse logger = init_logger(__name__) @@ -107,9 +113,11 @@ def post_process_online( ####################################### # offline APIs - def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: + def get_request_factory_offline( + self, ctx: ALLOfflineInputsContext + ) -> tuple[RequestFactory, int]: + assert isinstance(ctx, OfflinePluginInputsContext) assert isinstance(ctx.prompts, dict) and "data" in ctx.prompts - assert ctx.pooling_params is not None # Validate the request data is valid for the loaded plugin prompt_data = ctx.prompts.get("data") @@ -126,20 +134,35 @@ def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput prompts = self.io_processor.pre_process(prompt=validated_prompt) prompts_seq = prompt_to_seq(prompts) + num_requests = len(prompts_seq) + + pooling_params: PoolingParams | Sequence[PoolingParams] + if ctx.pooling_params is None: + pooling_params = PoolingParams() + else: + pooling_params = ctx.pooling_params + params_seq: list[PoolingParams] = [ self.io_processor.merge_pooling_params(param) for param in self._params_to_seq( - ctx.pooling_params, - len(prompts_seq), + pooling_params, + num_requests, ) ] for p in params_seq: if p.task is None: p.task = "plugin" - ctx.pooling_params = params_seq - ctx.prompts = prompts_seq - return super().pre_process_offline(ctx) + return super().get_request_factory_offline( + OfflineEncodeInputsContext( + prompts=prompts_seq, + pooling_params=params_seq, + pooling_task="plugin", + tokenization_kwargs=ctx.tokenization_kwargs, + lora_request=ctx.lora_request, + priorities=ctx.priorities, + ) + ) def post_process_offline( self, diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index db872b7f304d..e8b3efc32eae 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -2,24 +2,36 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time from collections.abc import Sequence -from typing import Any, TypeAlias +from typing import Any, TypeAlias, cast import torch.nn.functional as F -from vllm import PoolingParams, PoolingRequestOutput, TokensPrompt +from vllm import PoolingParams, PoolingRequestOutput, PromptType, TokensPrompt from vllm.inputs import EngineInput from vllm.renderers import TokenizeParams from vllm.renderers.hf import safe_apply_chat_template -from vllm.renderers.inputs.preprocess import extract_target_prompt +from vllm.renderers.inputs.preprocess import ( + extract_target_prompt, + parse_model_prompt, + prompt_to_seq, +) from vllm.tasks import PoolingTask from vllm.utils.mistral import is_mistral_tokenizer from ...chat_utils import ChatTemplateResolutionError from ..base.io_processor import PoolingIOProcessor from ..typing import ( - OfflineInputsContext, + ALLOfflineInputsContext, + EncodeChatRenderParams, + EncodeCMPLRenderParams, + OfflineEncodeInputsContext, OfflineOutputsContext, + OfflineScoringInputsContext, + PoolingEngineInput, PoolingServeContext, + RequestFactory, + RequestGenerator, + ScoringRenderParams, ) from .protocol import RerankRequest, ScoreRequest, ScoringRequest from .typing import ScoreData, ScoreInput, ScoringData @@ -213,25 +225,37 @@ def post_process_online( ####################################### # offline APIs - def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: - assert isinstance(ctx.prompts, ScoringData) - assert not isinstance(ctx.pooling_params, Sequence) - - tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( - **(ctx.tokenization_kwargs or {}) - ) + def get_request_factory_offline( + self, ctx: ALLOfflineInputsContext + ) -> tuple[RequestFactory, int]: + assert isinstance(ctx, OfflineScoringInputsContext) max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( pooling_params=ctx.pooling_params ) - scoring_data = ctx.prompts + scoring_data = ctx.scoring_data if max_tokens_per_query > 0 or max_tokens_per_doc > 0: scoring_data = self._truncate_scoring_data( scoring_data, max_tokens_per_query, max_tokens_per_doc ) - return self._pre_process(scoring_data, tok_params) + data_1 = score_data_to_prompts(scoring_data.data_1, "query", self.model_config) + data_2 = score_data_to_prompts( + scoring_data.data_2, "document", self.model_config + ) + prompts = data_1 + data_2 + + return super().get_request_factory_offline( + OfflineEncodeInputsContext( + pooling_task=self.pooling_task, + prompts=prompts, + tokenization_kwargs=ctx.tokenization_kwargs, + pooling_params=ctx.pooling_params, + lora_request=ctx.lora_request, + priorities=ctx.priorities, + ) + ) def post_process_offline( self, @@ -258,6 +282,26 @@ def _pre_process( prompts=data_1 + data_2, tok_params=tok_params, prompt_extras=prompt_extras ) + def _preprocess_cmpl_offline( + self, + prompts: PromptType | Sequence[PromptType], + tok_params: TokenizeParams, + prompt_extras: dict[str, Any] | None = None, + ) -> Sequence[EngineInput]: + prompts = prompt_to_seq(prompts) + parsed_prompts = [ + ( + prompt + if isinstance(prompt, bytes) + else parse_model_prompt(self.model_config, prompt) + ) + for prompt in prompts + ] + + return self.renderer.render_cmpl( + parsed_prompts, tok_params, prompt_extras=prompt_extras + ) + def _post_process(self, outputs: list[PoolingRequestOutput], n_queries: int): emb_data_1 = outputs[:n_queries] emb_data_2 = outputs[n_queries:] @@ -429,33 +473,95 @@ def pre_process_online(self, ctx: ScoringServeContext): ####################################### # offline APIs - def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]: - assert isinstance(ctx.prompts, ScoringData) - assert not isinstance(ctx.pooling_params, Sequence) + def get_request_factory_offline( + self, ctx: ALLOfflineInputsContext + ) -> tuple[RequestFactory, int]: + assert isinstance(ctx, OfflineScoringInputsContext) + + data_1 = ctx.scoring_data.data_1 + data_2 = ctx.scoring_data.data_2 + num_requests = len(data_2) + + if len(data_1) == 1: + data_1 = data_1 * num_requests tok_params = self.renderer.default_cmpl_tok_params.with_kwargs( **(ctx.tokenization_kwargs or {}) ) + prompt_extras = ctx.pooling_params.extra_kwargs + + seq_lora_requests = self._lora_request_to_seq(ctx.lora_request, num_requests) + seq_priority = self._priority_to_seq(ctx.priorities, num_requests) + + def request_factory() -> RequestGenerator: + for i in range(num_requests): + yield ScoringRenderParams( + data_1=data_1[i], + data_2=data_2[i], + chat_template=ctx.chat_template, + tok_params=tok_params, + prompt_extras=prompt_extras, + skip_mm_cache=False, + params=ctx.pooling_params, + lora_requests=seq_lora_requests[i], + priorities=seq_priority[i], + ) + + return request_factory, num_requests + + ####################################### + # helpers + + def render( + self, + render_params: EncodeCMPLRenderParams + | EncodeChatRenderParams + | ScoringRenderParams, + ) -> PoolingEngineInput: + if "data_1" not in render_params: + raise ValueError( + f"Unsupported render_params type {render_params.__class__.__name__}" + ) + render_params = cast(ScoringRenderParams, render_params) + + arrival_time = time.time() + + tok_params = render_params["tok_params"] + params = render_params["params"] + prompt_extras = render_params["prompt_extras"] + max_tokens_per_query, max_tokens_per_doc = self._get_token_limits( - pooling_params=ctx.pooling_params + pooling_params=params ) - prompt_extras = ctx.pooling_params.extra_kwargs if ctx.pooling_params else None - engine_inputs, pooling_params_list = self._pre_process( - ctx.prompts, - tok_params, - ctx.pooling_params, - ctx.chat_template, + _, engine_prompt = self.get_score_prompt( + data_1=render_params["data_1"], + data_2=render_params["data_2"], + encode_kwargs=tok_params.get_encode_kwargs(), + chat_template=render_params["chat_template"], max_tokens_per_query=max_tokens_per_query, max_tokens_per_doc=max_tokens_per_doc, - prompt_extras=prompt_extras, + chat_template_kwargs=prompt_extras.get("chat_template_kwargs") + if prompt_extras + else None, ) - ctx.pooling_params = pooling_params_list - return engine_inputs - ####################################### - # helpers + tok_params.apply_post_tokenization(self.tokenizer, engine_prompt) + + if token_type_ids := engine_prompt.pop("token_type_ids", None): + params = params.clone() + compressed = compress_token_type_ids(token_type_ids) + params.extra_kwargs = {"compressed_token_type_ids": compressed} + + engine_input = self.renderer.process_for_engine(engine_prompt, arrival_time) + + return PoolingEngineInput( + prompts=engine_input, + params=params, + lora_requests=render_params["lora_requests"], + priorities=render_params["priorities"], + ) def _pre_process( self, @@ -731,6 +837,49 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM name = "jina-reranking-scoring" pooling_task: PoolingTask = "token_embed" + def get_request_factory_offline( + self, ctx: ALLOfflineInputsContext + ) -> tuple[RequestFactory, int]: + assert isinstance(ctx, OfflineScoringInputsContext) + + scoring_data = ctx.scoring_data + prompt_extras = ctx.pooling_params.extra_kwargs + + queries = self.ensure_str(scoring_data.data_1) + docs = self.ensure_str(scoring_data.data_2) + chat_template_kwargs = ( + prompt_extras.get("chat_template_kwargs") if prompt_extras else None + ) + instruction = ( + chat_template_kwargs.get("instruction") if chat_template_kwargs else None + ) + + if len(queries) == 1: + prompts = [ + self.format_docs_prompts_func( + query=queries[0], docs=docs, instruction=instruction + ) + ] + else: + prompts = [ + self.format_docs_prompts_func( + query=q, docs=[d], instruction=instruction + ) + for q, d in zip(queries, docs) + ] + + return PoolingIOProcessor.get_request_factory_offline( + self, + OfflineEncodeInputsContext( + pooling_task=self.pooling_task, + prompts=prompts, + tokenization_kwargs=ctx.tokenization_kwargs, + pooling_params=ctx.pooling_params, + lora_request=ctx.lora_request, + priorities=ctx.priorities, + ), + ) + def _pre_process( self, scoring_data: ScoringData, diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 54d02c5b61fe..b44c9476e789 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -1,17 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import time -from collections.abc import AsyncGenerator, Sequence +from collections.abc import AsyncGenerator, Callable, Generator, Sequence from dataclasses import dataclass, field -from typing import Any, Generic, TypeAlias, TypeVar +from typing import Any, Generic, TypeAlias, TypedDict, TypeVar from fastapi import Request from pydantic import ConfigDict from vllm import PoolingParams, PoolingRequestOutput, PromptType +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam from vllm.inputs import DataPrompt, EngineInput from vllm.lora.request import LoRARequest +from vllm.renderers import ChatParams, TokenizeParams +from vllm.renderers.inputs import DictPrompt +from ...tasks import PoolingTask from .classify.protocol import ( ClassificationChatRequest, ClassificationCompletionRequest, @@ -34,7 +38,7 @@ PoolingResponse, ) from .scoring.protocol import ScoringRequest, ScoringResponse -from .scoring.typing import ScoringData +from .scoring.typing import ScoreData, ScoringData PoolingCompletionLikeRequest: TypeAlias = ( EmbeddingCompletionRequest @@ -111,13 +115,36 @@ class PoolingServeContext(Generic[PoolingRequestT]): @dataclass class OfflineInputsContext: - prompts: PromptType | Sequence[PromptType] | DataPrompt | ScoringData - pooling_params: PoolingParams | Sequence[PoolingParams] - tokenization_kwargs: dict[str, Any] | None = None - chat_template: str | None = None + pooling_task: PoolingTask + tokenization_kwargs: dict[str, Any] | None + lora_request: Sequence[LoRARequest | None] | None + priorities: Sequence[int] | None - ## for bi-encoder & late-interaction - n_queries: int | None = None + +@dataclass +class OfflineEncodeInputsContext(OfflineInputsContext): + prompts: PromptType | Sequence[PromptType] + pooling_params: PoolingParams | Sequence[PoolingParams] | None + + +@dataclass +class OfflineScoringInputsContext(OfflineInputsContext): + scoring_data: ScoringData + chat_template: str | None + pooling_params: PoolingParams + + +@dataclass +class OfflinePluginInputsContext(OfflineInputsContext): + prompts: DataPrompt + pooling_params: PoolingParams | Sequence[PoolingParams] | None + + +ALLOfflineInputsContext: TypeAlias = ( + OfflineEncodeInputsContext + | OfflineScoringInputsContext + | OfflinePluginInputsContext +) @dataclass @@ -126,3 +153,41 @@ class OfflineOutputsContext: ## for bi-encoder & late-interaction n_queries: int | None = None + + +class RenderParams(TypedDict): + tok_params: TokenizeParams + prompt_extras: dict[str, Any] | None + skip_mm_cache: bool + + params: PoolingParams + lora_requests: LoRARequest | None + priorities: int + + +class EncodeCMPLRenderParams(RenderParams): + prompts: DictPrompt + + +class EncodeChatRenderParams(RenderParams): + conversations: list["ChatCompletionMessageParam"] + chat_params: ChatParams + + +class ScoringRenderParams(RenderParams): + data_1: ScoreData + data_2: ScoreData + chat_template: str | None + + +class PoolingEngineInput(TypedDict): + prompts: EngineInput + params: PoolingParams + lora_requests: LoRARequest | None + priorities: int + + +RequestGenerator: TypeAlias = Generator[ + EncodeCMPLRenderParams | EncodeChatRenderParams | ScoringRenderParams +] +RequestFactory: TypeAlias = Callable[[], RequestGenerator] diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 11308d67c5ec..c22e70b014c4 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -266,8 +266,13 @@ class DerenderChatRequest(BaseModel): and ``parser.parse_delta()`` instead of ``parser.parse()``. """ + # --8<-- [start:derender-chat-request] model: str + """Served model name.""" + generate_response: GenerateResponse + """The complete token-in / token-out engine response to derender.""" + prompt_tokens: int | None = None """Prompt token count for usage; defaults to 0 if omitted. @@ -282,6 +287,7 @@ class DerenderChatRequest(BaseModel): request context they expect (request.tools, request.tool_choice, request._grammar_from_tool_parser, etc.). """ + # --8<-- [end:derender-chat-request] class DerenderCompletionRequest(BaseModel): @@ -292,8 +298,14 @@ class DerenderCompletionRequest(BaseModel): returned by /v1/completions/render. """ + # --8<-- [start:derender-completion-request] model: str + """Served model name.""" + generate_responses: list[GenerateResponse] + """One response per prompt, parallel to the list[GenerateRequest] + returned by /v1/completions/render.""" + prompt_tokens: list[int] | None = None """One prompt token count per response; each defaults to 0 if omitted. @@ -306,6 +318,7 @@ class DerenderCompletionRequest(BaseModel): Mirrors chat_request on DerenderChatRequest. Required by the parsing so parsers receive the full request context. """ + # --8<-- [end:derender-completion-request] @model_validator(mode="after") def _validate_prompt_tokens_length(self) -> "DerenderCompletionRequest": diff --git a/vllm/envs.py b/vllm/envs.py index 5b601cda8cb2..67d8fa5a2c65 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -127,7 +127,6 @@ VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True VLLM_ROCM_USE_AITER: bool = False VLLM_ROCM_USE_AITER_CUSTOM_AR: bool = True - VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False VLLM_ROCM_USE_AITER_LINEAR: bool = True VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False VLLM_ROCM_USE_AITER_MOE: bool = True @@ -182,8 +181,6 @@ VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False VLLM_V1_USE_OUTLINES_CACHE: bool = False - VLLM_TPU_BUCKET_PADDING_GAP: int = 0 - VLLM_TPU_MOST_MODEL_LEN: int | None = None VLLM_TPU_USING_PATHWAYS: bool = False VLLM_USE_DEEP_GEMM: bool = True VLLM_MOE_USE_DEEP_GEMM: bool = True @@ -1194,11 +1191,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_ROCM_USE_AITER_CUSTOM_AR": lambda: ( os.getenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "True").lower() in ("true", "1") ), - # Whether to use aiter paged attention. - # By default is disabled. - "VLLM_ROCM_USE_AITER_PAGED_ATTN": lambda: ( - os.getenv("VLLM_ROCM_USE_AITER_PAGED_ATTN", "False").lower() in ("true", "1") - ), # use aiter linear op if aiter ops are enabled # The following list of related ops # - scaled_mm (per-tensor / rowwise) @@ -1432,8 +1424,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY": lambda: os.getenv( "VLLM_RAY_EXTRA_ENV_VARS_TO_COPY", "" ), - # Whether to use S3 path for model loading in CI via RunAI Streamer - "VLLM_CI_USE_S3": lambda: os.environ.get("VLLM_CI_USE_S3", "0") == "1", # Use model_redirect to redirect the model name to a local folder. # `model_redirect` can be a json file mapping the model between # repo_id and local folder: @@ -1482,16 +1472,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_V1_USE_OUTLINES_CACHE": lambda: ( os.environ.get("VLLM_V1_USE_OUTLINES_CACHE", "0") == "1" ), - # Gap between padding buckets for the forward pass. So we have - # 8, we will run forward pass with [16, 24, 32, ...]. - "VLLM_TPU_BUCKET_PADDING_GAP": lambda: ( - int(os.environ["VLLM_TPU_BUCKET_PADDING_GAP"]) - if "VLLM_TPU_BUCKET_PADDING_GAP" in os.environ - else 0 - ), - "VLLM_TPU_MOST_MODEL_LEN": lambda: maybe_convert_int( - os.environ.get("VLLM_TPU_MOST_MODEL_LEN", None) - ), # Whether using Pathways "VLLM_TPU_USING_PATHWAYS": lambda: bool( "proxy" in os.getenv("JAX_PLATFORMS", "").lower() @@ -1667,15 +1647,6 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE": lambda: int( os.getenv("VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE", "163840") ), - # Specifies the thresholds of the communicated tensor sizes under which - # vllm should use flashinfer fused allreduce. The variable should be a - # JSON with the following format: - # { : } - # Unspecified world sizes will fall back to - # { 2: 64, 4: 1, : 0.5 } - "VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB": lambda: json.loads( - os.getenv("VLLM_FLASHINFER_ALLREDUCE_FUSION_THRESHOLDS_MB", "{}") - ), # MoE routing strategy selector. # See `RoutingSimulator.get_available_strategies()` # for available # strategies. @@ -2147,7 +2118,6 @@ def compile_factors() -> dict[str, object]: "VLLM_DP_MASTER_PORT", "VLLM_NIXL_SIDE_CHANNEL_HOST", "VLLM_RANDOMIZE_DP_DUMMY_INPUTS", - "VLLM_CI_USE_S3", "VLLM_MODEL_REDIRECT_PATH", "VLLM_HOST_IP", "VLLM_FORCE_AOT_LOAD", diff --git a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py index 45bd8f6fcf3e..75f240c2aac7 100644 --- a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py +++ b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py @@ -41,7 +41,13 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) result = torch.empty(input.shape, device=input.device, dtype=out_dtype) scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype) - scale_ub = torch.mean(input).to(scale_dtype) + # scale_ub clamps the per-token amax of |input|. Use a non-degenerate + # upper bound (midway between the mean and max of |input|) so clamping is + # partially active and the baseline comparison is meaningful. + # torch.mean(input) ~= 0 for the zero-mean input would collapse every + # scale to the floor and saturate the output. + input_abs = input.to(torch.float32).abs() + scale_ub = (0.5 * (input_abs.mean() + input_abs.amax())).to(scale_dtype) config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens}) inputs[config_key] = (result, input, scale, scale_ub) diff --git a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py index 3e02169db3b8..8bf352cc5128 100644 --- a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py @@ -48,7 +48,6 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) result = torch.empty(input.shape, device=input.device, dtype=out_dtype) scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype) - scale_ub = torch.mean(input).to(scale_dtype) residual = torch.randn_like(input) weight = torch.normal( mean=1.0, @@ -58,6 +57,16 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: device=input.device, ) epsilon = 1e-6 + # scale_ub clamps the per-token amax of the RMS-normed, weighted output. + # Use a non-degenerate upper bound (midway between the mean and max of + # that magnitude) so clamping is partially active and the baseline + # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean + # input would collapse every scale to the floor and saturate the output. + # Mirrors the reference normalization in baseline() below. + x = input.to(torch.float32) + residual.to(torch.float32) + rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) + x_norm_abs = ((x * rms).to(input.dtype) * weight).abs().to(torch.float32) + scale_ub = (0.5 * (x_norm_abs.mean() + x_norm_abs.amax())).to(scale_dtype) config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens}) inputs[config_key] = (result, input, weight, scale, epsilon, scale_ub, residual) diff --git a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py index da7651a1fa59..fe995b2c626f 100644 --- a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py @@ -55,7 +55,6 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: device=input.device, dtype=scale_dtype, ) - scale_ub = torch.mean(input).to(scale_dtype) residual = torch.randn_like(input) weight = torch.normal( mean=1.0, @@ -65,6 +64,16 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: device=input.device, ) epsilon = 1e-6 + # scale_ub clamps the per-group amax of the RMS-normed, weighted output. + # Use a non-degenerate upper bound (midway between the mean and max of + # that magnitude) so clamping is partially active and the baseline + # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean + # input would collapse every scale to the floor and saturate the output. + # Mirrors the reference normalization in baseline() below. + x = input.to(torch.float32) + residual.to(torch.float32) + rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) + x_norm_abs = ((x * rms).to(input.dtype) * weight).abs().to(torch.float32) + scale_ub = (0.5 * (x_norm_abs.mean() + x_norm_abs.amax())).to(scale_dtype) config_key = CaseKey( { diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py index 06b7f10af2fd..31386943e7cb 100644 --- a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -60,7 +60,14 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: device=input.device, dtype=scale_dtype, ) - scale_ub = torch.mean(input).to(scale_dtype) + # scale_ub clamps the per-group amax of the SiLU-and-mul activation. Use + # a non-degenerate upper bound (midway between the mean and max of the + # activation magnitude) so clamping is partially active and the baseline + # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean + # input would collapse every scale to the floor and saturate the output. + # Mirrors tests/kernels/helion/test_silu_and_mul_per_block_quant.py. + act_abs = SiluAndMul.forward_native(input.to(torch.float32)).abs() + scale_ub = (0.5 * (act_abs.mean() + act_abs.amax())).to(scale_dtype) config_key = CaseKey( { diff --git a/vllm/lora/layers/fused_moe.py b/vllm/lora/layers/fused_moe.py index 63a4ea9a8298..b447af87802e 100644 --- a/vllm/lora/layers/fused_moe.py +++ b/vllm/lora/layers/fused_moe.py @@ -57,6 +57,9 @@ def __init__(self, base_layer: MoERunner) -> None: # For non-gated MoE (is_act_and_mul=False), only 1 slice is needed # since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3) self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1 + # Set from lora_config.enable_moe_shared_loras in create_lora_weights. + # When True, w13 lora_A and w2 lora_B are stored once for all experts. + self.enable_moe_shared_loras = False # Mirrors per-(lora_id) layout of `self.lora_a_stacked` (built in # `create_lora_weights`) so `create_dummy_lora`'s n_slices fallback # matches `lora_a_stacked` length under EP. @@ -150,10 +153,21 @@ def _build_lora_context(self): local_num_experts=self.local_num_experts, punica_wrapper=self.punica_wrapper, use_tuned_config=bool(envs.VLLM_TUNED_CONFIG_FOLDER), + enable_moe_shared_loras=self.enable_moe_shared_loras, aux_stream=self._lora_stream if use_dual_stream else None, events=self._events if use_dual_stream else None, ) + @property + def _w13_a_num_experts(self) -> int: + """Expert-dim of the w13 lora_A buffer: 1 when shared.""" + return 1 if self.enable_moe_shared_loras else self.local_num_experts + + @property + def _w2_b_num_experts(self) -> int: + """Expert-dim of the w2 lora_B buffer: 1 when shared.""" + return 1 if self.enable_moe_shared_loras else self.local_num_experts + def _create_lora_a_weights( self, max_loras: int, @@ -163,7 +177,7 @@ def _create_lora_a_weights( torch.zeros( ( max_loras, - self.local_num_experts, + self._w13_a_num_experts, lora_config.max_lora_rank if not self.fully_sharded else divide(lora_config.max_lora_rank, self.tp_size), @@ -205,7 +219,7 @@ def _create_lora_b_weights(self, max_loras: int, lora_config: LoRAConfig): torch.zeros( ( max_loras, - self.local_num_experts, + self._w2_b_num_experts, self.hidden_size if not self.fully_sharded else divide(self.hidden_size, self.tp_size), @@ -247,6 +261,7 @@ def create_lora_weights( self._verify_ep_fs(lora_config) self.max_loras = lora_config.max_loras self.fully_sharded = lora_config.fully_sharded_loras + self.enable_moe_shared_loras = lora_config.enable_moe_shared_loras self.adapter_enabled = torch.tensor( [0] * (max_loras + 1), dtype=torch.int, device=self.device @@ -263,8 +278,11 @@ def create_lora_weights( for experts_id in range(self.local_num_experts): # For gated MoE: gate_proj (w1), down_proj (w2), up_proj (w3) # For non-gated MoE: up_proj (w1), down_proj (w2) + # Shared factors are collapsed (expert-dim 1): index 0. + w13_a_eid = 0 if self.enable_moe_shared_loras else experts_id + w2_b_eid = 0 if self.enable_moe_shared_loras else experts_id self.lora_a_stacked.append( - self.w13_lora_a_stacked[0][lora_id][experts_id] + self.w13_lora_a_stacked[0][lora_id][w13_a_eid] ) self.lora_a_stacked.append( self.w2_lora_a_stacked[0][lora_id][experts_id] @@ -273,14 +291,12 @@ def create_lora_weights( self.lora_b_stacked.append( self.w13_lora_b_stacked[0][lora_id][experts_id] ) - self.lora_b_stacked.append( - self.w2_lora_b_stacked[0][lora_id][experts_id] - ) + self.lora_b_stacked.append(self.w2_lora_b_stacked[0][lora_id][w2_b_eid]) # Only add w3 (up_proj) for gated MoE (_w13_slices == 2) if self._w13_slices == 2: self.lora_a_stacked.append( - self.w13_lora_a_stacked[1][lora_id][experts_id] + self.w13_lora_a_stacked[1][lora_id][w13_a_eid] ) self.lora_b_stacked.append( self.w13_lora_b_stacked[1][lora_id][experts_id] @@ -340,6 +356,22 @@ def _slice_w2_b(self, w2_lora_b: torch.Tensor) -> torch.Tensor: return w2_lora_b[:, start_idx:end_idx, :] + def _match_expert_dim( + self, src: torch.Tensor, buffer: torch.Tensor + ) -> torch.Tensor: + """Align a LoRA factor's expert-dim (dim 0) to its stacked buffer. + + Equal dims pass through. When the buffer is collapsed to expert-dim 1 + (a shared factor) but the source carries per-expert copies, keep + the first — every expert shares the same factor, so the copies are + identical for a real adapter and irrelevant (zeros) for the dummy. + """ + tgt = buffer.shape[1] + if src.shape[0] == tgt: + return src + assert tgt == 1, f"expert-dim mismatch: source {src.shape[0]} vs buffer {tgt}" + return src[:1] + def reset_lora(self, index: int): """Resets the lora weights at index back to 0.""" for pos in range(self._w13_slices): @@ -366,20 +398,24 @@ def set_lora( self.reset_lora(index) self.adapter_enabled[index] = 1 - num_experts = self.w13_lora_a_stacked[0].shape[1] - w1_lora_a, w2_lora_a, w3_lora_a = lora_a w1_lora_b, w2_lora_b, w3_lora_b = lora_b # EP slicing is done once at add time in # LoRAModelManager._slice_moe_lora_ep, so by here the cached # tensors already match the local-expert dim of the stacked buffers. - assert ( - num_experts - == w1_lora_a.shape[0] - == w2_lora_a.shape[0] - == w3_lora_a.shape[0] - ) + # For shared-loras adapters the w13 lora_A and w2 lora_B buffers are + # collapsed to expert-dim 1: match each factor to the expert-dim of + # the buffer it is copied into. A real shared factor already has + # expert-dim 1 (no-op); a per-expert source (e.g. the warmup dummy + # LoRA, which always stacks to num_experts) is collapsed to the shared + # slot since every expert shares the same factor. + w1_lora_a = self._match_expert_dim(w1_lora_a, self.w13_lora_a_stacked[0]) + w3_lora_a = self._match_expert_dim(w3_lora_a, self.w13_lora_a_stacked[-1]) + w2_lora_a = self._match_expert_dim(w2_lora_a, self.w2_lora_a_stacked[0]) + w1_lora_b = self._match_expert_dim(w1_lora_b, self.w13_lora_b_stacked[0]) + w3_lora_b = self._match_expert_dim(w3_lora_b, self.w13_lora_b_stacked[-1]) + w2_lora_b = self._match_expert_dim(w2_lora_b, self.w2_lora_b_stacked[0]) slliced_w1_lora_a = self._slice_w13_a(w1_lora_a) slliced_w1_lora_b = self._slice_w13_b(w1_lora_b) diff --git a/vllm/lora/lora_weights.py b/vllm/lora/lora_weights.py index f90724c5eb51..38a5875525da 100644 --- a/vllm/lora/lora_weights.py +++ b/vllm/lora/lora_weights.py @@ -227,6 +227,39 @@ def pack_moe( ) return obj + @classmethod + def pack_moe_stacked( + cls, + loras: GenericSequence["LoRALayerWeights | None"], + module_name: str, + ) -> "PackedLoRALayerWeights": + """Pack pre-stacked (3D) w1/w2/w3 expert LoRAs into a single LoRA. + + Unlike :meth:`pack_moe`, which stacks one 2D tensor per expert, this + expects each of the three input LoRAs to already carry the expert + dimension (``experts.w{1,2,3}``). It is used for "shared-outer" + adapters where some factors are shared across experts (expert-dim 1) + and cannot be reconstructed by stacking per-expert tensors. The + produced ``lora_a``/``lora_b`` are the ``[w1, w2, w3]`` lists that + ``FusedMoEWithLoRA.set_lora`` consumes directly. + """ + assert len(loras) == 3, ( + "shared-outer expert LoRA expects exactly w1/w2/w3 stacked tensors" + ) + w1_lora, w2_lora, w3_lora = loras + assert w1_lora is not None and w2_lora is not None and w3_lora is not None + rank = w1_lora.rank + lora_alpha = w1_lora.lora_alpha + scaling = lora_alpha / rank + return cls( + module_name, + rank, + [lora_alpha, lora_alpha, lora_alpha], + [w1_lora.lora_a, w2_lora.lora_a, w3_lora.lora_a], + [w1_lora.lora_b, w2_lora.lora_b, w3_lora.lora_b], + scaling=[scaling, scaling, scaling], + ) + def optimize(self) -> "PackedLoRALayerWeights": """Optimize the LoRA by merging the scaling into lora_b.""" for i in range(len(self.lora_b)): diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index 8b5d97c1e2ee..2c4be5678630 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -132,8 +132,13 @@ def __init__( and self.model.is_3d_moe_weight and not self._enable_mixed_moe_lora_format ) + # Shared MoE adapters: w13 lora_A / w2 lora_B shared across experts, + # stored as pre-stacked experts.w{1,2,3} tensors (startup opt-in). + self._enable_moe_shared_loras = is_moe and lora_config.enable_moe_shared_loras self.packed_modules_mapping = process_packed_modules_mapping( - self.model, force_2d_moe=self._enable_mixed_moe_lora_format + self.model, + force_2d_moe=self._enable_mixed_moe_lora_format, + enable_moe_shared_loras=self._enable_moe_shared_loras, ) self._is_non_gated_moe = is_moe and self.model.is_non_gated_moe self._use_ep = bool( @@ -761,6 +766,16 @@ def _create_merged_loras_inplace(self, lora_model: LoRAModel) -> None: if lora_model.check_lora_name(replaced_module_name): module_name = replaced_module_name if module_name.endswith(".experts"): + if self._enable_moe_shared_loras: + lora_model.loras[module_name] = ( + PackedLoRALayerWeights.pack_moe_stacked( + replacement_loras, + module_name, + ) + ) + for module in packed_module_names: + lora_model.loras.pop(module, None) + continue if self._is_non_gated_moe and len(replacement_loras) > 0: replacement_loras = self._pad_lora_pairs_to_triplets( replacement_loras @@ -1019,20 +1034,10 @@ def _slice_moe_lora_ep( module: FusedMoEWithLoRA, module_name: str, ) -> None: - """Slice the cached LoRA tensors down to this rank's local experts. - - The 2D MoE checkpoint enters as a list of per-(w1/w2/w3) tensors of - shape (num_experts, rank, in) / (num_experts, out, rank). When EP - is active each rank only owns local_num_experts; without this slice - the CPU LoRAModel keeps the full global weight and set_lora has to - re-slice on every activation. - - With the load-time / pack-time slicing in - ``_restrict_to_local_experts``, the stacked tensors already match - ``local_num_experts`` and the inner branch becomes a no-op. The - guard remains so checkpoints that bypassed the pre-slicing (e.g. - ``.bin``/``.pt`` adapters with weights mappers we don't recognize) - still get sliced here. + """Slice cached LoRA tensors down to this rank's local experts. + + Shared factors have expert dimension one and remain available on every + rank. Per-expert factors are narrowed to the local expert block. """ if not module.use_ep: return @@ -1046,16 +1051,13 @@ def _slice_moe_lora_ep( expert_start = ep_rank * local_num_experts expert_end = expert_start + local_num_experts - new_lora_a: list[torch.Tensor | None] = [] - new_lora_b: list[torch.Tensor | None] = [] - for a, b in zip(module_lora.lora_a, module_lora.lora_b): - if a is not None and b is not None and a.shape[0] == global_num_experts: - a = a[expert_start:expert_end].contiguous() - b = b[expert_start:expert_end].contiguous() - new_lora_a.append(a) - new_lora_b.append(b) - module_lora.lora_a = new_lora_a - module_lora.lora_b = new_lora_b + def _slice_local(t: torch.Tensor | None) -> torch.Tensor | None: + if t is not None and t.shape[0] == global_num_experts: + return t[expert_start:expert_end].contiguous() + return t + + module_lora.lora_a = [_slice_local(a) for a in module_lora.lora_a] + module_lora.lora_b = [_slice_local(b) for b in module_lora.lora_b] def _restrict_to_local_experts( self, module_name: str, new_module_names: list[str] diff --git a/vllm/lora/peft_helper.py b/vllm/lora/peft_helper.py index 1443efd4f0cd..a0bb8bf3c283 100644 --- a/vllm/lora/peft_helper.py +++ b/vllm/lora/peft_helper.py @@ -51,6 +51,8 @@ def _validate_features(self) -> list[str]: return error_msg def __post_init__(self): + if self.r <= 0: + raise ValueError(f"LoRA rank `r` must be a positive integer, got {self.r}.") if self.use_rslora: logger.info_once("Loading LoRA weights trained with rsLoRA.") self.vllm_lora_scaling_factor = self.lora_alpha / math.sqrt(self.r) diff --git a/vllm/lora/utils.py b/vllm/lora/utils.py index 6b9c66b980d2..c5d5765bd078 100644 --- a/vllm/lora/utils.py +++ b/vllm/lora/utils.py @@ -369,7 +369,9 @@ def get_adapter_absolute_path(lora_path: str) -> str: def process_packed_modules_mapping( - model: nn.Module, force_2d_moe: bool = False + model: nn.Module, + force_2d_moe: bool = False, + enable_moe_shared_loras: bool = False, ) -> dict[str, list[str]]: if is_moe_model(model): # This method generates and returns a dictionary mapping packed module @@ -382,7 +384,17 @@ def process_packed_modules_mapping( # the engine forces the universal 2D wrapper via # enable_mixed_moe_lora_format (so 3D models can also load 2D # adapters through FusedMoEWithLoRA). - if (not model.is_3d_moe_weight) or force_2d_moe: + if enable_moe_shared_loras: + # Shared MoE adapters store one pre-stacked tensor per + # expert-projection (experts.w1/w2/w3) rather than a per-expert + # tensor list, so the packed mapping references the three stack + # names directly (drives expected_lora_modules and the packing). + packed_modules_mapping["experts"] = [ + "experts.w1", + "experts.w2", + "experts.w3", + ] + elif (not model.is_3d_moe_weight) or force_2d_moe: # Filter out malformed entries: non-gated MoE has empty # ckpt_up_proj_name which results in weight_name containing ".." # (e.g., "experts.0.." instead of "experts.0.layer_name.") diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index baefed95d828..fcc50ffcb0ed 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -550,7 +550,7 @@ def choose_scaled_mm_linear_kernel( scope="global", ) - platform_kernels = possible_kernels[current_platform._enum] + platform_kernels = possible_kernels.get(current_platform._enum, []) # Apply --linear-backend filtering when set. linear_backend = _get_linear_backend() @@ -714,7 +714,7 @@ def choose_mp_linear_kernel( if _cc is not None: compute_capability = _cc[0] * 10 + _cc[1] - platform_kernels = _POSSIBLE_KERNELS[current_platform._enum] + platform_kernels = _POSSIBLE_KERNELS.get(current_platform._enum, []) # Apply --linear-backend filtering when set. linear_backend = _get_linear_backend() diff --git a/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py index c0a5c86b0af5..532a7d55e929 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py @@ -56,7 +56,9 @@ def apply_weights( out_shape = x.shape[:-1] + (layer.output_size_per_partition,) x_2d = x.reshape(-1, x.shape[-1]) - x_fp4, x_scale = flashinfer_mxfp4_quantize(x_2d.contiguous()) + x_fp4, x_scale = flashinfer_mxfp4_quantize( + x_2d.contiguous(), backend="cute-dsl" + ) out = flashinfer_scaled_fp4_mm( x_fp4, weight, diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index d8c4cf80ba96..375049d3b773 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -136,10 +136,12 @@ def set_default_quant_scales(layer: nn.Module, register_buffer: bool = False) -> # We also keep q/k/v_scale on host (cpu) memory for attention # backends that require the scales to be on host instead of on device. - # e.g. Flashinfer + # e.g. Flashinfer & AITER layer._q_scale_float = 1.0 layer._k_scale_float = 1.0 layer._v_scale_float = 1.0 + layer._k_scale_cpu = torch.tensor(1.0, dtype=torch.float32) + layer._v_scale_cpu = torch.tensor(1.0, dtype=torch.float32) layer._prob_scale_float = 1.0 # Initialize q/k/v range constants used by calc_kv_scales @@ -586,6 +588,8 @@ def calc_kv_scales(self, query, key, value): self._q_scale_float = self._q_scale.item() self._k_scale_float = self._k_scale.item() self._v_scale_float = self._v_scale.item() + self._k_scale_cpu.fill_(self._k_scale_float) + self._v_scale_cpu.fill_(self._v_scale_float) # We only calculate the scales once self.calculate_kv_scales = False diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 032f84984116..55ad7aedef23 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -240,6 +240,7 @@ kFp8StaticTensorSym, kNvfp4Dynamic, ) +from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer from vllm.utils.math_utils import cdiv, round_down @@ -393,6 +394,19 @@ def __init__( calculate_kv_scales = False self.quant_config = quant_config + if cache_config is not None and cache_config.kv_cache_dtype_skip_layers: + from vllm.model_executor.models.utils import extract_layer_index + + layer_idx = extract_layer_index(prefix) + if str(layer_idx) in cache_config.kv_cache_dtype_skip_layers: + kv_cache_dtype = "auto" + calculate_kv_scales = False + logger.debug( + "Layer %s: kv_cache_dtype=%s", + prefix, + kv_cache_dtype, + ) + dtype = torch.get_default_dtype() if attn_backend is not None: assert attn_backend.is_mla(), ( @@ -976,9 +990,9 @@ def process_weights_after_loading(self, act_dtype: torch.dtype): ) else: # Convert from (L, N, V) to (N, L, V) - self.W_UV = W_UV.transpose(0, 1) + replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True) # Convert from (L, N, P) to (N, P, L) - self.W_UK_T = W_UK.permute(1, 2, 0) + replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True) # If we should not load quant weights, we initialize the scales to 1.0 # as the default value. See [Note: Register q/k/v/prob scales in state dict] @@ -1011,6 +1025,8 @@ def calc_kv_scales( self._q_scale_float = self._q_scale.item() self._k_scale_float = self._k_scale.item() self._v_scale_float = self._v_scale.item() + self._k_scale_cpu.fill_(self._k_scale_float) + self._v_scale_cpu.fill_(self._v_scale_float) self.calculate_kv_scales = False def get_attn_backend(self) -> type[AttentionBackend]: @@ -1025,7 +1041,7 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: num_kv_heads=1, head_size=self.head_size, dtype=kv_cache_dtype, - cache_dtype_str=vllm_config.cache_config.cache_dtype, + cache_dtype_str=self.kv_cache_dtype, kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) @@ -1485,7 +1501,8 @@ def build_mla_chunked_context_metadata( chunk_ends = torch.min( context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk ) - chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0) + chunk_seq_lens = chunk_ends - chunk_starts + chunk_seq_lens.clamp_(min=0) cu_seq_lens_cpu = torch.zeros( num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True @@ -1535,9 +1552,8 @@ def build_mla_chunked_context_metadata( padded_local_context_lens_cpu.unsqueeze(0), local_chunk_starts + padded_local_max_context_chunk, ) - padded_local_chunk_seq_lens = (local_chunk_ends - local_chunk_starts).clamp( - min=0 - ) + padded_local_chunk_seq_lens = local_chunk_ends - local_chunk_starts + padded_local_chunk_seq_lens.clamp_(min=0) padded_local_cu_seq_lens_cpu = torch.zeros( num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True @@ -1734,6 +1750,7 @@ def __init__( self.determine_chunked_prefill_workspace_size(vllm_config) ) + use_packed_fp8_cache = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" if self.dcp_world_size > 1: # Note(hc): The local kvcache is incomplete when DCP is triggered, # an additional kvcache allgather across the DCP group is therefore @@ -1746,7 +1763,9 @@ def __init__( + self.chunked_prefill_workspace_size // self.dcp_world_size, self.model_config.get_head_size(), ), - dtype=self.model_config.dtype, + dtype=torch.bfloat16 + if use_packed_fp8_cache + else self.model_config.dtype, device=device, ) else: @@ -1755,7 +1774,7 @@ def __init__( self.chunked_prefill_workspace_size, self.model_config.get_head_size(), ), - dtype=self.q_data_type, + dtype=torch.bfloat16 if use_packed_fp8_cache else self.q_data_type, device=device, ) @@ -2100,7 +2119,16 @@ def _compute_prefill_context( for i in range(iters): toks = prefill_metadata.chunked_context.seq_tot[i] - if not use_fp8_prefill: + if self.kv_cache_dtype == "fp8_ds_mla": + ops.cp_gather_and_upconvert_fp8_kv_cache( + src_cache=kv_c_and_k_pe_cache, + dst=workspace[:toks], + block_table=prefill_metadata.block_table, + workspace_starts=prefill_metadata.chunked_context.cu_seq_lens[i], + batch_size=attn_metadata.num_prefills, + seq_starts=prefill_metadata.chunked_context.starts[i], + ) + elif not use_fp8_prefill: ops.gather_and_maybe_dequant_cache( src_cache=kv_c_and_k_pe_cache, dst=workspace, @@ -2215,9 +2243,16 @@ def _context_parallel_compute_prefill_context( padded_local_cu_seq_lens = ( prefill_metadata.chunked_context.padded_local_cu_seq_lens[i] ) - if is_quantized_kv_cache(self.kv_cache_dtype) and ( - self.kv_cache_dtype != "fp8_ds_mla" - ): + if self.kv_cache_dtype == "fp8_ds_mla": + ops.cp_gather_and_upconvert_fp8_kv_cache( + src_cache=kv_c_and_k_pe_cache, + dst=workspace[:toks], + block_table=prefill_metadata.block_table, + workspace_starts=padded_local_cu_seq_lens, + batch_size=attn_metadata.num_prefills, + seq_starts=prefill_metadata.chunked_context.starts[i], + ) + elif is_quantized_kv_cache(self.kv_cache_dtype): assert k_scale is not None ops.gather_and_maybe_dequant_cache( src_cache=kv_c_and_k_pe_cache, diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 4948dca6767e..35bbd4d3d4f3 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -3,7 +3,7 @@ """Shared forward_mha implementation and metadata builder for sparse MLA backends.""" from shutil import which -from typing import TYPE_CHECKING, Generic, TypeVar +from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar import numpy as np import torch @@ -38,6 +38,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): metadata_cls: type[T] + require_uniform_decodes: ClassVar[bool] = False def __init__( self, @@ -173,6 +174,7 @@ def build( num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold or 1, + require_uniform=self.require_uniform_decodes, ) ( prefill_query_start_loc, diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index f05998b1beb6..defca7c2c0bd 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -781,6 +781,7 @@ def _rms_norm_kernel( n_cols, eps, BLOCK_SIZE: tl.constexpr, + HAS_WEIGHT: tl.constexpr, ): """ Compute RMS normalization along the last dimension of a 2D tensor. @@ -813,18 +814,19 @@ def _rms_norm_kernel( col_idx = col_offset + tl.arange(0, BLOCK_SIZE) mask = col_idx < n_cols vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) - weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) # Compute in float32 then convert back to input dtype vals_f32 = vals.to(tl.float32) - weight_f32 = weight.to(tl.float32) - output_f32 = vals_f32 * inv_rms * weight_f32 + output_f32 = vals_f32 * inv_rms + if HAS_WEIGHT: + weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) + output_f32 = output_f32 * weight.to(tl.float32) output = output_f32.to(vals.dtype) tl.store(output_row_start_ptr + col_idx, output, mask=mask) def rms_norm_batch_invariant( input: torch.Tensor, - weight: torch.Tensor, + weight: torch.Tensor | None, eps: float = 1e-6, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: @@ -834,7 +836,8 @@ def rms_norm_batch_invariant( Args: input: Input tensor of shape (..., hidden_size) - weight: Weight tensor of shape (hidden_size,) + weight: Weight tensor of shape (hidden_size,), or None to skip the + per-channel multiply (``RMSNorm(has_weight=False)``) eps: Small constant for numerical stability residual: Optional residual tensor fused into the normalization path @@ -851,17 +854,18 @@ def rms_norm_batch_invariant( ops.fused_add_rms_norm(input, residual, weight, eps) return input, residual - assert weight.dim() == 1, "Weight must be 1-dimensional" - assert input.shape[-1] == weight.shape[0], ( - f"Input last dimension ({input.shape[-1]}) must match " - f"weight dimension ({weight.shape[0]})" - ) + if weight is not None: + assert weight.dim() == 1, "Weight must be 1-dimensional" + assert input.shape[-1] == weight.shape[0], ( + f"Input last dimension ({input.shape[-1]}) must match " + f"weight dimension ({weight.shape[0]})" + ) + weight = weight.contiguous() # Flatten all dimensions except the last one original_shape = input.shape input_2d = input.reshape(-1, input.shape[-1]) input_2d = input_2d.contiguous() - weight = weight.contiguous() n_rows, n_cols = input_2d.shape @@ -870,13 +874,14 @@ def rms_norm_batch_invariant( grid = (n_rows,) _rms_norm_kernel[grid]( input_2d, - weight, + weight if weight is not None else input_2d, output, input_2d.stride(0), output.stride(0), n_cols, eps, BLOCK_SIZE=BLOCK_SIZE, + HAS_WEIGHT=weight is not None, ) return output.reshape(original_shape) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index ad2ed510b33f..39b14fc249e2 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -887,6 +887,9 @@ def int4_w4a16_moe_quant_config( block_shape: list[int] | None = None, a1_gscale: torch.Tensor | None = None, a2_gscale: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for 16-bit float activations and int4 weights. @@ -897,6 +900,9 @@ def int4_w4a16_moe_quant_config( _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp, w1_bias), _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) @@ -950,6 +956,9 @@ def int8_w8a16_moe_quant_config( block_shape: list[int] | None = None, a1_gscale: torch.Tensor | None = None, a2_gscale: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for 16-bit float activations and int8 weights. @@ -960,6 +969,9 @@ def int8_w8a16_moe_quant_config( _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale), _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias), _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/int4_emulation_moe.py b/vllm/model_executor/layers/fused_moe/experts/int4_emulation_moe.py new file mode 100644 index 000000000000..7ec917e46984 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/int4_emulation_moe.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Int4 weight-only quantization emulation for MoE. + +Weights are dequantized from packed int4 to BF16 once at load time; +the forward pass then runs plain TritonExperts in BF16. +""" + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static, + kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, +) +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +class Int4EmulationTritonExperts(TritonExperts): + """Int4 W-only MoE that dequantizes weights to BF16 at load time. + + Weights arrive already dequantized (convert_to_wna16_moe_kernel_format + does the unpacking); apply() simply forwards to TritonExperts. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + logger.warning_once( + "Using Int4EmulationTritonExperts MoE backend. Int4 weights are " + "dequantized to BF16 at load time " + ) + # Weights are dequantized to BF16 before apply() is called, so + # TritonExperts must see them as plain float — clear the int4 dtype + # and scales so the hidden-size assertion and kernel dispatch treat + # them as unquantized. + self.quant_config._w1.dtype = None + self.quant_config._w2.dtype = None + self.quant_config._w1.scale = None + self.quant_config._w2.scale = None + + @staticmethod + def _supports_current_device() -> bool: + return current_platform.is_cuda_alike() + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return ( + weight_key + in ( + kInt4Static, + kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, + ) + and activation_key is None + ) + + @property + def quant_dtype(self) -> torch.dtype | str | None: + return None + + @property + def block_shape(self) -> list[int] | None: + return None + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + if w1.element_size() < 2: + raise RuntimeError( + "Int4EmulationTritonExperts.apply() received packed int4 weights " + "(element_size < 2). Weights must be dequantized to BF16 before " + "the forward pass via convert_to_wna16_moe_kernel_format." + ) + return super().apply( + output=output, + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=expert_tokens_meta, + apply_router_weight_on_input=apply_router_weight_on_input, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_context.py b/vllm/model_executor/layers/fused_moe/experts/lora_context.py index 117f744aeea9..48232913229d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_context.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_context.py @@ -43,6 +43,11 @@ class MoELoRAContext: # try_get_optimal_moe_lora_config for Triton kernel tile configs. use_tuned_config: bool + # Shared MoE LoRA: w13 lora_A and w2 lora_B are shared across all experts + # (stored collapsed with expert-dim 1). When True, LoRAExpertsMixin expands + # them to local_num_experts via a stride-0 view before the kernel. + enable_moe_shared_loras: bool = False + # Optional dual-stream support for overlapping each (base GEMM, LoRA) # pair. When aux_stream is None, the experts.apply() path runs the # original sequential schedule. When set, base GEMM runs on the default diff --git a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py index a47145b9493f..b8dc80ed1815 100644 --- a/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py +++ b/vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py @@ -52,10 +52,19 @@ def apply_w13_lora( torch.Tensor | None, torch.Tensor | None, ]: + w13_lora_a_stacked = lora_context.w13_lora_a_stacked + if lora_context.enable_moe_shared_loras: + # w13 lora_A is shared across experts (collapsed expert-dim 1); + # broadcast to local_num_experts via a stride-0 view. The kernel + # derives num_experts from shape[1] and indexes via stride(1)==0. + w13_lora_a_stacked = tuple( + a.expand(-1, lora_context.local_num_experts, -1, -1) + for a in w13_lora_a_stacked + ) return lora_context.punica_wrapper.add_lora_w13( y, x, - lora_context.w13_lora_a_stacked, + w13_lora_a_stacked, lora_context.w13_lora_b_stacked, topk_ids, topk_weights, @@ -92,11 +101,19 @@ def apply_w2_lora( top_k_num: int, add_inputs: bool = True, ) -> None: + w2_lora_b_stacked = lora_context.w2_lora_b_stacked + if lora_context.enable_moe_shared_loras: + # w2 lora_B is shared across experts (collapsed expert-dim 1); + # broadcast to local_num_experts via a stride-0 view. + w2_lora_b_stacked = tuple( + b.expand(-1, lora_context.local_num_experts, -1, -1) + for b in w2_lora_b_stacked + ) lora_context.punica_wrapper.add_lora_w2( y, x, lora_context.w2_lora_a_stacked, - lora_context.w2_lora_b_stacked, + w2_lora_b_stacked, topk_weights, sorted_token_ids_lora, expert_ids_lora, diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index edcea5361ec1..55ddd1a2c964 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -122,7 +122,9 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo return ( not moe_parallel_config.use_all2all_kernels or moe_parallel_config.use_ag_rs_all2all_kernels - ) and not moe_parallel_config.enable_eplb + ) and not ( + moe_parallel_config.enable_eplb or moe_parallel_config.is_sequence_parallel + ) class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular): diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index f046dfeaf26c..2a1443417328 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -33,6 +33,10 @@ logger = init_logger(__name__) +# Base scale for per-token NVFP4 activation quant; the kernel folds the +# per-token global scale (from the activation amax) on top of it. +_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0) + class TrtLlmNvFp4ExpertsBase: """ @@ -43,9 +47,13 @@ def __init__( self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, + per_token_activation: bool = False, ): self.moe_config = moe_config self.quant_config = quant_config + # Quantize the input here (deferred from prepare) to capture a per-token + # global scale, instead of a static one. + self.per_token_activation = per_token_activation self.routing_method_type = self.moe_config.routing_method self.topk = moe_config.experts_per_token @@ -211,6 +219,27 @@ def _supports_shape(hidden_dim: int) -> bool: def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard + @property + def expects_unquantized_inputs(self) -> bool: + return self.per_token_activation + + def _quantize_per_token_input( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """NVFP4-quantize activations with a per-token global scale. + + Returns ``(packed_fp4, block_scale, per_token_scale)``. + """ + from flashinfer import SfLayout, nvfp4_quantize + + hs_fp4, hs_block_scale, per_token_scale = nvfp4_quantize( + hidden_states, + _PER_TOKEN_BASE_GLOBAL_SCALE, + sfLayout=SfLayout.layout_linear, + per_token_activation=True, + ) + return hs_fp4, hs_block_scale, per_token_scale + def _get_chunk_size(self) -> int: MAX_GRID_Y = 65535 MAX_TILE_TOKENS_DIM = 128 @@ -255,6 +284,15 @@ def workspace_shapes( expert_tokens_meta: mk.ExpertTokensMetadata | None, activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + if self.per_token_activation: + # Deferred input quant leaves K unpacked here, breaking the + # workspace assumptions below. Per-token NVFP4 is only supported on + # the monolithic (non-EP) path for now. + raise NotImplementedError( + "NVFP4 per-token activation is only supported on the monolithic " + "(non-EP) FlashInfer TRTLLM MoE path." + ) + # The workspaces for this implementation are managed by flashinfer. workspace1 = (0,) workspace2 = (0,) @@ -286,6 +324,15 @@ def _invoke_kernel( assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None + # Per-token: input is unquantized, quantize it here. Otherwise it was + # already quantized in prepare() with the static global scale. + if self.per_token_activation: + hidden_states, block_scale, per_token_scale = ( + self._quantize_per_token_input(hidden_states) + ) + else: + block_scale, per_token_scale = a1q_scale, None + # Pack topk ids and weights into format expected by the kernel. packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) output1_scale_gate_scalar = self.quant_config.g1_alphas @@ -295,7 +342,7 @@ def _invoke_kernel( topk_ids=packed_tensor, routing_bias=None, hidden_states=hidden_states, - hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).reshape( + hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape( *hidden_states.shape[:-1], -1 ), gemm1_weights=w1, @@ -321,6 +368,7 @@ def _invoke_kernel( routing_method_type=1, # not used do_finalize=True, activation_type=activation_to_flashinfer_int(activation), + per_token_scale=per_token_scale, output=output, tune_max_num_tokens=min( fi_moe_largest_bucket(self.moe_config), self._get_chunk_size() @@ -346,7 +394,8 @@ def apply( apply_router_weight_on_input: bool, ): assert self._supports_activation(activation) - assert a1q_scale is not None + # Per-token defers input quant to _invoke_kernel, so a1q_scale is None. + assert a1q_scale is not None or self.per_token_activation M = hidden_states.shape[0] chunk_size = self._get_chunk_size() @@ -375,7 +424,7 @@ def apply( topk_ids[start:end], activation, global_num_experts, - a1q_scale[start:end], + None if a1q_scale is None else a1q_scale[start:end], ) @@ -439,7 +488,7 @@ def apply( import flashinfer assert self._supports_activation(activation) - assert a1q_scale is not None + assert a1q_scale is not None or self.per_token_activation assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None assert ( @@ -450,6 +499,14 @@ def apply( and self.routing_method_type != RoutingMethodType.Llama4 ) + # Per-token: input is unquantized, quantize it here (see modular apply). + if self.per_token_activation: + hidden_states, block_scale, per_token_scale = ( + self._quantize_per_token_input(hidden_states) + ) + else: + block_scale, per_token_scale = a1q_scale, None + output1_scale_gate_scalar = self.quant_config.g1_alphas # Invoke kernel. @@ -459,7 +516,7 @@ def apply( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, - hidden_states_scale=a1q_scale.view(torch.float8_e4m3fn).reshape( + hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape( *hidden_states.shape[:-1], -1 ), gemm1_weights=w1, @@ -485,5 +542,6 @@ def apply( routing_method_type=self.routing_method_type, do_finalize=True, activation_type=activation_to_flashinfer_int(activation), + per_token_scale=per_token_scale, tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config), )[0] diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index d0cc08ea141f..a80013f65baa 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -51,6 +51,7 @@ class WNA16MoEBackend(Enum): CPU = "CPU" FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" XPU = "XPU" + EMULATION = "EMULATION" def backend_to_kernel_cls( @@ -87,6 +88,12 @@ def backend_to_kernel_cls( ) return [CPUExpertsInt4] + elif backend == WNA16MoEBackend.EMULATION: + from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import ( + Int4EmulationTritonExperts, + ) + + return [Int4EmulationTritonExperts] else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -105,6 +112,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, WNA16MoEBackend.HUMMING, + WNA16MoEBackend.EMULATION, ] return _AVAILABLE_BACKENDS @@ -115,6 +123,7 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: "marlin": WNA16MoEBackend.MARLIN, "humming": WNA16MoEBackend.HUMMING, "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM, + "emulation": WNA16MoEBackend.EMULATION, } if backend := mapping.get(runner_backend): return backend @@ -215,6 +224,9 @@ def make_wna16_moe_quant_config( w2_bias: torch.Tensor | None = None, a1_gscale: torch.Tensor | None = None, a2_gscale: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, + gemm1_alpha: float | None = None, + gemm1_beta: float | None = None, ) -> FusedMoEQuantConfig: """Create the FusedMoEQuantConfig for 4 or 8-bit WNA16 MoE.""" if num_bits == 4: @@ -228,6 +240,9 @@ def make_wna16_moe_quant_config( block_shape=[0, group_size], a1_gscale=a1_gscale, a2_gscale=a2_gscale, + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) else: assert num_bits == 8 @@ -241,6 +256,9 @@ def make_wna16_moe_quant_config( block_shape=[0, group_size], a1_gscale=a1_gscale, a2_gscale=a2_gscale, + gemm1_clamp_limit=gemm1_clamp_limit, + gemm1_alpha=gemm1_alpha, + gemm1_beta=gemm1_beta, ) @@ -263,19 +281,23 @@ def make_wna16_moe_kernel( from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( CPUExpertsInt4, ) + from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import ( + Int4EmulationTritonExperts, + ) from vllm.model_executor.layers.fused_moe.experts.xpu_moe import ( XPUExpertsWNA16, ) # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts, - # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, and the Humming - # grouped/indexed experts. + # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, the Humming + # grouped/indexed experts, and Int4EmulationTritonExperts allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = ( MarlinExperts, BatchedMarlinExperts, TrtLlmMxint4ExpertsMonolithic, XPUExpertsWNA16, CPUExpertsInt4, + Int4EmulationTritonExperts, ) if backend == WNA16MoEBackend.HUMMING: allowed_experts += tuple(backend_to_kernel_cls(WNA16MoEBackend.HUMMING)) @@ -1017,6 +1039,243 @@ def _humming_wna16_weight_schema( ) +def _unpack_and_dequant_int4_gptq( + w_int32: torch.Tensor, + scale: torch.Tensor, + qzeros: torch.Tensor | None, + transpose_output: bool, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Unpack GPTQ-packed int4 weights and dequantize to output_dtype. + + Args: + w_int32: packed weights, shape [E, K_packed, N] where K_packed = K//8 + (8 nibbles per int32, LSB-first in the K dimension). + scale: per-group scales, shape [E, K//group_size, N], float16. + qzeros: optional asymmetric zero-points, shape [E, K//gs, N//8], int32. + None for symmetric (uint4b8 with implicit bias 8). + transpose_output: if True return [E, N, K]; if False return [E, K, N]. + output_dtype: target floating-point dtype (bfloat16 or float16). + + Returns: + Dequantized weight tensor in the requested layout. + """ + E, K_packed, N = w_int32.shape + K = K_packed * 8 + + # Unpack: [E, K_packed, N] -> [E, K_packed, N, 8] via bit-shifts. + # The nibble index (last dim) enumerates K rows within each packed column, + # so we must fuse K_packed and the nibble dim, not N and the nibble dim. + # Permute to [E, K_packed, 8, N] before reshaping to [E, K, N]. + shifts = torch.arange(8, device=w_int32.device, dtype=torch.int32) * 4 + nibbles = (w_int32.unsqueeze(-1) >> shifts) & 0xF # [E, K_packed, N, 8] + + # Reshape to [E, K, N]: fuse K_packed and nibble index (dim 1 and 3) + w = nibbles.permute(0, 1, 3, 2).reshape(E, K, N).to(torch.int16) + + if qzeros is None: + # Symmetric uint4b8: subtract bias so the range is [-8, 7] + w = w - 8 + else: + # Asymmetric: unpack zero-points (same 8-nibble packing) and subtract + # qzeros shape: [E, K//gs, N//8] int32 + gs = K // scale.shape[1] + n_gs = scale.shape[1] + zp_shifts = torch.arange(8, device=qzeros.device, dtype=torch.int32) * 4 + zp_nibbles = (qzeros.unsqueeze(-1) >> zp_shifts) & 0xF # [E, n_gs, N//8, 8] + zp = zp_nibbles.reshape(E, n_gs, N).to(torch.int16) # [E, n_gs, N] + zp = zp.repeat_interleave(gs, dim=1) # [E, K, N] + w = w - zp + + # Broadcast scale [E, K//gs, N] -> [E, K, N] + gs = K // scale.shape[1] + scale_broadcast = scale.repeat_interleave(gs, dim=1).to(output_dtype) + + w_dequant = w.to(output_dtype) * scale_broadcast # [E, K, N] + + if transpose_output: + return w_dequant.permute(0, 2, 1).contiguous() # [E, N, K] + return w_dequant.contiguous() # [E, K, N] + + +def _unpack_and_dequant_int4_awq( + w_int32: torch.Tensor, + scale: torch.Tensor, + qzeros: torch.Tensor | None, + transpose_output: bool, + output_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Unpack AWQ-packed int4 weights and dequantize to output_dtype. + + AWQ packs along the N (column) dimension with an interleave permutation + [0,2,4,6,1,3,5,7] applied before packing, so unpacking must undo that. + + Args: + w_int32: packed weights, shape [E, K, N_packed] where N_packed = N//8 + (8 nibbles per int32, packed along N with AWQ interleaving). + scale: per-group scales, shape [E, K//group_size, N], float16. + qzeros: asymmetric zero-points, shape [E, K//gs, N_packed], int32. + None for symmetric (uint4b8 with implicit bias 8). + transpose_output: if True return [E, N, K]; if False return [E, K, N]. + output_dtype: target floating-point dtype (bfloat16 or float16). + + Returns: + Dequantized weight tensor in the requested layout. + """ + E, K, N_packed = w_int32.shape + N = N_packed * 8 + + # Unpack 8 nibbles per int32 along the N dimension (LSB-first) + shifts = torch.arange(8, device=w_int32.device, dtype=torch.int32) * 4 + # [E, K, N_packed, 8] -> [E, K, N_packed*8] = [E, K, N_interleaved] + nibbles = (w_int32.unsqueeze(-1) >> shifts) & 0xF + w_interleaved = nibbles.reshape(E, K, N) # [E, K, N] but column-interleaved + + # Undo AWQ interleave: packed order is [0,2,4,6,1,3,5,7] within each group + # of 8. Inverse: position i in packed -> original column interleave[i]. + # To reverse: we need the inverse permutation so that + # w[:, :, inv_interleave] = w_interleaved gives the natural column order. + interleave = torch.tensor([0, 2, 4, 6, 1, 3, 5, 7], device=w_int32.device) + inv_interleave = torch.empty_like(interleave) + inv_interleave[interleave] = torch.arange(8, device=w_int32.device) + + # Apply inverse interleave within each group of 8 columns + w_reshaped = w_interleaved.reshape(E, K, N // 8, 8) # [E, K, groups, 8] + w_reordered = w_reshaped[:, :, :, inv_interleave] # undo interleave + w = w_reordered.reshape(E, K, N).to(torch.int16) # [E, K, N] + + if qzeros is None: + w = w - 8 + else: + # qzeros: [E, K//gs, N_packed] int32, same AWQ column packing + gs = K // scale.shape[1] + n_gs = scale.shape[1] + zp_nibbles = (qzeros.unsqueeze(-1) >> shifts) & 0xF # [E, n_gs, N_packed, 8] + zp_interleaved = zp_nibbles.reshape(E, n_gs, N) + zp_reshaped = zp_interleaved.reshape(E, n_gs, N // 8, 8) + zp_reordered = zp_reshaped[:, :, :, inv_interleave] + zp = zp_reordered.reshape(E, n_gs, N).to(torch.int16) # [E, n_gs, N] + zp = zp.repeat_interleave(gs, dim=1) # [E, K, N] + w = w - zp + + gs = K // scale.shape[1] + scale_broadcast = scale.repeat_interleave(gs, dim=1).to(output_dtype) # [E, K, N] + + w_dequant = w.to(output_dtype) * scale_broadcast # [E, K, N] + + if transpose_output: + return w_dequant.permute(0, 2, 1).contiguous() # [E, N, K] + return w_dequant.contiguous() # [E, K, N] + + +def _process_weights_emulation_gptq( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_qzeros: torch.Tensor | None, + w2_qzeros: torch.Tensor | None, +) -> tuple: + """Dequantize int4 weights to BF16 for the emulation backend. + + Inputs are in GPTQ packed format: + w13: [E, K//8, 2*N] int32 (gate+up proj stacked on dim 2) + w2: [E, N//8, K] int32 + w13_scale: [E, K//gs, 2*N] float16 + w2_scale: [E, N//gs, K] float16 + + Outputs (what TritonExperts expects): + w13_out: [E, 2*N, K] bfloat16 + w2_out: [E, K, N] bfloat16 + """ + # w13: packed along K (dim 1), output cols are 2*N (dim 2) + # transpose_output=True yields [E, 2*N, K] + w13_bf16 = _unpack_and_dequant_int4_gptq( + w13, w13_scale, w13_qzeros, transpose_output=True + ) + + # w2: packed along N (dim 1 is N//8), output cols are K (dim 2) + # After unpacking we get [E, N, K]; we want [E, K, N] for TritonExperts + # transpose_output=False gives [E, N, K], then we permute once more + w2_unpacked = _unpack_and_dequant_int4_gptq( + w2, w2_scale, w2_qzeros, transpose_output=False + ) # [E, N, K] + w2_bf16 = w2_unpacked.permute(0, 2, 1).contiguous() # [E, K, N] + + dummy = torch.ones(1, dtype=torch.float16, device=w13.device) + return ( + w13_bf16, # w13_qweight (now bf16, not int32) + w2_bf16, # w2_qweight (now bf16, not int32) + dummy, # w13_scales (unused; nulled out in Int4EmulationTritonExperts) + dummy, # w2_scales (unused) + None, # w13_g_idx + None, # w2_g_idx + None, # w13_g_idx_sort_indices + None, # w2_g_idx_sort_indices + None, # w13_qzeros + None, # w2_qzeros + None, # w13_input_global_scale + None, # w2_input_global_scale + None, # w13_bias + None, # w2_bias + ) + + +def _process_weights_emulation_awq( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_qzeros: torch.Tensor | None, + w2_qzeros: torch.Tensor | None, +) -> tuple: + """Dequantize AWQ int4 weights to BF16 for the emulation backend. + + AWQ inputs: + w13: [E, K, 2*N//8] int32 (packed along N, gate+up on dim 2) + w2: [E, N, K//8] int32 (packed along K) + w13_scale: [E, K//gs, 2*N] float16 + w2_scale: [E, N//gs, K] float16 + + Outputs (what TritonExperts expects): + w13_out: [E, 2*N, K] bfloat16 + w2_out: [E, K, N] bfloat16 + """ + # w13: AWQ-packed along N (dim 2), K is unpacked in dim 1 + # _unpack_and_dequant_int4_awq with transpose_output=True yields [E, 2*N, K] + w13_bf16 = _unpack_and_dequant_int4_awq( + w13, w13_scale, w13_qzeros, transpose_output=True + ) + + # w2: AWQ packs along K (dim 2 is K//8), N is unpacked in dim 1. + # AWQ w2 is [E, N, K//8] — same column-pack format applied to the K dim. + # _unpack_and_dequant_int4_awq expects [E, rows, N_packed] where the + # packed dim is columns. Treat dim 1 as rows and dim 2 as N_packed: + # unpacking gives [E, N, K]. Then permute to [E, K, N]. + w2_unpacked = _unpack_and_dequant_int4_awq( + w2, w2_scale, w2_qzeros, transpose_output=False + ) # [E, N, K] + w2_bf16 = w2_unpacked.permute(0, 2, 1).contiguous() # [E, K, N] + + dummy = torch.ones(1, dtype=torch.float16, device=w13.device) + return ( + w13_bf16, + w2_bf16, + dummy, + dummy, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, @@ -1203,5 +1462,25 @@ def convert_to_wna16_moe_kernel_format( w13_bias_out, w2_bias_out, ) + elif backend == WNA16MoEBackend.EMULATION: + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig + + if isinstance(quant_config, AutoAWQConfig): + return _process_weights_emulation_awq( + w13, + w2, + w13_scale, + w2_scale, + w13_qzeros, + w2_qzeros, + ) + return _process_weights_emulation_gptq( + w13, + w2, + w13_scale, + w2_scale, + w13_qzeros, + w2_qzeros, + ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index f295163568d6..7852b4db8474 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -529,6 +529,7 @@ def make_nvfp4_moe_kernel( backend: NvFp4MoeBackend, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, layer: torch.nn.Module | None = None, + per_token_activation: bool = False, ) -> mk.FusedMoEKernel: # Create Prepare/Finalize. prepare_finalize = maybe_make_prepare_finalize( @@ -546,6 +547,8 @@ def make_nvfp4_moe_kernel( if backend == NvFp4MoeBackend.HUMMING: assert layer is not None extra_kwargs = {"layer": layer} + if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM and per_token_activation: + extra_kwargs["per_token_activation"] = True # Create Experts. if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py index e49d8b2624ab..74341e7681f8 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py @@ -161,8 +161,8 @@ def finalize( ep_size, self.runtime_max_tokens_per_rank, hidden_size ) - combined_output = self.all2all_manager.moe_alltoall.combine( # type: ignore[attr-defined] + self.all2all_manager.combine_into( # type: ignore[attr-defined] payload=fused_expert_output, runtime_max_tokens_per_rank=self.runtime_max_tokens_per_rank, + output=output, ) - output.copy_(combined_output) diff --git a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py new file mode 100644 index 000000000000..2860c64bb3f7 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py @@ -0,0 +1,449 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CuteDSL BF16x3 router GEMM. + +Computes ``X @ W.T`` for BF16 ``X`` with shape ``[N, K]`` and FP32 router +weights ``W`` with shape ``[M, K]`` by decomposing each FP32 weight value into +three BF16 residual terms inside the kernel, then accumulating the three BF16 +MMA results into FP32 TMEM output. +""" + +from functools import cache + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync +from cutlass.cutlass_dsl import T, dsl_user_op +from quack.compile_utils import make_fake_tensor + +from vllm.cute_utils import _tcgen05, simple_tma_copy +from vllm.triton_utils import tl, triton +from vllm.utils import math_utils + +__all__ = ["bf16x3_router_gemm"] + + +@dsl_user_op +def _decompose_fp32x2_to_3xbf16x2( + w0: Float32, + w1: Float32, + *, + loc=None, + ip=None, +) -> tuple[Uint32, Uint32, Uint32]: + # this PTX snippets does the following + # out0 = BF16(in); res = in - FP32(out0) + # out1 = BF16(res); res = res - FP32(out1) + # out2 = BF16(res) + # + # for normal FP32, this decomposition is exact + # i.e. in = FP32(out0) + FP32(out1) + FP32(out2) + # + out = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32()]), + [w0.ir_value(loc=loc, ip=ip), w1.ir_value(loc=loc, ip=ip)], + "{\n\t" + ".reg .b32 r1_lo, r1_hi, r2_lo, r2_hi;\n\t" + ".reg .b32 w0_lo, w0_hi, w1_lo, w1_hi;\n\t" + ".reg .b64 a_pair, w0_pair, w1_pair, r1_pair, r2_pair;\n\t" + "cvt.rn.bf16x2.f32 $0, $4, $3;\n\t" + "shl.b32 w0_lo, $0, 16;\n\t" + "and.b32 w0_hi, $0, 0xffff0000;\n\t" + "mov.b64 a_pair, {$3, $4};\n\t" + "mov.b64 w0_pair, {w0_lo, w0_hi};\n\t" + "sub.rn.f32x2 r1_pair, a_pair, w0_pair;\n\t" + "mov.b64 {r1_lo, r1_hi}, r1_pair;\n\t" + "cvt.rn.bf16x2.f32 $1, r1_hi, r1_lo;\n\t" + "shl.b32 w1_lo, $1, 16;\n\t" + "and.b32 w1_hi, $1, 0xffff0000;\n\t" + "mov.b64 w1_pair, {w1_lo, w1_hi};\n\t" + "sub.rn.f32x2 r2_pair, r1_pair, w1_pair;\n\t" + "mov.b64 {r2_lo, r2_hi}, r2_pair;\n\t" + "cvt.rn.bf16x2.f32 $2, r2_hi, r2_lo;\n\t" + "}\n", + "=r,=r,=r,f,f", + has_side_effects=False, + is_align_stack=False, + loc=loc, + ip=ip, + ) + return ( + Uint32(llvm.extractvalue(T.i32(), out, [0], loc=loc, ip=ip)), + Uint32(llvm.extractvalue(T.i32(), out, [1], loc=loc, ip=ip)), + Uint32(llvm.extractvalue(T.i32(), out, [2], loc=loc, ip=ip)), + ) + + +class Sm100BF16x3RouterGemm: + def __init__(self, BN: int = 128) -> None: + self.cta_tile = (BN, 128, 64) + self.num_stages = 2 + self.num_warps = 10 + + @cute.jit + def _make_tma(self, tensor: cute.Tensor, BM: int, BK: int): + op = cpasync.CopyBulkTensorTileG2SOp() + swizzle_128B = cute.make_swizzle(3, 4, 3) + elems = 128 * 8 // tensor.element_type.width # 128B + slayout = cute.make_layout( + (BM, (elems, BK // elems), self.num_stages), + stride=(elems, (1, BM * elems), BM * BK), + ) + slayout = cute.make_composed_layout(swizzle_128B, 0, slayout) + return cpasync.make_tiled_tma_atom(op, tensor, slayout, (BM, BK)) + + @cute.jit + def __call__( + self, + X: cute.Tensor, + W: cute.Tensor, + out: cute.Tensor, + split_k: Int32, + stream: CUstream, + ): + BN, BM, BK = self.cta_tile + W_tma = self._make_tma(W, BM, BK) + X_tma = self._make_tma(X, BN, BK) + + grid_m = cute.ceil_div(W.shape[0], BM) + grid_n = cute.ceil_div(X.shape[0], BN) + + self.kernel(X_tma, W_tma, out).launch( + grid=(grid_m, grid_n, split_k), + block=(self.num_warps * 32, 1, 1), + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel(self, X_tma: cpasync.TmaInfo, W_tma: cpasync.TmaInfo, out: cute.Tensor): + tid, _, _ = cute.arch.thread_idx() + bid_m, bid_n, bid_k = cute.arch.block_idx() + _, _, split_k = cute.arch.grid_dim() + + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + BN, BM, BK = self.cta_tile + num_stages = self.num_stages + + N, K = X_tma.tma_tensor.shape + M, _ = W_tma.tma_tensor.shape + k_tiles = cute.ceil_div(K, BK) + + smem = cutlass.utils.SmemAllocator() + sX = smem.allocate_tensor( + BFloat16, + X_tma.smem_layout.outer, + byte_alignment=128, + swizzle=X_tma.smem_layout.inner, + ) + sW = smem.allocate_tensor( + Float32, + W_tma.smem_layout.outer, + byte_alignment=128, + swizzle=W_tma.smem_layout.inner, + ) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + w_full_mbar = smem.allocate_array(Int64, num_stages) + mma_mbar = smem.allocate_array(Int64, 1) + taddr = smem.allocate(Int32, 4) + + BAR_TMEM_ALLOC = 1 + BAR_PREP = 2 + BAR_EPI = 3 + + # tmem "allocation" + # acc_main is for the 1st BF16 term. acc_res is for the 2nd and 3rd + # BF16 terms, which are much smaller than the first term. + acc_main = 0 + acc_res = BN + w_tmem_base = BN * 2 + + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 1) + cute.arch.mbarrier_init(w_full_mbar + i, 128) + cute.arch.mbarrier_init(mma_mbar, 1) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(X_tma.atom) + cpasync.prefetch_descriptor(W_tma.atom) + cute.arch.sync_threads() + cute.arch.griddepcontrol_wait() + + if warp_id == 9: + # TMA warp + stage_id = 0 + parity = 1 + + # (BM, BK, K/BK) + gW_tiles = cute.local_tile(W_tma.tma_tensor, (BM, BK), (bid_m, None)) + gX_tiles = cute.local_tile(X_tma.tma_tensor, (BN, BK), (bid_n, None)) + + for tile_k in cutlass.range(bid_k, k_tiles, split_k, unroll=1): + cute.arch.mbarrier_wait(tma_empty_mbar + stage_id, parity) + mbar = tma_full_mbar + stage_id + with cute.arch.elect_one(): + stage_bytes = BN * BK * 2 + BM * BK * 4 + cute.arch.mbarrier_arrive_and_expect_tx(mbar, stage_bytes) + simple_tma_copy( + W_tma.atom, + gW_tiles[None, None, tile_k], + sW[None, None, stage_id], + mbar, + ) + simple_tma_copy( + X_tma.atom, + gX_tiles[None, None, tile_k], + sX[None, None, stage_id], + mbar, + ) + + stage_id = (stage_id + 1) % num_stages + if stage_id == 0: + parity ^= 1 + + elif warp_id == 8: + # MMA warp + stage_id = 0 + parity = 0 + + idesc = _tcgen05.make_bf16_idesc(BM, BN) + sdesc = _tcgen05.make_sdesc_128B_swizzle(0) + + for tile_k in cutlass.range(bid_k, k_tiles, split_k, unroll=1): + cute.arch.mbarrier_wait(tma_full_mbar + stage_id, parity) + cute.arch.mbarrier_wait(w_full_mbar + stage_id, parity) + _tcgen05.fence_after_thread_sync() + + w_tmem = w_tmem_base + stage_id * (BK // 2 * 3) + x_desc = sdesc | (sX[None, None, stage_id].iterator.toint() >> 4) + + for k in cutlass.range_constexpr(BK // 16): + enable_d = (tile_k > bid_k) or (k > 0) + _tcgen05.mma_ts_f16( + acc_main, w_tmem + k * 8, x_desc, idesc, enable_d + ) + _tcgen05.mma_ts_f16( + acc_res, w_tmem + 32 + k * 8, x_desc, idesc, enable_d + ) + _tcgen05.mma_ts_f16( + acc_res, w_tmem + 64 + k * 8, x_desc, idesc, True + ) + x_desc += 32 >> 4 + + _tcgen05.commit(tma_empty_mbar + stage_id) + + stage_id = (stage_id + 1) % self.num_stages + if stage_id == 0: + parity ^= 1 + + _tcgen05.commit(mma_mbar) + + elif warp_id >= 4: + # prep warps: decompose FP32 W into 3xBF16 + warp_id_ = warp_id % 4 + + stage_id = 0 + parity = 0 + + # ld.shared.v4.f32 + op = cute.nvgpu.CopyUniversalOp() + cp_atom = cute.make_copy_atom(op, Float32, num_bits_per_copy=128) + + # sW_view: ((4, 1), (BK/4, num_stages)) + row = warp_id_ * 32 + lane_id + sW_view = cute.zipped_divide(sW[row, None, None], (4, 1)) + + for _ in cutlass.range(bid_k, k_tiles, split_k, unroll=1): + if warp_id_ == 0: + cute.arch.mbarrier_wait(tma_full_mbar + stage_id, parity) + cute.arch.barrier(barrier_id=BAR_PREP, number_of_threads=128) + + row = warp_id_ * 32 + lane_id + w_tmem = w_tmem_base + stage_id * (BK // 2 * 3) + for kblock in cutlass.range_constexpr(BK // 4): + w0 = cute.make_rmem_tensor(2, Uint32) + w1 = cute.make_rmem_tensor(2, Uint32) + w2 = cute.make_rmem_tensor(2, Uint32) + + w_tmp = cute.make_rmem_tensor(4, Float32) + cute.copy(cp_atom, sW_view[None, (kblock, stage_id)], w_tmp) + w0[0], w1[0], w2[0] = _decompose_fp32x2_to_3xbf16x2( + w_tmp[0], w_tmp[1] + ) + w0[1], w1[1], w2[1] = _decompose_fp32x2_to_3xbf16x2( + w_tmp[2], w_tmp[3] + ) + + tcol = kblock * 2 + _tcgen05.st(warp_id_ * 32, w_tmem + 0 + tcol, "32x32b", 2, w0) + _tcgen05.st(warp_id_ * 32, w_tmem + 32 + tcol, "32x32b", 2, w1) + _tcgen05.st(warp_id_ * 32, w_tmem + 64 + tcol, "32x32b", 2, w2) + + _tcgen05.wait_st() + _tcgen05.fence_before_thread_sync() + cute.arch.mbarrier_arrive(w_full_mbar + stage_id) + + stage_id = (stage_id + 1) % self.num_stages + if stage_id == 0: + parity ^= 1 + + else: + # epilogue warps + if warp_id == 0: + _tcgen05.alloc(taddr) + cute.arch.barrier(barrier_id=BAR_TMEM_ALLOC, number_of_threads=128) + + if warp_id == 0: + cute.arch.mbarrier_wait(mma_mbar, 0) + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=128) + _tcgen05.fence_after_thread_sync() + + cute.arch.griddepcontrol_launch_dependents() + + WIDTH = 8 + for i in cutlass.range_constexpr(BN // WIDTH): + tcol = i * WIDTH + main_regs = cute.make_rmem_tensor(WIDTH, Float32) + res_regs = cute.make_rmem_tensor(WIDTH, Float32) + main_regs.store(_tcgen05.ld(warp_id * 32, tcol, "32x32b", WIDTH)) + res_regs.store(_tcgen05.ld(warp_id * 32, BN + tcol, "32x32b", WIDTH)) + _tcgen05.wait_ld() + + # CuteDSL will codegen add.f32x2 + for j in cutlass.range(WIDTH, vectorize=True): + main_regs[j] += res_regs[j] + + w_row_idx = bid_m * BM + tid + for j in cutlass.range_constexpr(WIDTH): + x_row_idx = bid_n * BN + i * WIDTH + j + if x_row_idx < N and w_row_idx < M: + out[bid_k, x_row_idx, w_row_idx] = main_regs[j] + + cute.arch.barrier(barrier_id=BAR_EPI, number_of_threads=128) + if warp_id == 0: + _tcgen05.dealloc() + + @cache + @staticmethod + def compile(BN: int = 128, K: int = 6144): + N = cute.sym_int() + M = cute.sym_int() + SPLIT_K = cute.sym_int() + X = make_fake_tensor(BFloat16, (N, K), divisibility=8) + W = make_fake_tensor(Float32, (M, K), divisibility=4) + out = make_fake_tensor(Float32, (SPLIT_K, N, M), divisibility=1) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = Sm100BF16x3RouterGemm(BN) + return cute.compile( + kernel, X, W, out, Int32(1), stream, options="--enable-tvm-ffi" + ) + + +@triton.jit +def _splitk_reduce_kernel( + partials, + out, + N, + M: tl.constexpr, + split_stride, + k_splits, + BN: tl.constexpr, + BM: tl.constexpr, + BS: tl.constexpr, + USE_PDL: tl.constexpr, +): + pid_n = tl.program_id(0) + pid_m = tl.program_id(1) + offs_n = pid_n * BN + tl.arange(0, BN) + offs_m = pid_m * BM + tl.arange(0, BM) + offs_s = tl.arange(0, BS) + + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + vals = tl.load( + partials + + offs_s[:, None, None] * split_stride + + offs_n[None, :, None] * M + + offs_m[None, None, :], + mask=( + (offs_s[:, None, None] < k_splits) + & (offs_n[None, :, None] < N) + & (offs_m[None, None, :] < M) + ), + other=0.0, + ) + acc = tl.sum(vals, axis=0) + tl.store( + out + offs_n[:, None] * M + offs_m[None, :], + acc, + mask=(offs_n[:, None] < N) & (offs_m[None, :] < M), + ) + + +def splitk_reduce_triton(partials: torch.Tensor, out: torch.Tensor): + split_k, N, M = partials.shape + block_s = 1 << (split_k - 1).bit_length() + split_stride = partials.stride(0) + if block_s >= 64: + BN, BM = 1, 32 + elif block_s >= 8: + BN, BM = 1, 256 + else: + BN, BM = min(16, 32 // block_s), 32 + grid = (triton.cdiv(N, BN), triton.cdiv(M, BM)) + _splitk_reduce_kernel[grid]( + partials, + out, + N, + M, + split_stride, + split_k, + BN=BN, + BM=BM, + BS=block_s, + USE_PDL=True, + num_warps=4, + launch_pdl=True, + ) + + +def bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor: + """Return ``X @ W.T`` using the SM100 BF16x3 router GEMM kernel.""" + N, K = X.shape + M, _ = W.shape + num_sms = torch.cuda.get_device_properties(X.device).multi_processor_count + + # next power of 2 within 8 and 128 + BN = triton.next_power_of_2(N) + BN = min(max(BN, 8), 128) + + BM = 128 + BK = 64 + k_tiles = math_utils.cdiv(K, BK) + grid_m = math_utils.cdiv(M, BM) + grid_n = math_utils.cdiv(N, BN) + + base_ctas = grid_m * grid_n + split_k = min(k_tiles, max(1, num_sms // base_ctas)) + + partials = X.new_empty(split_k, N, M, dtype=torch.float32) + Sm100BF16x3RouterGemm.compile(BN, K)(X, W, partials, split_k) + + if split_k == 1: + return partials.squeeze(0) + + out = X.new_empty(N, M, dtype=torch.float32) + splitk_reduce_triton(partials, out) + return out diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index 058a96ed6b51..1ddcaa50e83b 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -170,13 +170,12 @@ def fused_topk_bias( hash_indices_table: torch.Tensor | None = None, routed_scaling_factor: float = 1.0, ): - # The topk kernel dispatches dtype based on topk_ids (set by - # indices_type) and assumes input_tokens/hash_indices_table match. - if indices_type is not None: - if input_tokens is not None and input_tokens.dtype != indices_type: - input_tokens = input_tokens.to(dtype=indices_type) - if hash_indices_table is not None and hash_indices_table.dtype != indices_type: - hash_indices_table = hash_indices_table.to(dtype=indices_type) + if ( + input_tokens is not None + and hash_indices_table is not None + and input_tokens.dtype != hash_indices_table.dtype + ): + input_tokens = input_tokens.to(dtype=hash_indices_table.dtype) if not rocm_aiter_ops.is_fused_moe_enabled(): assert hidden_states.size(0) == gating_output.size(0), ( @@ -304,6 +303,7 @@ def fused_topk_bias( scores_for_choice = scores.view(-1, n_routed_experts) # For batch invariance, use sorted=True to ensure deterministic expert selection if hash_indices_table is not None: + assert input_tokens is not None topk_indices = hash_indices_table[input_tokens] else: use_sorted = envs.VLLM_BATCH_INVARIANT diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index e55cd5f4f2ae..39c4f4d81578 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -4,11 +4,15 @@ from torch.nn.parameter import Parameter import vllm._custom_ops as ops +from vllm.config import get_current_vllm_config_or_none +from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op +logger = init_logger(__name__) + @PluggableLayer.register("gate_linear") class GateLinear(ReplicatedLinear): @@ -19,8 +23,9 @@ class GateLinear(ReplicatedLinear): 2. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256) 3. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32, (H, E) in {(3072, 256), (6144, 128), (6144, 256)}) - 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) - 5. F.linear via ReplicatedLinear (ultimate fallback) + 4. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight) + 5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 6. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init (e.g. when the required dtype depends on the expert quantization @@ -86,6 +91,11 @@ def __init__( self._dsv3_max_batch = 16 if is_hopper else 8 # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight) + vllm_config = get_current_vllm_config_or_none() + enable_bf16x3_router_gemm = ( + vllm_config is not None + and vllm_config.kernel_config.enable_bf16x3_router_gemm + ) self.allow_fp32_router_gemm = ( not bias and self.weight.dtype == torch.float32 @@ -93,6 +103,16 @@ def __init__( and (is_hopper or is_blackwell) and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES ) + self.allow_bf16x3_router_gemm = ( + not bias + and self.weight.dtype == torch.float32 + and current_platform.is_cuda() + and is_blackwell + and input_size % 8 == 0 + and enable_bf16x3_router_gemm + ) + if self.allow_bf16x3_router_gemm: + logger.info_once("Enabled experimental SM100 BF16x3 router GEMM.") # cuBLAS bf16→fp32 eligibility self.allow_cublas_router_gemm = ( @@ -173,15 +193,26 @@ def forward( torch.float32, torch.bfloat16, ): - output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight) + output = torch.ops.vllm.fp32_router_gemm_dispatch( + x, self.weight, self.allow_bf16x3_router_gemm + ) return output, None - # Tier 4: cuBLAS bf16→fp32 + # Tier 4: experimental bf16x3 CuteDSL kernel for fp32 router weights + if self.allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: + from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501 + bf16x3_router_gemm, + ) + + output = bf16x3_router_gemm(x, self.weight) + return output, None + + # Tier 5: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: output = torch.mm(x, self.weight.T, out_dtype=torch.float32) return output, None - # Tier 5: F.linear (ReplicatedLinear) + # Tier 6: F.linear (ReplicatedLinear) if self.out_dtype is not None and x.dtype != self.weight.dtype: x = x.to(self.weight.dtype) output, output_bias = super().forward(x) @@ -194,22 +225,34 @@ def forward( def fp32_router_gemm_dispatch_impl( - x: torch.Tensor, weight: torch.Tensor + x: torch.Tensor, + weight: torch.Tensor, + allow_bf16x3_router_gemm: bool, ) -> torch.Tensor: """ Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, - otherwise fall back to F.linear. + otherwise optionally run the experimental BF16x3 kernel for medium/large + SM100 router batches, then fall back to F.linear. This must be wrapped in a custom op because our torch.compile integration does not support runtime dispatching on num_tokens. """ if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: return ops.fp32_router_gemm(x, weight) - else: - return torch.nn.functional.linear(x.float(), weight) + + if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: + from vllm.model_executor.layers.fused_moe.router.bf16x3_router_gemm_cutedsl import ( # noqa: E501 + bf16x3_router_gemm, + ) + + return bf16x3_router_gemm(x, weight) + + return torch.nn.functional.linear(x.float(), weight) def fp32_router_gemm_dispatch_fake( - x: torch.Tensor, weight: torch.Tensor + x: torch.Tensor, + weight: torch.Tensor, + allow_bf16x3_router_gemm: bool, ) -> torch.Tensor: return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.float32) diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 8418245b825b..5f9e763af02c 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -102,9 +102,12 @@ def forward_cuda( assert self.variance_size_override is None, ( "Batch invariance is not supported for variance_size_override" ) + pass_weight = ( + self.pass_weight_add if residual is not None else self.pass_weight + ) return rms_norm_batch_invariant( x, - self.weight.data, + self.weight.data if pass_weight else None, self.variance_epsilon, residual=residual, ) @@ -283,7 +286,9 @@ def forward_native( def forward_cuda( self, x: torch.Tensor, z: torch.Tensor | None = None ) -> torch.Tensor: - from vllm.model_executor.layers.fla.ops.layernorm_guard import rmsnorm_fn + from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( + rmsnorm_fn, + ) return rmsnorm_fn( x, diff --git a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py index 23d7070cc807..f95aac54fdfe 100644 --- a/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py @@ -15,16 +15,16 @@ from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader from vllm.model_executor.utils import set_weight_attrs -from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig -from vllm.utils.torch_utils import direct_register_custom_op -from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata - -from ...fla.ops.kda import ( +from vllm.third_party.flash_linear_attention.ops.kda import ( FusedRMSNormGated, chunk_kda_with_fused_gate, fused_kda_gate, fused_recurrent_kda, ) +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.utils.torch_utils import direct_register_custom_op +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata + from ...linear import ( ColumnParallelLinear, ReplicatedLinear, diff --git a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py index 65da90eccb79..a72707c7a119 100644 --- a/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py @@ -13,10 +13,6 @@ ) from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.custom_op import PluggableLayer -from vllm.model_executor.layers.fla.ops import ( - chunk_gated_delta_rule, - fused_recurrent_gated_delta_rule, -) from vllm.model_executor.layers.layernorm import RMSNormGated from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -37,6 +33,10 @@ ) from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops import ( + chunk_gated_delta_rule, + fused_recurrent_gated_delta_rule, +) from vllm.triton_utils import tl, triton from vllm.triton_utils.allocation import set_triton_allocator from vllm.utils.torch_utils import direct_register_custom_op diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 9e286c692d94..567c063bb3b9 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Qwen3-Next/Qwen3.5 model.""" -import functools from typing import Literal import torch @@ -21,16 +20,6 @@ from vllm.forward_context import ForwardContext, get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp, PluggableLayer -from vllm.model_executor.layers.fla.ops import ( - chunk_gated_delta_rule as fla_chunk_gated_delta_rule, -) -from vllm.model_executor.layers.fla.ops import ( - fused_post_conv_prep, - fused_recurrent_gated_delta_rule_packed_decode, - fused_sigmoid_gating_delta_rule_update, -) -from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd -from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE from vllm.model_executor.layers.layernorm import RMSNormGated from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -56,6 +45,16 @@ ) from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +from vllm.third_party.flash_linear_attention.ops import ( + chunk_gated_delta_rule as fla_chunk_gated_delta_rule, +) +from vllm.third_party.flash_linear_attention.ops import ( + fused_post_conv_prep, + fused_recurrent_gated_delta_rule_packed_decode, + fused_sigmoid_gating_delta_rule_update, +) +from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd +from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import ( @@ -83,70 +82,6 @@ logger = init_logger(__name__) -# TODO(arpera): remove ``_is_libs_cu13_install_intact`` and its caller in -# ``_resolve_gdn_prefill_backend`` once the upstream packaging bug is -# fixed and the broken wheels are yanked / superseded on PyPI: -# https://github.com/NVIDIA/cutlass/issues/3170 -# https://github.com/NVIDIA/cutlass/issues/3259 -@functools.cache -def _is_libs_cu13_install_intact() -> bool: - """Return True if every file installed by ``nvidia-cutlass-dsl-libs-cu13`` - matches the SHA-256 declared in its wheel ``RECORD``. - - ``nvidia-cutlass-dsl-libs-base`` and ``nvidia-cutlass-dsl-libs-cu13`` - both ship into the shared ``nvidia_cutlass_dsl/`` namespace and - write many of the same on-disk paths (the runtime ``.so``, the MLIR - Python bindings, cuTe-DSL Python sources, ...) with different - content. Whichever wheel extracts last wins; with a parallel - installer (e.g. ``uv``) the order is racy and the resulting venv - can end up with a mix of files from both variants. The - ``-libs-base`` variant fails MLIR legalization when JIT-compiling - the FlashInfer Blackwell GDN prefill kernel, and any other - cuTe-DSL-based kernel can break too if on-disk files diverge from - what ``-libs-cu13``'s wheel expects. Tracked upstream at: - - * https://github.com/NVIDIA/cutlass/issues/3170 - * https://github.com/NVIDIA/cutlass/issues/3259 - - This helper re-hashes every file the ``-libs-cu13`` wheel claims to - own and compares against its declared SHA-256. Returns False on any - error (uninstalled, missing RECORD, missing file, hash mismatch). - Result is cached per-process. - """ - import hashlib - import importlib.metadata - - import pybase64 as base64 - - try: - dist = importlib.metadata.distribution("nvidia-cutlass-dsl-libs-cu13") - except importlib.metadata.PackageNotFoundError: - return False - - files = dist.files - if not files: - return False - - for pkg_path in files: - file_hash = pkg_path.hash - # Skip RECORD rows without a hash (RECORD itself, generated - # ``.pyc`` files, ...) and any non-SHA-256 hash modes. - if file_hash is None or not file_hash.value: - continue - if file_hash.mode != "sha256": - continue - try: - with open(pkg_path.locate(), "rb") as f: - digest = hashlib.sha256(f.read()).digest() - except OSError: - return False - actual = base64.urlsafe_b64encode(digest).decode().rstrip("=") - if actual != file_hash.value: - return False - - return True - - def _resolve_gdn_prefill_backend( vllm_config: VllmConfig, ) -> tuple[str, Literal["triton", "flashinfer", "cutedsl"]]: @@ -157,9 +92,7 @@ def _resolve_gdn_prefill_backend( * ``platform == cuda``; * one of the following: - Hopper (SM90) — no further constraints; - - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``, - and an intact ``nvidia-cutlass-dsl-libs-cu13`` install on disk - (see :func:`_is_libs_cu13_install_intact`). + - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``. In-tree CuteDSL GDN prefill kernel is chosen when: * "cutedsl" is requested; (opt-in only) @@ -190,19 +123,8 @@ def _resolve_gdn_prefill_backend( and head_k_dim == 128 and current_platform.get_cuda_runtime_major() >= 13 ): - supports_flashinfer = _is_libs_cu13_install_intact() + supports_flashinfer = True supports_cutedsl = True - if not supports_flashinfer: - logger.warning_once( - "FlashInfer Blackwell GDN requires an intact nvidia-cutlass-dsl" - "-libs-cu13 install, but some on-disk files do not match the " - "SHA-256 declared in its RECORD (install-order race in " - "nvidia-cutlass-dsl packaging -- see " - "https://github.com/NVIDIA/cutlass/issues/3170 and " - "https://github.com/NVIDIA/cutlass/issues/3259). Falling back " - "to Triton/FLA. Repair with: pip install --force-reinstall " - "--no-deps nvidia-cutlass-dsl-libs-cu13" - ) if backend in ["flashinfer", "auto"] and supports_flashinfer: return backend, "flashinfer" diff --git a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py index a00fbc74bf8a..b69595e34069 100644 --- a/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py +++ b/vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py @@ -16,10 +16,6 @@ ) from vllm.forward_context import get_forward_context from vllm.model_executor.custom_op import PluggableLayer -from vllm.model_executor.layers.fla.ops.layernorm_guard import ( - RMSNormGated, - layernorm_fn, -) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -33,6 +29,10 @@ linear_attention_decode, ) from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import ( + RMSNormGated, + layernorm_fn, +) from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py index 2066c04f5224..ba61ec68b5cc 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py @@ -16,7 +16,7 @@ _tcgen05, cvt, fence_before_tma_store, - mma_bf16, + mma_sync, simple_tma_copy, ) @@ -536,14 +536,16 @@ def store_ab_abg( zeros_f32.fill(0.0) Ai_bf16 = cute.make_rmem_tensor(8, BFloat16) - mma_B_bf16 = cute.make_rmem_tensor(8, BFloat16) - M_bf16 = cute.make_rmem_tensor(8, BFloat16) + mma_B_bf16 = cute.make_rmem_tensor((4, 2), BFloat16) + M_bf16 = cute.make_rmem_tensor((4, 2), BFloat16) acc = cute.make_rmem_tensor((4, 2), Float32) - # share the same storage + # ldmatrix copies require rank-1 destination views. + mma_B_bf16_ldsm = cute.group_modes(mma_B_bf16, 0, 2) + + # Packed aliases for BF16x2 arithmetic. Ai = cute.recast_tensor(Ai_bf16, Uint32) - mma_B = cute.logical_divide(cute.recast_tensor(mma_B_bf16, Uint32), 2) - M = cute.logical_divide(cute.recast_tensor(M_bf16, Uint32), 2) + M = cute.recast_tensor(M_bf16, Uint32) # construct rmem-backed identity matrix eye = cute.make_rmem_tensor(4, Uint32) @@ -563,7 +565,11 @@ def store_ab_abg( Ai_f32 = cute.logical_divide(cvt.bf16x2_to_fp32x2(Ai), 4) # M is holding -(I+A), stay constant throughout the iterations - cute.copy(ldsm_trans_atom, sA_ldsm[None, (warp_id_, warp_id_)], M_bf16) + cute.copy( + ldsm_trans_atom, + sA_ldsm[None, (warp_id_, warp_id_)], + cute.group_modes(M_bf16, 0, 2), + ) for i in cutlass.range_constexpr(4): M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i]) @@ -572,8 +578,8 @@ def store_ab_abg( # First MMA: -AiM = Ai @ (-M) cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)]) cute.arch.sync_warp() - acc[None, 0] = mma_bf16(Ai, M[None, 0], zeros_f32) - acc[None, 1] = mma_bf16(Ai, M[None, 1], zeros_f32) + acc[None, 0] = mma_sync(Ai_bf16, M_bf16[None, 0], zeros_f32) + acc[None, 1] = mma_sync(Ai_bf16, M_bf16[None, 1], zeros_f32) Ai_bf16.store(acc.load().to(BFloat16)) # Second MMA: Ai_new = 2Ai + (-AiM) @ Ai @@ -582,10 +588,14 @@ def store_ab_abg( cute.copy( ldsm_trans_atom, sA_ldsm[None, (warp_id_, warp_id_)], - mma_B_bf16, + mma_B_bf16_ldsm, + ) + Ai_f32[None, 0] = mma_sync( + Ai_bf16, mma_B_bf16[None, 0], Ai_f32[None, 0] + ) + Ai_f32[None, 1] = mma_sync( + Ai_bf16, mma_B_bf16[None, 1], Ai_f32[None, 1] ) - Ai_f32[None, 0] = mma_bf16(Ai, mma_B[None, 0], Ai_f32[None, 0]) - Ai_f32[None, 1] = mma_bf16(Ai, mma_B[None, 1], Ai_f32[None, 1]) Ai_bf16.store(Ai_f32.load().to(BFloat16)) cute.copy(stsm_atom, Ai_bf16, sAi_ldsm[None, (warp_id_, warp_id_)]) @@ -613,23 +623,24 @@ def store_ab_abg( neg_Ai = cute.make_rmem_tensor(4, Uint32) for i in cutlass.range_constexpr(4): neg_Ai[i] = _bf16x2_neg(Ai[i]) + neg_Ai_bf16 = cute.recast_tensor(neg_Ai, BFloat16) cute.copy( ldsm_trans_atom, sA_ldsm[None, (warp_id_, warp_id_ - 1)], - mma_B_bf16, + mma_B_bf16_ldsm, ) - acc[None, 0] = mma_bf16(neg_Ai, mma_B[None, 0], zeros_f32) - acc[None, 1] = mma_bf16(neg_Ai, mma_B[None, 1], zeros_f32) + acc[None, 0] = mma_sync(neg_Ai_bf16, mma_B_bf16[None, 0], zeros_f32) + acc[None, 1] = mma_sync(neg_Ai_bf16, mma_B_bf16[None, 1], zeros_f32) Ai_bf16.store(acc.load().to(BFloat16)) cute.copy( ldsm_trans_atom, sAi_ldsm[None, (warp_id_ - 1, warp_id_ - 1)], - mma_B_bf16, + mma_B_bf16_ldsm, ) - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32) + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32) Ai_bf16.store(acc.load().to(BFloat16)) store_ab_abg( acc, @@ -660,10 +671,10 @@ def store_ab_abg( cute.copy( ldsm_trans_atom, sAi_ldsm[None, (tile_col, tile_col)], - mma_B_bf16, + mma_B_bf16_ldsm, ) - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32) + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32) cute.copy( ldsm_atom, sA_ldsm[None, (warp_id_, tile_col + 1)], Ai_bf16 @@ -671,10 +682,10 @@ def store_ab_abg( cute.copy( ldsm_trans_atom, sAi_ldsm[None, (tile_col + 1, tile_col)], - mma_B_bf16, + mma_B_bf16_ldsm, ) - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0]) - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1]) + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], acc[None, 0]) + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], acc[None, 1]) tmp = cute.make_rmem_tensor(8, BFloat16) tmp.store(acc.load().to(BFloat16)) @@ -687,10 +698,10 @@ def store_ab_abg( cute.copy( ldsm_trans_atom, sAi_ldsm[None, (warp_id_, tile_col)], - mma_B_bf16, + mma_B_bf16_ldsm, ) - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32) + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32) tmp.store(acc.load().to(BFloat16)) cute.copy(stsm_atom, tmp, sAi_ldsm[None, (warp_id_, tile_col)]) store_ab_abg( @@ -708,15 +719,27 @@ def store_ab_abg( # warp3: Ai30 = -Ai33 @ (A30 @ Ai00 + A31 @ Ai10 + A32 @ Ai20) if warp_id_ == 3: cute.copy(ldsm_atom, sA_ldsm[None, (3, 0)], Ai_bf16) - cute.copy(ldsm_trans_atom, sAi_ldsm[None, (0, 0)], mma_B_bf16) - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) + cute.copy( + ldsm_trans_atom, + sAi_ldsm[None, (0, 0)], + mma_B_bf16_ldsm, + ) + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32) + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32) for i in cutlass.range_constexpr(1, 3): cute.copy(ldsm_atom, sA_ldsm[None, (3, i)], Ai_bf16) - cute.copy(ldsm_trans_atom, sAi_ldsm[None, (i, 0)], mma_B_bf16) - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], acc[None, 0]) - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], acc[None, 1]) + cute.copy( + ldsm_trans_atom, + sAi_ldsm[None, (i, 0)], + mma_B_bf16_ldsm, + ) + acc[None, 0] = mma_sync( + Ai_bf16, mma_B_bf16[None, 0], acc[None, 0] + ) + acc[None, 1] = mma_sync( + Ai_bf16, mma_B_bf16[None, 1], acc[None, 1] + ) tmp = cute.make_rmem_tensor(8, BFloat16) tmp.store(acc.load().to(BFloat16)) @@ -726,9 +749,13 @@ def store_ab_abg( cute.copy(ldsm_atom, sAi_ldsm[None, (3, 3)], Ai_bf16) for i in cutlass.range_constexpr(4): Ai[i] = _bf16x2_neg(Ai[i]) - cute.copy(ldsm_trans_atom, sAi_ldsm[None, (3, 0)], mma_B_bf16) - acc[None, 0] = mma_bf16(Ai, mma_B[None, 0], zeros_f32) - acc[None, 1] = mma_bf16(Ai, mma_B[None, 1], zeros_f32) + cute.copy( + ldsm_trans_atom, + sAi_ldsm[None, (3, 0)], + mma_B_bf16_ldsm, + ) + acc[None, 0] = mma_sync(Ai_bf16, mma_B_bf16[None, 0], zeros_f32) + acc[None, 1] = mma_sync(Ai_bf16, mma_B_bf16[None, 1], zeros_f32) tmp.store(acc.load().to(BFloat16)) cute.copy(stsm_atom, tmp, sAi_ldsm[None, (3, 0)]) store_ab_abg( diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 866bc30a151a..55e815b2a90e 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -42,6 +42,7 @@ "fp8_per_block", "fp8_per_channel", "int8_per_channel_weight_only", + "nvfp4_per_token", "mxfp8", ] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 58104fa7d256..1e49ec3387e4 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -774,6 +774,9 @@ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfi w2_bias=getattr(layer, "w2_bias", None), a1_gscale=getattr(layer, "w13_input_global_scale", None), a2_gscale=getattr(layer, "w2_input_global_scale", None), + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), ) def select_gemm_impl( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 6f1d20bd7f88..260a8309ff7d 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -1200,6 +1200,10 @@ def _to_scalar(tensor: torch.Tensor) -> float: layer._v_scale_float = _to_scalar(layer.v_scale) layer._q_scale_float = _to_scalar(layer.q_scale) + # Sync host (cpu) scale copies read by AITER fused kernels. + layer._k_scale_cpu.fill_(layer._k_scale_float) + layer._v_scale_cpu.fill_(layer._v_scale_float) + # Discard all placeholders. del layer.k_scale del layer.v_scale diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index 00221485233f..7c03baf1e005 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -142,6 +142,23 @@ def get_moe_method( return CompressedTensorsW4A16FlydslMoEMethod( weight_quant, input_quant, layer.moe_config ) + elif moe_backend == "emulation": + # Although this is called 'Marlin', actually it selects + # emulation backend by calling select_wna16_moe_backend. + # TODO: we need to update CompressedTensorsWNA16MoeMethod + # to honor "--moe-backend" option + from .compressed_tensors_moe_wna16_marlin import ( + CompressedTensorsWNA16MarlinMoEMethod, + ) + + logger.info_once( + "Using CompressedTensorsWNA16MarlinMoEMethod " + "(emulation backend requested)" + ) + return CompressedTensorsWNA16MarlinMoEMethod( + weight_quant, input_quant, layer.moe_config, layer_name + ) + from .compressed_tensors_moe_wna16 import ( CompressedTensorsWNA16MoEMethod, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py index f8faddbd07bf..f2159b0eb2ac 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py @@ -108,7 +108,8 @@ def __init__( # grouped actorder isn't supported by this kernel assert weight_quant.actorder != "group" assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE" + "Only symmetric quantization is supported for MoE. " + "Try --moe-backend emulation." ) def create_weights( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index cfeacc902f41..2dabf7a5154b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -48,7 +48,8 @@ def __init__( # grouped actorder isn't supported by this kernel assert weight_quant.actorder != "group" assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE" + "Only symmetric quantization is supported for MoE. " + "Try --moe-backend emulation." ) def create_weights( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index 8af36bcb1022..2b3317d00f3c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -467,10 +467,16 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Marlin-specific parameters (not needed for Flashinfer) if not is_flashinfer: - replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) - replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) - replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + if w13_g_idx_processed is not None: + replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) + if w2_g_idx_processed is not None: + replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) + if w13_g_idx_sort_indices is not None: + replace_parameter( + layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices + ) + if w2_g_idx_sort_indices is not None: + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) # Register input global scales if present if w13_input_global_scale is not None: @@ -484,8 +490,11 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: torch.nn.Parameter(w2_input_global_scale, requires_grad=False), ) - if self.experts_cls is not None and issubclass( - self.experts_cls, FusedMoEExpertsModular + # Marlin workspace — only needed for Marlin-family backends, not emulation. + if ( + self.experts_cls is not None + and issubclass(self.experts_cls, FusedMoEExpertsModular) + and self.wna16_backend != WNA16MoEBackend.EMULATION ): layer.workspace = marlin_make_workspace_new( layer.w13_weight_g_idx.device, 4 @@ -528,6 +537,9 @@ def get_fused_moe_quant_config( num_bits=self.num_bits, w1_zp=getattr(layer, "w13_weight_zero_point", None), w2_zp=getattr(layer, "w2_weight_zero_point", None), + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py index aee482a14563..06715baddadd 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa4.py @@ -96,7 +96,7 @@ def get_min_capability(cls) -> int: return 75 def _build_input_quant_config(self) -> dict | None: - """Build the config dict that HummingInputSchema.from_config expects.""" + """Build the config dict that BaseInputSchema.from_config expects.""" if self.input_quant is None: return None iq = self.input_quant @@ -112,7 +112,7 @@ def _build_input_quant_config(self) -> dict | None: "dynamic": iq.dynamic, "group_size": iq.group_size or 0, "quant_method": "compressed-tensors", - "format": self.quant_format, + "format": "int-quantized", } def create_weights( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py index 57988f66f7f5..4ca8a1c12c6d 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8.py @@ -96,7 +96,7 @@ def get_min_capability(cls) -> int: return 75 def _build_input_quant_config(self) -> dict | None: - """Build the config dict that HummingInputSchema.from_config expects.""" + """Build the config dict that BaseInputSchema.from_config expects.""" if self.input_quant is None: return None iq = self.input_quant @@ -112,7 +112,7 @@ def _build_input_quant_config(self) -> dict | None: "dynamic": iq.dynamic, "group_size": iq.group_size or 0, "quant_method": "compressed-tensors", - "format": self.quant_format, + "format": "int-quantized", } def create_weights( diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index 9bbbf41115e3..f73852e0cfa6 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -217,8 +217,17 @@ def is_layer_skipped(self, config: dict[str, Any], prefix: str): if hasattr(self, "hf_to_vllm_mapper"): ignored_layers = self.hf_to_vllm_mapper.apply_list(ignored_layers) - if any(module_name in prefix for module_name in ignored_layers): - return True + for module_name in ignored_layers: + # compressed-tensors style entries may be regex patterns prefixed + # with "re:" (e.g. "re:vision_tower.*"). These must be regex-matched; + # plain substring matching never matches the literal "re:..." string, + # so ignored layers get silently quantized. Non-"re:" entries keep the + # existing substring behavior (e.g. bitsandbytes modules_to_not_convert). + if module_name.startswith("re:"): + if re.match(module_name[3:], prefix): + return True + elif module_name in prefix: + return True if "lm_head" in prefix: return True diff --git a/vllm/model_executor/layers/quantization/kv_cache.py b/vllm/model_executor/layers/quantization/kv_cache.py index 100632686b0e..380ad33f28f1 100644 --- a/vllm/model_executor/layers/quantization/kv_cache.py +++ b/vllm/model_executor/layers/quantization/kv_cache.py @@ -144,6 +144,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer._v_scale.copy_(v_scale) layer._k_scale_float = k_scale layer._v_scale_float = v_scale + # Sync host (cpu) scale copies read by AITER fused kernels. + layer._k_scale_cpu.fill_(k_scale) + layer._v_scale_cpu.fill_(v_scale) if k_scale == 1.0 and v_scale == 1.0 and "e5m2" not in layer.kv_cache_dtype: logger.warning_once( "Using KV cache scaling factor 1.0 for fp8_e4m3. " diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index b0a70e102420..3f8a741d357d 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -40,6 +40,9 @@ Mxfp8OnlineLinearMethod, Mxfp8OnlineMoEMethod, ) +from vllm.model_executor.layers.quantization.online.nvfp4 import ( + Nvfp4OnlineMoEMethod, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Static128BlockSym, @@ -47,6 +50,7 @@ kFp8StaticTensorSym, kInt8StaticChannelSym, kMxfp8Dynamic, + kNvfp4Static, ) logger = init_logger(__name__) @@ -68,6 +72,7 @@ kFp8StaticChannelSym: Fp8PtpcOnlineMoEMethod, kMxfp8Dynamic: Mxfp8OnlineMoEMethod, kInt8StaticChannelSym: Int8OnlineMoEMethod, + kNvfp4Static: Nvfp4OnlineMoEMethod, } diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py new file mode 100644 index 000000000000..872643ba2666 --- /dev/null +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch.nn import Module + +from vllm._custom_ops import scaled_fp4_quant +from vllm.model_executor.layers.fused_moe import RoutedExperts +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + convert_to_nvfp4_moe_kernel_format, + make_nvfp4_moe_kernel, + make_nvfp4_moe_quant_config, + select_nvfp4_moe_backend, +) +from vllm.model_executor.layers.quantization.online.moe_base import ( + OnlineMoEMethodBase, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( + FLOAT4_E2M1_MAX, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform + +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max + + +def _quantize_moe_weight_to_nvfp4( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize stacked MoE expert weights ``(E, N, K)`` to NVFP4. + + One FP32 global scale per expert plus per-block (group-16) FP8 scales, + matching the ModelOpt NVFP4 checkpoint layout. Returns packed FP4 weights + ``(E, N, K // 2)``, block scales ``(E, N, K // 16)``, and the per-expert + global scale ``(E,)`` stored as ``amax / (fp4_max * fp8_max)``. + """ + assert weight.dim() == 3, f"expected 3D expert weights, got {weight.shape}" + num_experts, n, k = weight.shape + assert k % 16 == 0, f"last dim must be a multiple of 16, got {k}" + + amax = weight.abs().amax(dim=(1, 2)).to(torch.float32).clamp_min(1e-8) + global_scale = (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) / amax + weight_scale_2 = (1.0 / global_scale).to(torch.float32) + + # scaled_fp4_quant(w, g) == scaled_fp4_quant(w * g, 1), so fold each + # expert's scale in and quantize all experts in one call (fp32 to keep the + # large scale precise), rather than looping per expert. + scaled = (weight.float() * global_scale[:, None, None]).to(weight.dtype) + scaled = scaled.reshape(-1, k) + one = torch.ones((), device=weight.device, dtype=torch.float32) + qweight, block_scale = scaled_fp4_quant(scaled, one, is_sf_swizzled_layout=False) + return ( + qweight.reshape(num_experts, n, k // 2), + block_scale.reshape(num_experts, n, k // 16), + weight_scale_2, + ) + + +class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase): + """Online NVFP4 MoE quantization with per-token activation scales. + + Quantizes fp16/bf16 expert weights to NVFP4 at load time; the FlashInfer + TRTLLM kernel computes per-token activation scales at runtime. Blackwell + (SM100) only. + """ + + def __init__( + self, + *, + layer: torch.nn.Module, + ): + if not current_platform.is_device_capability_family(100): + raise ValueError( + "nvfp4_per_token online quantization requires a Blackwell (SM100) GPU." + ) + super().__init__(layer.moe_config) + self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( + config=self.moe, + weight_key=kNvfp4Static, + activation_key=kNvfp4Dynamic, + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + self._quantize_weights(layer) + self._setup_kernel(layer) + + layer._already_called_process_weights_after_loading = True + + def _quantize_weights(self, layer: Module) -> None: + w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w13_weight) + w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(layer.w2_weight) + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w13_weight_scale_2", w13_scale_2) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w2_weight_scale", w2_scale) + replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) + + # Neutral (1.0) activation global scales: the kernel derives per-token + # scales at runtime, so the output scalars reduce to the weight scales. + ones = torch.ones(layer.num_experts, device=w13.device, dtype=torch.float32) + replace_parameter(layer, "w13_input_scale", ones) + replace_parameter(layer, "w2_input_scale", ones.clone()) + + def _setup_kernel(self, layer: RoutedExperts) -> None: + ( + w13, + w13_scale, + w13_scale_2, + a13_scale, + w2, + w2_scale, + w2_scale_2, + a2_scale, + ) = convert_to_nvfp4_moe_kernel_format( + nvfp4_backend=self.nvfp4_backend, + layer=layer, + w13=layer.w13_weight, + w13_scale=layer.w13_weight_scale, + w13_scale_2=layer.w13_weight_scale_2, + a13_scale=layer.w13_input_scale, + w2=layer.w2_weight, + w2_scale=layer.w2_weight_scale, + w2_scale_2=layer.w2_weight_scale_2, + a2_scale=layer.w2_input_scale, + is_act_and_mul=self.moe.is_act_and_mul, + ) + + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w13_weight_scale_2", w13_scale_2) + replace_parameter(layer, "w13_input_scale", a13_scale) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w2_weight_scale", w2_scale) + replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) + replace_parameter(layer, "w2_input_scale", a2_scale) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.experts_cls is not None + self.moe_kernel = make_nvfp4_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + backend=self.nvfp4_backend, + routing_tables=layer._expert_routing_tables(), + layer=layer, + per_token_activation=True, + ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) + + def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: + return make_nvfp4_moe_quant_config( + backend=self.nvfp4_backend, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_scale_2=layer.w13_weight_scale_2, + w2_scale_2=layer.w2_weight_scale_2, + a13_scale=layer.w13_input_scale, + a2_scale=layer.w2_input_scale, + swiglu_limit=getattr(layer, "swiglu_limit", None), + layer=layer, + ) diff --git a/vllm/model_executor/layers/quantization/quark/quark.py b/vllm/model_executor/layers/quantization/quark/quark.py index fbd61e28cd2a..694116827c54 100644 --- a/vllm/model_executor/layers/quantization/quark/quark.py +++ b/vllm/model_executor/layers/quantization/quark/quark.py @@ -9,7 +9,10 @@ from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import RoutedExperts +from vllm.model_executor.layers.fused_moe import ( + RoutedExperts, + UnquantizedFusedMoEMethod, +) from vllm.model_executor.layers.linear import ( LinearBase, LinearMethodBase, @@ -146,8 +149,13 @@ def get_quant_method( # Check if the layer is skipped for quantization. exclude_layers = cast(list[str], self.quant_config.get("exclude")) if should_ignore_layer( - prefix, ignore=exclude_layers, fused_mapping=self.packed_modules_mapping + prefix, + ignore=exclude_layers, + fused_mapping=self.packed_modules_mapping, + check_children=isinstance(layer, RoutedExperts), ): + if isinstance(layer, RoutedExperts): + return UnquantizedFusedMoEMethod(layer.moe_config) if ( "self_attn" not in prefix # only quantize attention projections or not getattr(self, "dynamic_mxfp4_quant", False) diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index 70a7e81cc455..ea63d1bcf122 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -9,7 +9,7 @@ import torch import torch.nn.functional as F -from vllm._aiter_ops import rocm_aiter_ops +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( dequant_mxfp4, @@ -36,19 +36,15 @@ logger = init_logger(__name__) -try: - from aiter.ops.shuffle import shuffle_weight - from aiter.ops.triton.gemm_afp4wfp4 import ( - gemm_afp4wfp4, - gemm_afp4wfp4_preshuffled_weight_scales, - ) - from aiter.ops.triton.quant import dynamic_mxfp4_quant - +# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP +# which can force the engine core to spawn instead of fork. +# is_aiter_found_and_supported() checks platform + arch + library availability via +# find_spec/amdsmi, so it stays HIP-free. +# Actual aiter imports are deferred to the functions/methods that need them, +# where HIP initialization is expected. +if is_aiter_found_and_supported(): from vllm.utils.torch_utils import direct_register_custom_op - if rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled(): - from aiter import gemm_a4w4, per_1x32_f4_quant_hip - def gemm_with_dynamic_quant( x: torch.Tensor, weight: torch.Tensor, @@ -57,6 +53,15 @@ def gemm_with_dynamic_quant( out_dtype: torch.dtype | None = torch.bfloat16, x_scales: torch.Tensor | None = None, ) -> torch.Tensor: + from aiter.ops.triton.gemm_afp4wfp4 import ( + gemm_afp4wfp4, + gemm_afp4wfp4_preshuffled_weight_scales, + ) + from aiter.ops.triton.quant import dynamic_mxfp4_quant + + if rocm_use_aiter_fp4_asm_gemm: + from aiter import gemm_a4w4, per_1x32_f4_quant_hip + M = x.shape[0] N = weight.shape[0] K = weight.shape[1] @@ -137,13 +142,12 @@ def gemm_with_dynamic_quant_fake( fake_impl=gemm_with_dynamic_quant_fake, dispatch_key=current_platform.dispatch_key, ) -except (ImportError, AttributeError, RuntimeError): - if current_platform.is_rocm(): - logger.warning( - "AITER is not found or QuarkOCP_MX is not supported on the current " - "platform. QuarkOCP_MX quantization will not be available." - ) - dynamic_mxfp4_quant = gemm_afp4wfp4 = None +elif current_platform.is_rocm(): + logger.warning( + "AITER is not found or not supported on the current platform, " + "QuarkOCP_MX will fall back to emulation." + "Native MXFP4/MXFP6 acceleration will not be available." + ) class QuarkOCP_MX(QuarkScheme): @@ -211,8 +215,8 @@ def __init__( rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() ) - if not self.emulate and (dynamic_mxfp4_quant is None or gemm_afp4wfp4 is None): - # Currently need these kernels if not emulating + if not self.emulate and not is_aiter_found_and_supported(): + # Currently need AITER kernels if not emulating raise NotImplementedError( f"{self.__class__.__name__} requires AITER to be installed " "for non-emulation mode! Please refer to " @@ -261,6 +265,8 @@ def get_min_capability(cls) -> int: def process_dynamic_mxfp4_weights_after_loading( self, layer: torch.nn.Module ) -> None: + from aiter.ops.triton.quant import dynamic_mxfp4_quant + w_q, w_s = dynamic_mxfp4_quant(layer.weight) layer.weight_scale = torch.nn.Parameter(w_s.T.contiguous(), requires_grad=False) layer.weight = torch.nn.Parameter(w_q, requires_grad=False) @@ -279,6 +285,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if self.dynamic_mxfp4_quant: self.process_dynamic_mxfp4_weights_after_loading(layer) elif self.rocm_use_aiter_fp4_asm_gemm: + from aiter.ops.shuffle import shuffle_weight + # shuffle weight scale weight_scale_shuffle = layer.weight_scale.data sm, sn = weight_scale_shuffle.shape diff --git a/vllm/model_executor/layers/quantization/quark/utils.py b/vllm/model_executor/layers/quantization/quark/utils.py index ee55e5d39e70..8f6ce4d91cb6 100644 --- a/vllm/model_executor/layers/quantization/quark/utils.py +++ b/vllm/model_executor/layers/quantization/quark/utils.py @@ -27,10 +27,26 @@ def should_ignore_layer( layer_name: str | None, ignore: Iterable[str], fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), + *, + check_children: bool = False, ) -> bool: if layer_name is None: return False + # MoE layers are currently all-or-nothing: if any child is ignored, + # the parent layer must be ignored as well. For example, the + # amd/GLM-5.2-MXFP4 config ignores children like + # model.layers.78.mlp.experts.*.down_proj, while the layer checked + # here is the parent model.layers.N.mlp.experts. + # See: + # https://huggingface.co/amd/GLM-5.2-MXFP4/blob/main/config.json#L793-L795 + if check_children and any( + target == layer_name or target.startswith(layer_name + ".") + for target in ignore + if not target.startswith("re:") + ): + return True + # layer_name = model.layers.0.self_attn.qkv_proj # proj_name = qkv_proj proj_name = layer_name.split(".")[-1] diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 2fd21c2dff8d..6f0e237785e0 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -353,7 +353,13 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( layer.moe_config.hidden_dim = padded_hidden # Align weights for FI NVFP4 MoE kernels. - min_alignment = 16 if is_gated else 128 + # FlashInfer's TRT-LLM block-scale shuffle asserts the gate/up row dim + # (= up_mult * padded_intermediate, up_mult=2 when gated) is a multiple of + # 128. So gated needs padded_intermediate % 64 (2*64=128); the old value 16 + # left 2*intermediate a multiple of only 32, so an NVFP4 MoE whose rank-local + # intermediate is not 128-aligned at TP>1 (e.g. Gemma-4-26B-A4B at tp4) hit + # `assert M % 128 == 0`. Padded rows are zero -> outputs unchanged. + min_alignment = 64 if is_gated else 128 w13, w13_scale, w2, w2_scale, padded_intermediate = ( align_fp4_moe_weights_for_fi( w13, w13_scale, w2, w2_scale, is_act_and_mul, min_alignment diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index 4ed7adcf43c4..e08e80ec4bb7 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -441,6 +441,7 @@ def prepare_humming_layer( input_quant_config: dict | None = None, ): from vllm.utils.humming import ( + BaseInputSchema, BaseWeightSchema, HummingInputSchema, HummingMethod, @@ -448,7 +449,7 @@ def prepare_humming_layer( weight_schema = BaseWeightSchema.from_config(quant_config) if input_quant_config is not None: - input_schema = HummingInputSchema.from_config(input_quant_config) + input_schema = BaseInputSchema.from_config(input_quant_config) else: input_schema = HummingInputSchema() @@ -463,13 +464,19 @@ def prepare_humming_layer( shape_k_stacks = [input_size_per_partition] shape_n_stacks = layer.output_partition_sizes - # Step 1: convert weight to humming standard format + # Step 1: convert weight and input schemas to humming standard format weight_schema, tensors = weight_schema.convert_humming( tensors=dict(layer.named_parameters()), shape_n_stacks=shape_n_stacks, shape_k_stacks=shape_k_stacks, param_dtype=layer.params_dtype, ) + input_schema, _ = input_schema.convert_humming( + tensors={}, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=layer.params_dtype, + ) layer.weight_schema = weight_schema diff --git a/vllm/model_executor/models/deepseek_eagle3.py b/vllm/model_executor/models/deepseek_eagle3.py index 127734808dc8..71ce971d530b 100644 --- a/vllm/model_executor/models/deepseek_eagle3.py +++ b/vllm/model_executor/models/deepseek_eagle3.py @@ -12,7 +12,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig, get_current_vllm_config -from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor @@ -36,8 +35,6 @@ process_eagle_weight, ) -logger = init_logger(__name__) - class DeepseekV2Eagle3DecoderLayer(nn.Module): """ diff --git a/vllm/model_executor/models/deepseek_mtp.py b/vllm/model_executor/models/deepseek_mtp.py index da7f77cae0f9..a5dd6ae8f9f9 100644 --- a/vllm/model_executor/models/deepseek_mtp.py +++ b/vllm/model_executor/models/deepseek_mtp.py @@ -11,7 +11,6 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import tensor_model_parallel_all_gather -from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -38,16 +37,15 @@ ) from .utils import get_pp_missing_layer_names, maybe_prefix -logger = init_logger(__name__) - def _restore_full_token_layout_if_needed( hidden_states: torch.Tensor, residual: torch.Tensor, num_tokens: int, + is_sequence_parallel: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Restore full token rows for the MTP proposer after SP MoE layers.""" - if hidden_states.shape[0] == num_tokens: + if not is_sequence_parallel and hidden_states.shape[0] == num_tokens: return hidden_states, residual combined_states = torch.cat([hidden_states, residual], dim=-1) @@ -142,6 +140,7 @@ def forward( hidden_states, residual, positions.shape[0], + is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, ) hidden_states = residual + hidden_states # pre-final-norm (logits hidden) # Recycle the post-final-norm hidden into the next draft step. diff --git a/vllm/model_executor/models/gemma4_dspark.py b/vllm/model_executor/models/gemma4_dspark.py new file mode 100644 index 000000000000..792f10884296 --- /dev/null +++ b/vllm/model_executor/models/gemma4_dspark.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gemma4 DSpark draft model for speculative decoding.""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm import _custom_ops as ops +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer +from .qwen3_dflash import DFlashQwen3Model +from .qwen3_dspark import DSparkMarkovHead, Qwen3DSparkForCausalLM +from .utils import extract_layer_index, maybe_prefix + + +class Gemma4DSparkAttention(Gemma4MTPAttention): + """Gemma4 attention with its own KV cache and K/V projections.""" + + def __init__( + self, + config, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + is_full = config.layer_types[extract_layer_index(prefix)] == "full_attention" + head_dim = ( + getattr(config, "global_head_dim", config.head_dim) + if is_full + else config.head_dim + ) + use_k_eq_v = is_full and getattr(config, "attention_k_eq_v", False) + num_kv_heads = ( + getattr(config, "num_global_key_value_heads", config.num_key_value_heads) + if use_k_eq_v + else config.num_key_value_heads + ) + super().__init__( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + max_position_embeddings=config.max_position_embeddings, + cache_config=cache_config, + quant_config=quant_config, + attn_logits_soft_cap=getattr(config, "attn_logit_softcapping", None), + prefix=prefix, + ) + self.is_kv_shared_layer = False + self.use_k_eq_v = use_k_eq_v + self.kv_size = self.num_kv_heads * self.head_dim + attn_bias = getattr(config, "attention_bias", False) + self.k_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=attn_bias, + quant_config=quant_config, + prefix=f"{prefix}.k_proj", + ) + self.v_proj = ( + None + if use_k_eq_v + else ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=attn_bias, + quant_config=quant_config, + prefix=f"{prefix}.v_proj", + ) + ) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.v_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps, has_weight=False) + + def _kv_proj( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + k, _ = self.k_proj(hidden_states) + k_normed = self.k_norm(k.unflatten(-1, (self.num_kv_heads, self.head_dim))) + v_src = k if self.use_k_eq_v else self.v_proj(hidden_states)[0] + v_normed = self.v_norm(v_src.unflatten(-1, (self.num_kv_heads, self.head_dim))) + return k_normed.flatten(-2, -1), v_normed.flatten(-2, -1) + + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor, **kwargs + ) -> torch.Tensor: + q, _ = self.q_proj(hidden_states) + q = self.q_norm(q.unflatten(-1, (self.num_heads, self.head_dim))).flatten( + -2, -1 + ) + k, v = self._kv_proj(hidden_states) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class Gemma4DSparkDecoderLayer(Gemma4MTPDecoderLayer): + """Gemma4 MTP decoder layer using the KV-owning DSpark attention.""" + + def __init__( + self, + config, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__( + config, cache_config=cache_config, quant_config=quant_config, prefix=prefix + ) + get_current_vllm_config().compilation_config.static_forward_context.pop( + f"{prefix}.self_attn.attn", None + ) + self.self_attn = Gemma4DSparkAttention( + config, cache_config, quant_config, prefix=f"{prefix}.self_attn" + ) + + +@support_torch_compile +class Gemma4DSparkModel(DFlashQwen3Model): + """Gemma4 DSpark draft backbone (Gemma4 layers + DSpark Markov head).""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + nn.Module.__init__(self) + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + self.use_aux_hidden_state = True + self.target_layer_ids = tuple( + getattr(config, "dspark_target_layer_ids", None) or config.target_layer_ids + ) + current_vllm_config = get_current_vllm_config() + cache_config = current_vllm_config.cache_config + quant_config = current_vllm_config.quant_config + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + self.register_buffer( + "normalizer", + torch.tensor(config.hidden_size**0.5, dtype=vllm_config.model_config.dtype), + persistent=False, + ) + self.fc = ReplicatedLinear( + config.hidden_size * len(self.target_layer_ids), + config.hidden_size, + bias=False, + return_bias=False, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "fc"), + ) + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.layers = nn.ModuleList( + Gemma4DSparkDecoderLayer( + config, + cache_config, + quant_config, + prefix=maybe_prefix(prefix, f"layers.{i}"), + ) + for i in range(config.num_hidden_layers) + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + draft_vocab_size = ( + getattr(config, "draft_vocab_size", None) or config.vocab_size + ) + self.markov_head = DSparkMarkovHead( + config.vocab_size, + draft_vocab_size, + config.markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) * self.normalizer + + def _build_fused_kv_buffers(self) -> None: + layers_attn = [layer.self_attn for layer in self.layers] + attn0 = layers_attn[0] + assert all(a.use_k_eq_v for a in layers_attn), ( + "Gemma4 DSpark fused precompute assumes uniform attention_k_eq_v layers" + ) + self._build_context_kv_buffers(layers_attn, attn0.k_proj.bias is not None) + self._rope_head_size = attn0.rotary_emb.head_size + self._rope_cos_sin_cache = attn0.rotary_emb.cos_sin_cache + self._rope_is_neox = attn0.rotary_emb.is_neox_style + self._num_attn_layers = len(layers_attn) + self._kv_size = attn0.kv_size + self._head_dim = attn0.head_dim + self._num_kv_heads = attn0.num_kv_heads + self._rms_norm_eps = attn0.q_norm.variance_epsilon + self._attn_layers = [layer.self_attn.attn for layer in self.layers] + + def _build_context_kv_buffers( + self, layers_attn: list[nn.Module], has_bias: bool + ) -> None: + self._hidden_norm_weight = self.hidden_norm.weight.data + self._fused_k_weight = torch.cat([a.k_proj.weight for a in layers_attn], dim=0) + self._fused_k_bias: torch.Tensor | None = ( + torch.cat([a.k_proj.bias for a in layers_attn], dim=0) if has_bias else None + ) + self._k_norm_weights = torch.stack( + [a.k_norm.weight.data for a in layers_attn], dim=0 + ).contiguous() + # v_norm has no learnable scale; ones matching the K-norm call shape. + self._v_norm_weights = torch.ones( + len(layers_attn), + layers_attn[0].head_dim, + dtype=self._k_norm_weights.dtype, + device=self._k_norm_weights.device, + ) + + def _project_context_kv( + self, + context_states: torch.Tensor, + num_ctx: int, + num_layers: int, + num_kv_heads: int, + head_dim: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Project once via k_proj for all layers. K is raw (the inherited path + # applies k_norm + RoPE); V = v_norm(same projection), no RoPE (k_eq_v). + normed = torch.empty_like(context_states) + ops.rms_norm( + normed, context_states, self._hidden_norm_weight, self._rms_norm_eps + ) + all_k_flat = F.linear(normed, self._fused_k_weight, self._fused_k_bias) + all_k = ( + all_k_flat.view(num_ctx, num_layers, num_kv_heads, head_dim) + .permute(1, 0, 2, 3) + .contiguous() + ) + all_v = torch.empty_like(all_k) + ops.rms_norm(all_v, all_k, self._v_norm_weights, self._rms_norm_eps) + return all_k, all_v + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + hidden_states = ( + self.embed_input_ids(input_ids) if input_embeds is None else input_embeds + ) + for layer in self.layers: + hidden_states, _ = layer(positions, hidden_states, None) + return self.norm(hidden_states) + + +class Gemma4DSparkForCausalLM(Qwen3DSparkForCausalLM): + """Gemma4 DSpark speculator over a self-contained draft checkpoint.""" + + dspark_shares_target_embeddings = False + packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]} + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + nn.Module.__init__(self) + assert vllm_config.speculative_config is not None + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + self.model = Gemma4DSparkModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor( + self.config.vocab_size, + soft_cap=getattr(self.config, "final_logit_softcapping", None), + ) + self.draft_id_to_target_id = None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked = [("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1)] + params = dict(self.named_parameters()) + params.update(dict(self.named_buffers())) + loaded: set[str] = set() + for name, w in weights: + if "confidence_head" in name: + continue + if "lm_head" not in name: + name = "model." + name + for pn, wn, shard in stacked: + if wn in name and (mapped := name.replace(wn, pn)) in params: + params[mapped].weight_loader(params[mapped], w, shard) + loaded.add(mapped) + break + else: + if name in params: + p = params[name] + getattr(p, "weight_loader", default_weight_loader)(p, w) + loaded.add(name) + self.model._build_fused_kv_buffers() + return loaded diff --git a/vllm/model_executor/models/llama_eagle3.py b/vllm/model_executor/models/llama_eagle3.py index 2d859bd4918b..549e8b7bf63a 100644 --- a/vllm/model_executor/models/llama_eagle3.py +++ b/vllm/model_executor/models/llama_eagle3.py @@ -60,7 +60,7 @@ def __init__( self.self_attn.total_num_kv_heads, bias=qkv_bias, quant_config=quant_config, - prefix=maybe_prefix(prefix, "qkv_proj"), + prefix=maybe_prefix(prefix, "self_attn.qkv_proj"), ) self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 9a93a6065adb..00382cf0a878 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -138,42 +138,44 @@ def check_interleaved_audio_video( num_audio: int, ) -> bool: """ - Check if video and audio positions are interleaved in the multimodal region. - - Returns True only for the use_audio_in_video=True case, where video and - audio tokens alternate within a single contiguous region with no gaps. - - A simple range-overlap check produces false positives when multiple - non-interleaved requests are batched together: audio tokens from request N - fall between video tokens from request N and request N+1, making the - global ranges overlap even though each individual request is non-interleaved. - - To distinguish true interleaving from this batching artefact we require - that every position in the combined [first_VA, last_VA] range is occupied - by either a video or an audio token (no text/image gaps). + Check if video and audio positions are interleaved in any per-video span. + + For use_audio_in_video=True, each video placeholder is expanded into one + local span containing only video/audio pad tokens, bounded by non-pad + tokens such as audio_start/audio_end. Check each contiguous V/A span + independently instead of requiring all V/A tokens in the whole sequence to + form one global dense range; otherwise multi-video requests can be + misclassified as non-interleaved because of the boundary tokens between + videos. """ if num_video == 0 or num_audio == 0: return False - video_pos = is_video.nonzero(as_tuple=True)[0] - audio_pos = is_audio.nonzero(as_tuple=True)[0] - - # Quick range-overlap pre-check (necessary but not sufficient). - if not ( - video_pos[0].item() < audio_pos[-1].item() - and audio_pos[0].item() < video_pos[-1].item() - ): + va_pos = (is_video | is_audio).nonzero(as_tuple=True)[0] + if len(va_pos) == 0: return False - # Density check: for true use_audio_in_video interleaving every position - # in the combined span is a video or audio token. Batched non-interleaved - # requests have text/image tokens between the per-request V and A blocks. - # combined_start/end encompass all V/A tokens, so num_video + num_audio - # equals the number of V/A tokens in range; compare directly to span size. - combined_start = min(video_pos[0].item(), audio_pos[0].item()) - combined_end = max(video_pos[-1].item(), audio_pos[-1].item()) - total_in_range = combined_end - combined_start + 1 - return (num_video + num_audio) == total_in_range + span_start = 0 + for span_end in range(1, len(va_pos) + 1): + is_last = span_end == len(va_pos) + if not is_last and va_pos[span_end].item() == va_pos[span_end - 1].item() + 1: + continue + + span = va_pos[span_start:span_end] + span_is_video = is_video[span] + span_is_audio = is_audio[span] + if span_is_video.any() and span_is_audio.any(): + video_offsets = span_is_video.nonzero(as_tuple=True)[0] + audio_offsets = span_is_audio.nonzero(as_tuple=True)[0] + if ( + video_offsets[0].item() < audio_offsets[-1].item() + and audio_offsets[0].item() < video_offsets[-1].item() + ): + return True + + span_start = span_end + + return False def merge_interleaved_embeddings( @@ -182,58 +184,69 @@ def merge_interleaved_embeddings( is_video: torch.Tensor, is_audio: torch.Tensor, is_multimodal: torch.Tensor, - num_video: int, - num_audio: int, ) -> torch.Tensor: """ Merge embeddings for interleaved audio-in-video sequences. When use_audio_in_video=True, video and audio tokens are interleaved in the token sequence, but embeddings are provided as separate contiguous - tensors (video first, then audio). This function reorders video and audio - embeddings to match sequence position order and scatters them efficiently. + tensors. This function scatters each modality by the ``modality`` attribute + attached to each embedding tensor (set during encoder gather) and also + supports image embeddings in the same interleaved request. Args: inputs_embeds: The input embeddings tensor to merge into. - multimodal_embeddings: List of embedding tensors (video, audio, other). + multimodal_embeddings: List of embedding tensors (video, audio, image). is_video: Boolean mask for video token positions. is_audio: Boolean mask for audio token positions. is_multimodal: Boolean mask for all multimodal token positions. - num_video: Total count of video tokens. - num_audio: Total count of audio tokens. Returns: The merged inputs_embeds tensor with multimodal embeddings scattered to their correct positions. """ - # Categorize embeddings by modality based on token counts. - # Embeddings come grouped by modality but order varies (e.g., image, video, audio - # or video, audio depending on input kwargs order). - video_embeds: list[torch.Tensor] = [] - audio_embeds: list[torch.Tensor] = [] - other_embeds: list[torch.Tensor] = [] - video_remaining = num_video - audio_remaining = num_audio - - for emb in multimodal_embeddings: - n = emb.shape[0] - if video_remaining > 0 and n <= video_remaining: - video_embeds.append(emb) - video_remaining -= n - elif audio_remaining > 0 and n <= audio_remaining: - audio_embeds.append(emb) - audio_remaining -= n - else: - other_embeds.append(emb) + from vllm.multimodal.utils import get_mm_embedding_modalities + + def _merge_embedding_group( + mask: torch.Tensor, + embeddings: Sequence[torch.Tensor], + modality: str, + ) -> None: + num_expected_tokens = mask.sum().item() + num_actual_tokens = sum(embedding.shape[0] for embedding in embeddings) + if num_actual_tokens != num_expected_tokens: + raise ValueError( + f"Attempted to assign {num_actual_tokens} {modality} tokens " + f"to {num_expected_tokens} placeholders" + ) + if embeddings: + inputs_embeds[mask] = torch.cat(list(embeddings), dim=0) + + # Qwen Omni multimodal placeholders are image/video/audio; after excluding + # video and audio positions, the remaining multimodal positions are images. + is_image = is_multimodal & ~is_video & ~is_audio + + embedding_modalities = get_mm_embedding_modalities(multimodal_embeddings) + + video_embeds = [ + embedding + for embedding, modality in zip(multimodal_embeddings, embedding_modalities) + if modality == "video" + ] + audio_embeds = [ + embedding + for embedding, modality in zip(multimodal_embeddings, embedding_modalities) + if modality == "audio" + ] + image_embeds = [ + embedding + for embedding, modality in zip(multimodal_embeddings, embedding_modalities) + if modality == "image" + ] - # Scatter each modality to its positions - if video_embeds: - inputs_embeds[is_video] = torch.cat(video_embeds, dim=0) - if audio_embeds: - inputs_embeds[is_audio] = torch.cat(audio_embeds, dim=0) - if other_embeds: - other_mask = is_multimodal & ~is_video & ~is_audio - inputs_embeds[other_mask] = torch.cat(other_embeds, dim=0) + _merge_embedding_group(is_video, video_embeds, "video") + _merge_embedding_group(is_audio, audio_embeds, "audio") + _merge_embedding_group(is_image, image_embeds, "image") return inputs_embeds @@ -1497,8 +1510,6 @@ def embed_input_ids( is_video, is_audio, is_multimodal, - num_video, - num_audio, ) # Default: standard merge (no interleaving), same as parent class diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 726329c7805f..32a622567ceb 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -1833,9 +1833,11 @@ def embed_input_ids( # both the deepstack path and the final embedding merge. video_token_id = self.config.video_token_id audio_token_id = self.config.audio_token_id + image_token_id = self.config.image_token_id input_ids_cpu = input_ids.cpu() is_video = is_multimodal & (input_ids_cpu == video_token_id) is_audio = is_multimodal & (input_ids_cpu == audio_token_id) + is_image = is_multimodal & (input_ids_cpu == image_token_id) num_video = is_video.sum().item() num_audio = is_audio.sum().item() @@ -1856,9 +1858,9 @@ def embed_input_ids( multimodal_embeddings_multiscale = [] if is_interleaved: - # Use input_ids-based mask for correct vision positions - # when audio and video tokens are interleaved. - is_vision = is_video.clone() + # Use input_ids-based mask for all vision positions when + # audio and video tokens are interleaved. + is_vision = is_video | is_image else: is_vision = torch.zeros_like(is_multimodal) mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0] @@ -1917,8 +1919,6 @@ def embed_input_ids( is_video, is_audio, is_multimodal, - num_video, - num_audio, ) # Default: standard merge (no interleaving), same as parent class. diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index cab03a8a6a26..d28aac974b54 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -159,6 +159,11 @@ "vllm.models.minimax_m3", "MiniMaxM3SparseForCausalLM", ), + "InklingForCausalLM": ("vllm.models.inkling", "InklingForCausalLM"), + "InklingForConditionalGeneration": ( + "vllm.models.inkling", + "InklingForConditionalGeneration", + ), "Ministral3ForCausalLM": ("mistral", "MistralForCausalLM"), "MistralForCausalLM": ("mistral", "MistralForCausalLM"), "MistralLarge3ForCausalLM": ("mistral_large_3", "MistralLarge3ForCausalLM"), @@ -604,6 +609,7 @@ "DSparkDraftModel": ("vllm.models.deepseek_v4", "DSparkDeepseekV4ForCausalLM"), "Qwen3DSparkModel": ("qwen3_dspark", "Qwen3DSparkForCausalLM"), "DFlashLagunaForCausalLM": ("laguna_dflash", "DFlashLagunaForCausalLM"), + "Gemma4DSparkModel": ("gemma4_dspark", "Gemma4DSparkForCausalLM"), "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"), @@ -625,6 +631,7 @@ "DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"), "MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"), "BailingMoeV25MTPModel": ("bailing_moe_mtp", "BailingMoeV25MTPModel"), + "InklingMTPModel": ("vllm.models.inkling", "InklingMTP"), "Gemma4MTPModel": ("gemma4_mtp", "Gemma4MTP"), "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"), "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"), diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index c964cb721d65..24cad1da20d2 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -587,22 +587,23 @@ def init_parameters(self, module: nn.Module, dtype: torch.dtype | None = None): self.model: "PreTrainedModel" = AutoModel.from_config(...) ``` """ + dtype = dtype or self.model_config.dtype + device = self.device_config.device - def _init_parameters(module: nn.Module, dtype: torch.dtype | None): + def _init_parameters(module: nn.Module): for name, param in module.named_parameters(recurse=False): - if param.device == torch.device("meta"): - new_param = nn.Parameter( - torch.empty_like( - param.data, - dtype=dtype or self.model_config.dtype, - device=self.device_config.device, - ) - ) - setattr(module, name, new_param) + # Already on device, nothing to do + if param.device != torch.device("meta"): + continue + # Already a vLLM parameter, nothing to do + if hasattr(param, "weight_loader"): + continue + data = torch.empty_like(param.data, dtype=dtype, device=device) + setattr(module, name, nn.Parameter(data=data)) for child in module.children(): - _init_parameters(child, dtype) + _init_parameters(child) - _init_parameters(module, dtype) + _init_parameters(module) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.get_input_embeddings()(input_ids) diff --git a/vllm/model_executor/warmup/cutedsl_warmup.py b/vllm/model_executor/warmup/cutedsl_warmup.py index 5978e91a6eac..991eab22b8a8 100644 --- a/vllm/model_executor/warmup/cutedsl_warmup.py +++ b/vllm/model_executor/warmup/cutedsl_warmup.py @@ -10,7 +10,9 @@ from dataclasses import dataclass import torch +from tqdm import tqdm +from vllm.distributed import is_global_first_rank from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.tracing import instrument @@ -76,6 +78,8 @@ def _compile_cutedsl_warmup_units( compile_units: Iterable[CuTeDSLCompileUnit], ) -> int: compiled = 0 + if is_global_first_rank(): + compile_units = tqdm(compile_units, desc="Compiling CuTeDSL kernels") with torch.inference_mode(): for unit in compile_units: unit.compile() diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py index 5b4900b50ce2..6ca8c94fff0a 100644 --- a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -15,7 +15,6 @@ from vllm.logger import init_logger from vllm.tracing import instrument -from vllm.utils.math_utils import cdiv logger = init_logger(__name__) @@ -39,23 +38,6 @@ ) -def _compute_mhc_pre_num_split( - *, - num_tokens: int, - hidden_size: int, - hc_mult: int, - num_sms: int, -) -> int: - block_k = 64 - block_m = 64 - k = hc_mult * hidden_size - grid_size = cdiv(num_tokens, block_m) - split_k = num_sms // grid_size - num_block_k = cdiv(k, block_k) - split_k = min(split_k, num_block_k // 4) - return max(split_k, 1) - - def _normalize_token_sizes( token_sizes: Iterable[int], *, diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index c55ad148b2da..d4b9a735d2cd 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -43,6 +43,7 @@ logger = init_logger(__name__) _LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = ( + (6144, 264), # Inkling (7168, 256), # DSV3 (7168, 384), # DSV4-Pro (14400, 256), # DSV4-Flash diff --git a/vllm/model_executor/warmup/qwen_triton_warmup.py b/vllm/model_executor/warmup/qwen_triton_warmup.py index de1a985c6f79..9b7d76207b41 100644 --- a/vllm/model_executor/warmup/qwen_triton_warmup.py +++ b/vllm/model_executor/warmup/qwen_triton_warmup.py @@ -274,7 +274,7 @@ def _warm_causal_conv1d_fwd_kernel( def _warm_fused_post_conv_kernel( device: torch.device, config: _QwenGDNWarmupConfig ) -> None: - from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import ( + from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import ( # noqa: E501 fused_post_conv_prep, ) @@ -304,7 +304,7 @@ def _warm_fused_sigmoid_gating_delta_rule_update_kernel( device: torch.device, config: _QwenGDNWarmupConfig, ) -> None: - from vllm.model_executor.layers.fla.ops.fused_sigmoid_gating import ( + from vllm.third_party.flash_linear_attention.ops.fused_sigmoid_gating import ( fused_sigmoid_gating_delta_rule_update, ) diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index 22ced9fa94d8..353aedc8ceed 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -7,7 +7,11 @@ import torch from vllm.config import VllmConfig -from vllm.distributed import get_pp_group +from vllm.distributed import ( + get_pp_group, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -32,6 +36,7 @@ is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, + sequence_parallel_chunk, ) from vllm.sequence import IntermediateTensors @@ -39,6 +44,18 @@ from .fused_ops import fused_allreduce_rms_norm +def _all_gather_sp_states( + hidden_states: torch.Tensor, + residual: torch.Tensor, + num_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + # combine hidden_states and residual and all gather once + combined_states = torch.cat([hidden_states, residual], dim=-1) + combined_states = tensor_model_parallel_all_gather(combined_states, 0)[:num_tokens] + hidden_states, residual = combined_states.chunk(2, dim=-1) + return hidden_states, residual.contiguous() + + class DeepseekV32DecoderLayer(torch.nn.Module): def __init__( self, @@ -92,6 +109,12 @@ def __init__( prefix=f"{prefix}.mlp", reduce_results=False, ) + self.use_sequence_parallel_moe = ( + parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size == 1 + and isinstance(self.mlp, DeepseekV2MoE) + ) + self.tp_size = parallel_config.tensor_parallel_size self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -104,25 +127,53 @@ def forward( hidden_states: torch.Tensor, residual: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: + full_num_tokens = positions.shape[0] + input_is_sequence_parallel = ( + self.use_sequence_parallel_moe + and residual is not None + and hidden_states.shape[0] != full_num_tokens + ) + if residual is None: # First layer: hidden_states is the (already reduced) embedding. residual = hidden_states hidden_states = self.input_layernorm(hidden_states) + elif input_is_sequence_parallel: + hidden_states, residual = self.input_layernorm(hidden_states, residual) else: # The previous layer's MLP/MoE output is left un-reduced; fuse its # all-reduce into this input_layernorm. hidden_states, residual = fused_allreduce_rms_norm( hidden_states, residual, self.input_layernorm ) - # self_attn's o_proj runs reduce_results=False; fuse its all-reduce with - # the post-attention RMSNorm. + if input_is_sequence_parallel: + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0) + hidden_states = hidden_states[:full_num_tokens] + + # self_attn's o_proj runs reduce_results=False; reduce before RMSNorm. hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) - hidden_states, residual = fused_allreduce_rms_norm( - hidden_states, residual, self.post_attention_layernorm - ) + if self.use_sequence_parallel_moe: + # small trick using minus, eg. -17 % 8 = 7 + sp_pad = (-hidden_states.shape[0]) % self.tp_size + # pad if not divisible by world size + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad)) + hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0) + if not input_is_sequence_parallel: + residual = sequence_parallel_chunk(residual) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + # MLP/MoE runs un-reduced; its all-reduce is fused into the next layer's # input_layernorm (or the model's final norm). - hidden_states = self.mlp(hidden_states) + if self.use_sequence_parallel_moe: + hidden_states = self.mlp(hidden_states, already_sequence_parallel=True) + else: + hidden_states = self.mlp(hidden_states) return hidden_states, residual @@ -206,12 +257,26 @@ def forward( residual = intermediate_tensors["residual"] aux_hidden_states = [] + full_num_tokens = positions.shape[0] for idx, layer in enumerate( islice(self.layers, self.start_layer, self.end_layer), start=self.start_layer, ): + if ( + hidden_states.shape[0] != full_num_tokens + and not layer.use_sequence_parallel_moe + ): + hidden_states, residual = _all_gather_sp_states( + hidden_states, residual, full_num_tokens + ) if idx in self.aux_hidden_state_layers: - aux_hidden_states.append(hidden_states + residual) + aux_hidden_state = hidden_states + residual + if aux_hidden_state.shape[0] != full_num_tokens: + aux_hidden_state = tensor_model_parallel_all_gather( + aux_hidden_state, 0 + ) + aux_hidden_state = aux_hidden_state[:full_num_tokens] + aux_hidden_states.append(aux_hidden_state) hidden_states, residual = layer(positions, hidden_states, residual) if not get_pp_group().is_last_rank: @@ -219,8 +284,15 @@ def forward( {"hidden_states": hidden_states, "residual": residual} ) - # Last layer's MoE output is un-reduced; fuse its all-reduce into norm. - hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm) + if hidden_states.shape[0] != full_num_tokens: + hidden_states, residual = _all_gather_sp_states( + hidden_states, residual, full_num_tokens + ) + hidden_states, _ = self.norm(hidden_states, residual) + else: + hidden_states, _ = fused_allreduce_rms_norm( + hidden_states, residual, self.norm + ) if len(aux_hidden_states) > 0: return hidden_states, aux_hidden_states return hidden_states diff --git a/vllm/models/deepseek_v32/nvidia/mtp.py b/vllm/models/deepseek_v32/nvidia/mtp.py index 8edb3283d07d..d3d8e5aae9c5 100644 --- a/vllm/models/deepseek_v32/nvidia/mtp.py +++ b/vllm/models/deepseek_v32/nvidia/mtp.py @@ -96,11 +96,11 @@ def forward( hidden_states, residual, positions.shape[0], + is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe, ) - # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the - # main model fuses that all-reduce into the next norm, but here the - # recycle hidden is consumed directly, so reduce it now. - hidden_states = tensor_model_parallel_all_reduce(hidden_states) + if not self.mtp_block.use_sequence_parallel_moe: + # Without sequence parallelism, the MoE output is left un-reduced. + hidden_states = tensor_model_parallel_all_reduce(hidden_states) # Recycle the POST-final-norm hidden into the next draft step. The # residual-add is fused into the final RMSNorm so it is computed # exactly once, and the result is returned for both tuple positions: diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index 05662bbd321a..c7dd948f79f8 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -21,10 +21,9 @@ from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP elif current_platform.is_xpu(): + from .xpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment] - - DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment, misc] else: from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, diff --git a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py index 9a5e478e315f..a2085cd220f1 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py @@ -19,6 +19,7 @@ and N_QUANT_BLOCKS ue8m0 bytes. """ +from functools import lru_cache from typing import Any import torch @@ -296,6 +297,360 @@ def _fused_kv_compress_norm_rope_insert_sparse_attn( tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope) +# ============================================================================= +# Split kernels variant of the head=512 compressor (deep cr=128 gather). +# - compress gather: instead of launching one program per token, split along +# the head dimension to maximize CU occupancy. The head dimension split +# does not require cross-group reduction +# - finalize norm rope quant store: same as the single pass kernel due to its +# per-token nature +# Mirrors the CUDA cutedsl split kernel where num_splits is occupancy-targeted. +# Currently only tested and validated on ROCm gfx950 +# ============================================================================= +@lru_cache(maxsize=1) +def _n_cu() -> int: + return torch.cuda.get_device_properties(0).multi_processor_count + + +def _pick_compress_num_splits( + num_actual: int, compress_ratio: int, head_dim: int +) -> int: + """Occupancy-targeted column splits for the cr>=128 head=512 compressor. + + Sizes the per-token fan-out so (estimated computing tokens) * num_splits ~ + #CU, capped by head tiling at a 32-wide min tile, as a power-of-2 divisor of + head_dim. + """ + max_splits = head_dim // 32 + est_compute = max(1, num_actual // compress_ratio) + target = -(-_n_cu() // est_compute) # ceil(#CU / est_compute) + ns = 1 + while ns * 2 <= min(target, max_splits) and head_dim % (ns * 2) == 0: + ns *= 2 + return ns + + +@triton.jit +def _compress_gather_split_sparse_attn( + state_cache_ptr, + state_cache_stride0, + state_cache_stride1, + positions_ptr, + slot_mapping_ptr, + token_to_req_indices_ptr, + block_table_ptr, + block_table_stride, + block_size, + scratch_ptr, + scratch_stride, + HEAD_SIZE: tl.constexpr, + STATE_WIDTH: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + NUM_SPLITS: tl.constexpr, + HEAD_TILE: tl.constexpr, # HEAD_SIZE // NUM_SPLITS +): + """Stage 1: per-(token, head-split) compress gather, write to fp32 scratch + + No-overlap gather (cr>=128) on rows [0, COMPRESS_RATIO) + """ + pid = tl.program_id(0) + token_idx = pid // NUM_SPLITS + split_idx = pid % NUM_SPLITS + + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + start = position - COMPRESS_RATIO + 1 + rows = tl.arange(0, COMPRESS_RATIO) + pos = start + rows + mask_pos = pos >= 0 + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + pos // block_size, + mask=mask_pos, + other=0, + ).to(tl.int64) + block_offsets = pos % block_size + + col = split_idx * HEAD_TILE + tl.arange(0, HEAD_TILE) + row_base = ( + state_cache_ptr + + block_numbers * state_cache_stride0 + + block_offsets * state_cache_stride1 + ) + cmask = mask_pos[:, None] + + score = tl.load( + row_base[:, None] + STATE_WIDTH + col[None, :], + mask=cmask, + other=float("-inf"), + ) + score = tl.softmax(score, dim=0) + kv = tl.load(row_base[:, None] + col[None, :], mask=cmask, other=0.0) + compressed = tl.sum(kv * score, axis=0) # [HEAD_TILE] fp32 + tl.store(scratch_ptr + token_idx * scratch_stride + col, compressed) + + +@triton.jit +def _finalize_norm_rope_quant_store_sparse_attn( + scratch_ptr, + scratch_stride, + positions_ptr, + slot_mapping_ptr, + rms_norm_weight_ptr, + rms_norm_eps, + cos_sin_cache_ptr, + cos_sin_stride, + k_cache_ptr, + kv_slot_mapping_ptr, + kv_cache_block_size, + HEAD_SIZE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + ROPE_HEAD_DIM: tl.constexpr, + FP8_MAX: tl.constexpr, + QUANT_BLOCK: tl.constexpr, + TOKEN_STRIDE: tl.constexpr, + SCALE_DIM: tl.constexpr, + KV_BLOCK_STRIDE: tl.constexpr, +): + """Stage 2: read compressed_kv[512] from scratch buffer, then + RMSNorm + FP8 quant (nope) + RoPE + bf16 store + """ + token_idx = tl.program_id(0) + slot_id = tl.load(slot_mapping_ptr + token_idx) + if slot_id < 0: + return + position = tl.load(positions_ptr + token_idx) + if (position + 1) % COMPRESS_RATIO != 0: + return + + block = tl.arange(0, TRITON_BLOCK_SIZE) + mask = block < HEAD_SIZE + compressed_kv = tl.load( + scratch_ptr + token_idx * scratch_stride + block, mask=mask, other=0.0 + ) + + rms_w = tl.load(rms_norm_weight_ptr + block, mask=mask, other=0.0) + variance = tl.sum(compressed_kv * compressed_kv, axis=0) / HEAD_SIZE + rrms = tl.rsqrt(variance + rms_norm_eps) + normed = compressed_kv * rrms * rms_w + + kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx) + if kv_slot_idx < 0: + return + kv_block_idx = kv_slot_idx // kv_cache_block_size + kv_pos_in_block = kv_slot_idx % kv_cache_block_size + cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE + fp8_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE + scale_ptr = ( + cache_block_ptr + + kv_cache_block_size * TOKEN_STRIDE + + kv_pos_in_block * SCALE_DIM + ) + + NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM + HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2 + N_QUANT_BLOCKS: tl.constexpr = TRITON_BLOCK_SIZE // QUANT_BLOCK + N_NOPE_BLOCKS: tl.constexpr = NOPE_HEAD_DIM // QUANT_BLOCK + INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX + + quant_input = normed.to(tl.bfloat16).to(tl.float32) + quant_2d = tl.reshape(quant_input, (N_QUANT_BLOCKS, QUANT_BLOCK)) + block_absmax = tl.maximum(tl.max(tl.abs(quant_2d), axis=1), 1e-4) + raw_scales = block_absmax * INV_FP8_MAX + exponents = tl.ceil(tl.log2(raw_scales)) + inv_scales = tl.exp2(-exponents) + x_scaled = quant_2d * tl.reshape(inv_scales, (N_QUANT_BLOCKS, 1)) + x_clamped = tl.clamp(x_scaled, -FP8_MAX, FP8_MAX) + x_uint8 = tl.reshape( + x_clamped.to(tl.float8e4nv).to(tl.uint8, bitcast=True), + (TRITON_BLOCK_SIZE,), + ) + tl.store(fp8_ptr + block, x_uint8, mask=block < NOPE_HEAD_DIM) + + scale_idx = tl.arange(0, N_QUANT_BLOCKS) + encoded = tl.maximum(tl.minimum(exponents + 127.0, 255.0), 0.0) + tl.store( + scale_ptr + scale_idx, encoded.to(tl.uint8), mask=scale_idx < N_NOPE_BLOCKS + ) + tl.store(scale_ptr + N_NOPE_BLOCKS, tl.zeros((), dtype=tl.uint8)) + + NUM_PAIRS: tl.constexpr = TRITON_BLOCK_SIZE // 2 + NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2 + even, odd = tl.split(tl.reshape(normed, (NUM_PAIRS, 2))) + pair_idx = tl.arange(0, NUM_PAIRS) + rope_pair_local = pair_idx - NOPE_PAIRS + is_rope_pair = rope_pair_local >= 0 + cs_idx = tl.maximum(rope_pair_local, 0) + compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO + cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride + cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0) + sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0) + new_even = even * cos_v - odd * sin_v + new_odd = odd * cos_v + even * sin_v + result = tl.interleave(new_even, new_odd) + bf16_ptr = (fp8_ptr + NOPE_HEAD_DIM).to(tl.pointer_type(tl.bfloat16)) + rope_local = block - NOPE_HEAD_DIM + is_rope = (block >= NOPE_HEAD_DIM) & mask + tl.store(bf16_ptr + rope_local, result.to(tl.bfloat16), mask=is_rope) + + +def _launch_two_stage_sparse_attn_compressor( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + compress_ratio: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + kv_slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, + head_dim: int, + rope_head_dim: int, + num_actual: int, + compress_scratch: torch.Tensor, +) -> None: + num_splits = _pick_compress_num_splits(num_actual, compress_ratio, head_dim) + head_tile = head_dim // num_splits + scratch = compress_scratch[:num_actual] + _compress_gather_split_sparse_attn[(num_actual * num_splits,)]( + state_cache, + state_cache.stride(0), + state_cache.stride(1), + positions, + slot_mapping, + token_to_req_indices, + block_table, + block_table.stride(0), + block_size, + scratch, + scratch.stride(0), + HEAD_SIZE=head_dim, + STATE_WIDTH=state_width, + COMPRESS_RATIO=compress_ratio, + NUM_SPLITS=num_splits, + HEAD_TILE=head_tile, + ) + _finalize_norm_rope_quant_store_sparse_attn[(num_actual,)]( + scratch, + scratch.stride(0), + positions, + slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + cos_sin_cache.stride(0), + kv_cache, + kv_slot_mapping, + kv_cache.shape[1], + HEAD_SIZE=head_dim, + TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim), + COMPRESS_RATIO=compress_ratio, + ROPE_HEAD_DIM=rope_head_dim, + FP8_MAX=448.0, + QUANT_BLOCK=quant_block, + TOKEN_STRIDE=token_stride, + SCALE_DIM=scale_dim, + KV_BLOCK_STRIDE=kv_cache.stride(0), + ) + + +def compress_norm_rope_store_two_stage_triton( + state_cache: torch.Tensor, + num_actual: int, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_size: int, + state_width: int, + cos_sin_cache: torch.Tensor, + kv_cache: torch.Tensor, + k_cache_metadata: Any, + pdl_kwargs: dict, + head_dim: int, + rope_head_dim: int, + compress_ratio: int, + overlap: bool, + use_fp4_cache: bool, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + quant_block: int, + token_stride: int, + scale_dim: int, + num_decode_tokens: int, + compress_scratch: torch.Tensor, +) -> None: + """Two-stage split compressor dispatch for head=512 cr>=128 (no-overlap) + + Run the occupancy-fanned two-stage split for prefill [num_decodee_tokens:] + to fill the CUs, and use the original single-pass launcher + for decode [0, num_decode_tokens) + """ + num_decodes = min(max(num_decode_tokens, 0), num_actual) + num_prefills = num_actual - num_decodes + if num_prefills > 0: + _launch_two_stage_sparse_attn_compressor( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices[num_decodes:], + positions=positions[num_decodes:], + slot_mapping=slot_mapping[num_decodes:], + block_table=block_table, + block_size=block_size, + state_width=state_width, + compress_ratio=compress_ratio, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + kv_slot_mapping=k_cache_metadata.slot_mapping[num_decodes:], + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + num_actual=num_prefills, + compress_scratch=compress_scratch, + ) + if num_decodes > 0: + compress_norm_rope_store_triton( + state_cache=state_cache, + num_actual=num_decodes, + token_to_req_indices=token_to_req_indices, + positions=positions, + slot_mapping=slot_mapping, + block_table=block_table, + block_size=block_size, + state_width=state_width, + cos_sin_cache=cos_sin_cache, + kv_cache=kv_cache, + k_cache_metadata=k_cache_metadata, + pdl_kwargs=pdl_kwargs, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + compress_ratio=compress_ratio, + overlap=overlap, + use_fp4_cache=use_fp4_cache, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + quant_block=quant_block, + token_stride=token_stride, + scale_dim=scale_dim, + ) + + # ============================================================================= # Indexer path (head=128, all FP8, single quant block) # ============================================================================= diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 24838c237ce5..13f327f6bc19 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.linear import MergedColumnParallelLinear from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import ( compress_norm_rope_store_triton, + compress_norm_rope_store_two_stage_triton, ) from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE from vllm.models.deepseek_v4.common.ops.save_partial_states import ( @@ -27,6 +28,7 @@ CommonAttentionMetadata, MultipleOf, ) +from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, @@ -34,6 +36,12 @@ ) +def _prefer_two_stage_compressor() -> bool: + # Platforms that favor the triton variant of two-stage compressor split. + # Currently only tested on ROCm + return current_platform.is_rocm() + + class CompressorBackend(AttentionBackend): def __init__(self): super().__init__() @@ -81,6 +89,7 @@ class CompressorMetadata: block_size: int token_to_req_indices: torch.Tensor | None = None # [num_tokens] + num_decode_tokens: int | None = None class CompressorMetadataBuilder(AttentionMetadataBuilder): @@ -107,11 +116,17 @@ def build( token_to_req_indices = common_attn_metadata.token_to_req_indices( self.token_to_req_indices ) + num_decode_tokens = None + if _prefer_two_stage_compressor(): + _, _, num_decode_tokens, _ = split_decodes_and_prefills( + common_attn_metadata, decode_threshold=1 + ) return CompressorMetadata( block_table=common_attn_metadata.block_table_tensor.clamp_(min=0), slot_mapping=common_attn_metadata.slot_mapping, block_size=self.block_size, token_to_req_indices=token_to_req_indices, + num_decode_tokens=num_decode_tokens, ) @@ -213,6 +228,25 @@ def __init__( self.overlap = compress_ratio == 4 self.coff = 1 + self.overlap + # The head=512 cr>=128 no-overlap deep gather uses the two-stage + # compressor, which needs an fp32 scratch [max_batched, 512] for + # the intermediate compressed_kv. + # Currently only tested on ROCm + self._use_two_stage_fused_compressor = ( + _prefer_two_stage_compressor() and head_dim == 512 and not self.overlap + ) + self.max_num_batched_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self._compress_scratch: torch.Tensor | None = None + if self._use_two_stage_fused_compressor: + self._compress_scratch = torch.empty( + self.max_num_batched_tokens, + self.head_dim, + dtype=torch.float32, + device=self.device, + ) + state_dtype = torch.float32 self.ape = nn.Parameter( torch.empty( @@ -364,6 +398,15 @@ def forward( store_full_fp8=store_full_fp8, fp8_scale=fp8_scale, ) + elif self._use_two_stage_fused_compressor: + # head=512 cr>=128 (no overlap): two-pass split compressor on the + # prefill suffix, single-pass on the decode prefix. + assert state_metadata.num_decode_tokens is not None + compress_norm_rope_store_fn = compress_norm_rope_store_two_stage_triton + extra_kwargs = { + "num_decode_tokens": state_metadata.num_decode_tokens, + "compress_scratch": self._compress_scratch, + } else: # Indexer path (head_dim == 128) or non-CUDA GPUs (AMD, XPU, etc.). compress_norm_rope_store_fn = compress_norm_rope_store_triton diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 1aaf3f1a141c..4523d1875ebf 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -209,7 +209,7 @@ def build( cm.num_actual_tokens, cm.query_start_loc, cm.seq_lens, - cm.block_table_tensor.clamp(min=0), + cm.block_table_tensor.clamp_(min=0), int(self.kv_cache_spec.storage_block_size), self.compress_ratio, out=self.compressed_slot_mapping_buffer, diff --git a/vllm/models/deepseek_v4/xpu/dspark.py b/vllm/models/deepseek_v4/xpu/dspark.py new file mode 100644 index 000000000000..63333a9c0a2e --- /dev/null +++ b/vllm/models/deepseek_v4/xpu/dspark.py @@ -0,0 +1,436 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DSpark draft model for DeepSeek-V4 on Intel XPU. + +Minimal XPU port of nvidia/dspark.py. Replaces tilelang MHC kernels with +the platform-agnostic custom ops (HCHeadOp, MHCPostOp) already used by the +XPU MTP path, and uses the XPU Triton-based qnorm_rope_kv_fp8_insert for +context KV precomputation. +""" + +from collections.abc import Iterable + +import regex as re +import torch +import torch.nn as nn + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import HCHeadOp, MHCPostOp +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.models.qwen3_dspark import DSparkMarkovHead +from vllm.model_executor.models.utils import maybe_prefix + +from .model import ( + DeepseekV4DecoderLayer, + make_deepseek_v4_expert_params_mapping, +) + +logger = init_logger(__name__) + +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") + + +class DSparkDeepseekV4Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + self.hidden_size = config.hidden_size + self.hc_mult = config.hc_mult + self.hc_eps = config.hc_eps + self.rms_norm_eps = config.rms_norm_eps + self.num_hidden_layers = config.num_hidden_layers + self.target_layer_ids = tuple(config.dspark_target_layer_ids) + + self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 + + # Shared with target (aliased by speculator loading utility). + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + + self.main_proj = ReplicatedLinear( + config.hidden_size * len(self.target_layer_ids), + config.hidden_size, + bias=False, + return_bias=False, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "main_proj"), + ) + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + current_vllm_config = get_current_vllm_config() + self.layers = nn.ModuleList( + [ + DeepseekV4DecoderLayer( + current_vllm_config, + prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), + ) + for i in range(self.num_dspark_layers) + ] + ) + + # Heads: final norm + hc_head, and the Markov head + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + draft_vocab_size = ( + getattr(config, "draft_vocab_size", None) or config.vocab_size + ) + self.markov_head = DSparkMarkovHead( + config.vocab_size, + draft_vocab_size, + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + + # XPU MHC ops (replaces tilelang) + self.mhc_post_op = MHCPostOp() + self.hc_head_op = HCHeadOp() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: + """main_x = main_norm(main_proj(concat of target aux hidden states)).""" + return self.main_norm(self.main_proj(aux_hidden_states)) + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + main_x: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mappings: list[torch.Tensor | None] | None = None, + ) -> None: + """Insert the sliding-window context KV for every draft layer. + + Each layer derives its context KV from the SAME projected target hidden + ``main_x``, via that layer's own wkv + kv_norm + RoPE + quant, then + writes it at the layer's context slots. + """ + for i, layer in enumerate(self.layers): + slot_mapping = ( + None if context_slot_mappings is None else context_slot_mappings[i] + ) + attn = layer.attn + # wkv part of the fused wq_a|wkv projection (q_lora part discarded) + qr_kv, _ = attn.fused_wqa_wkv(main_x) + kv = qr_kv[..., attn.q_lora_rank :] + kv = attn.kv_norm(kv) + if slot_mapping is None: + continue + _insert_context_kv(attn, kv, context_positions, slot_mapping) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). + hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) + + residual = post_mix = res_mix = None + for layer in self.layers: + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + # mhc_post: merge hyper-connection copies + hidden_states = self.mhc_post_op(hidden_states, residual, post_mix, res_mix) + # hc_head: reduces hc copies; return pre-norm head hidden + hidden_states = self.hc_head_op( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + return hidden_states + + +def _insert_context_kv( + attn: nn.Module, + kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. + + On XPU, we reuse the same xpu_qnorm_rope_kv_fp8_insert kernel used in + the forward path, passing a dummy q (result discarded). + """ + from vllm.models.deepseek_v4.xpu.xpu_qnorm_rope_kv_fp8_insert import ( + xpu_qnorm_rope_kv_fp8_insert, + ) + + swa_cache = attn.swa_cache_layer.kv_cache + block_size = attn.swa_cache_layer.block_size + cos_sin_cache = attn.rotary_emb.cos_sin_cache + n_ctx = kv.shape[0] + + # Dummy q — we only care about the KV insert side effect. + dummy_q = torch.empty( + (n_ctx, attn.n_local_heads, attn.head_dim), + dtype=kv.dtype, + device=kv.device, + ) + xpu_qnorm_rope_kv_fp8_insert( + dummy_q, + kv, + swa_cache, + slot_mapping, + positions, + cos_sin_cache, + attn.eps, + block_size, + ) + + +class DSparkDeepseekV4ForCausalLM(nn.Module): + """XPU DSpark draft model entry point for DeepSeek-V4.""" + + has_own_embed_tokens = False + has_own_lm_head = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + self.model = DSparkDeepseekV4Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + # Shared with the target (aliased by the speculator's load utility). + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(self.config.vocab_size) + + # --- Hooks used by the speculator --- + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: + return self.model.combine_hidden_states(aux_hidden_states) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers] + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mappings: list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv( + context_states, context_positions, context_slot_mappings + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Base logits U_k = lm_head(norm(head_hidden)).""" + return self.logits_processor(self.lm_head, self.model.norm(hidden_states)) + + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.compute_logits(hidden_states) + + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: + return draft_ids # full-vocab: draft ids are target ids + + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.embed(token_ids) + + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.bias(markov_embed, self.logits_processor) + + # --- Weight loading --- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load ``mtp.{0,1,2}.*`` draft weights from the target checkpoint.""" + first_layer = self.model.layers[0] + use_mega_moe = first_layer.ffn.use_mega_moe + if use_mega_moe: + expert_mapping = make_deepseek_v4_expert_params_mapping( + self.config.n_routed_experts + ) + else: + expert_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + expert_scale_suffix = ( + ".weight_scale" + if getattr(self.config, "expert_dtype", "fp4") == "fp4" + else ".weight_scale_inv" + ) + + stacked_params_mapping = [ + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_local_head = self.config.num_attention_heads // tp_size + head_start = n_local_head * tp_rank + head_end = n_local_head * (tp_rank + 1) + + for name, loaded_weight in weights: + mapped = self._remap_dspark_name(name) + if mapped is None: + continue + name = mapped + + # .scale -> per-method scale suffix + if name.endswith(".scale"): + suffix = ( + expert_scale_suffix + if _EXPERT_SCALE_RE.search(name) + else ".weight_scale_inv" + ) + name = name.removesuffix(".scale") + suffix + + # Expert weights + if ".experts." in name: + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for param_name, weight_name, expert_id, shard_id in expert_mapping: + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if name_mapped not in params_dict: + continue + param = params_dict[name_mapped] + success = param.weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(name_mapped) + break + continue + + # Stacked params (decoder-layer only) + is_layer_param = name.startswith("model.layers.") + for param_name, weight_name, stacked_shard_id in stacked_params_mapping: + if not is_layer_param or weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name not in params_dict: + break + param = params_dict[name] + param.weight_loader(param, loaded_weight, stacked_shard_id) + loaded_params.add(name) + break + else: + if "attn_sink" in name: + if name not in params_dict: + continue + narrow = loaded_weight[head_start:head_end] + params_dict[name][: narrow.shape[0]].copy_(narrow) + loaded_params.add(name) + continue + if ".shared_experts.w2" in name: + name = name.replace( + ".shared_experts.w2", ".shared_experts.down_proj" + ) + if name.endswith(".ffn.gate.bias"): + name = name.replace( + ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" + ) + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + self._finalize_moe() + logger.info_once("DSpark XPU draft model loaded: %d params", len(loaded_params)) + return loaded_params + + def _finalize_moe(self) -> None: + for layer in self.model.layers: + layer.ffn.finalize_mega_moe_weights() + + def _remap_dspark_name(self, name: str) -> str | None: + """Map checkpoint ``mtp.{i}.*`` name to this model's parameter path.""" + m = re.match(r"mtp\.(\d+)\.(.*)", name) + if m is None: + return None + stage = int(m.group(1)) + rest = m.group(2) + if rest.startswith("confidence_head."): + return None + head_prefixes = ( + "norm.", + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + "markov_head.", + ) + if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( + head_prefixes + ): + return f"model.{rest}" + return f"model.layers.{stage}.{rest}" diff --git a/vllm/models/deepseek_v4/xpu/model.py b/vllm/models/deepseek_v4/xpu/model.py index 1e5a574bed4a..e8449b9c058b 100644 --- a/vllm/models/deepseek_v4/xpu/model.py +++ b/vllm/models/deepseek_v4/xpu/model.py @@ -44,7 +44,11 @@ VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + SupportsEagle3, + SupportsPP, +) from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, @@ -975,7 +979,7 @@ def forward( @support_torch_compile -class DeepseekV4Model(nn.Module): +class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -1113,7 +1117,11 @@ def forward( input_ids = input_ids.to(torch.int64) residual, post_mix, res_mix = None, None, None - for layer in islice(self.layers, self.start_layer, self.end_layer): + aux_hidden_states: list[torch.Tensor] = [] + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): hidden_states, residual, post_mix, res_mix = layer( hidden_states, positions, @@ -1122,6 +1130,9 @@ def forward( res_mix, residual, ) + if idx + 1 in self.aux_hidden_state_layers: + aux_recon = layer.hc_post(hidden_states, residual, post_mix, res_mix) + aux_hidden_states.append(aux_recon.mean(dim=1)) # The fused path defers the final hc_post to the next layer's # fused_post_pre. After the last layer we must apply it explicitly. if layer is not None: @@ -1143,6 +1154,8 @@ def forward( self.hc_eps, ) hidden_states = self.norm(hidden_states) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1300,7 +1313,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: ) -class DeepseekV4ForCausalLM(nn.Module, SupportsPP): +class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3): model_cls = DeepseekV4Model # Default mapper assumes the original FP4-expert checkpoint layout. diff --git a/vllm/models/inkling/__init__.py b/vllm/models/inkling/__init__.py new file mode 100644 index 000000000000..32e58905e255 --- /dev/null +++ b/vllm/models/inkling/__init__.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .nvidia.model import ( + InklingForCausalLM, + InklingForConditionalGeneration, + ) + from .nvidia.mtp import InklingMTP + +__all__ = [ + "InklingForConditionalGeneration", + "InklingForCausalLM", + "InklingMTP", +] + + +def __getattr__(name: str): + if name == "InklingMTP": + from .nvidia import mtp + + return mtp.InklingMTP + if name in __all__: + from .nvidia import model + + return getattr(model, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/models/inkling/common/__init__.py b/vllm/models/inkling/common/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/inkling/common/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/inkling/common/mm_preprocess.py b/vllm/models/inkling/common/mm_preprocess.py new file mode 100644 index 000000000000..eb3c7b591ff5 --- /dev/null +++ b/vllm/models/inkling/common/mm_preprocess.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling multimodal preprocessing.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from typing import Any, cast + +import numpy as np +import regex as re +import torch +from transformers.feature_extraction_utils import BatchFeature + +from vllm.config.multimodal import ( + AudioDummyOptions, + BaseDummyOptions, + ImageDummyOptions, +) +from vllm.inputs import MultiModalDataDict +from vllm.multimodal.inputs import ( + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.transformers_utils.processors.inkling import ( + AUDIO_MARKER_ID, + AUDIO_TOKEN_ID, + IMAGE_MARKER_ID, + IMAGE_TOKEN_ID, + InklingAudioFeatureExtractor, + InklingImageProcessor, + InklingProcessor, +) + +from ..configs import InklingMMConfig + +# Maximum audio tokens accepted per clip. At the dMel rate of 20 tokens/s +# (50 ms hop) this is ~10 minutes of audio. It bounds the persistent per-request +# buffers and the encoder/memory budget; longer clips are rejected up front. +MAX_AUDIO_TOKENS = 12_000 + + +class InklingMultiModalDataParser(MultiModalDataParser): + def _parse_audio_data(self, data: Any) -> Any: + if isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 2: + raise ValueError( + "Inkling raw 2-D audio has an ambiguous channel layout. " + "Provide encoded audio or a list of mono waveforms." + ) + return super()._parse_audio_data(data) + + +def inkling_vision_enabled(config: InklingMMConfig) -> bool: + return getattr(config.vision_config, "decoder_dmodel", None) is not None + + +def inkling_audio_enabled(config: InklingMMConfig) -> bool: + return getattr(config.audio_config, "decoder_dmodel", None) is not None + + +class InklingProcessingInfo(BaseProcessingInfo): + def get_hf_config(self) -> InklingMMConfig: + return self.ctx.get_hf_config(InklingMMConfig) + + def get_hf_processor(self, **kwargs: object) -> InklingProcessor: + config = self.get_hf_config() + vision_config = config.vision_config + audio_config = config.audio_config + + image_processor = InklingImageProcessor( + patch_size=getattr(vision_config, "patch_size", None) or 40, + ) + + if inkling_audio_enabled(config): + audio_params = { + "n_mels": audio_config.n_mel_bins, + "num_dmel_bins": audio_config.mel_vocab_size, + "dmel_min_value": audio_config.dmel_min_value, + "dmel_max_value": audio_config.dmel_max_value, + } + else: + audio_params = {} + audio_extractor = InklingAudioFeatureExtractor(params=audio_params) + + return InklingProcessor( + image_processor=image_processor, + audio_feature_extractor=audio_extractor, + tokenizer=self.get_tokenizer(), + ) + + def get_data_parser(self) -> MultiModalDataParser: + # Audio inputs must be resampled to the dMel feature extractor's rate + # before `process_audios` (see InklingAudioFeatureExtractor._decode_one). + # Without a target_sr the default parser raises on any audio input. + extractor = self.get_hf_processor().audio_feature_extractor + return InklingMultiModalDataParser( + target_sr=extractor.params.sample_rate, + target_channels=1, + expected_hidden_size=self._get_expected_hidden_size(), + ) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + config = self.get_hf_config() + limits: dict[str, int | None] = {} + if inkling_vision_enabled(config): + limits["image"] = None + if inkling_audio_enabled(config): + limits["audio"] = None + return limits + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int] | None: + # Let vLLM profile dummy inputs to determine the max token counts; the + # image patch count is data-dependent, and the dummy audio is sized to + # MAX_AUDIO_TOKENS so audio is profiled/budgeted at its allowed maximum. + return None + + +class InklingDummyInputsBuilder(BaseDummyInputsBuilder[InklingProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + # One placeholder per media item; the processor expands each into N + # copies once feature row counts are known. + num_images = mm_counts.get("image", 0) + num_audios = mm_counts.get("audio", 0) + # Use spellings the renderer would emit; tokenization is bypassed in + # _call_hf_processor (we build input_ids directly), so the exact text + # only needs to be a stable per-item marker. + return ("<|content_image|>" * num_images) + ( + "<|content_audio_input|>" * num_audios + ) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + config = self.info.get_hf_config() + num_images = mm_counts.get("image", 0) + num_audios = mm_counts.get("audio", 0) + + mm_data: dict[str, Any] = {} + if num_images: + patch_size = getattr(config.vision_config, "patch_size", 40) + # A square image ~4 patches wide so the dummy emits several patches. + side = patch_size * 4 + image_overrides = mm_options.get("image") + mm_data["image"] = self._get_dummy_images( + width=side, + height=side, + num_images=num_images, + overrides=cast(ImageDummyOptions | None, image_overrides), + ) + if num_audios: + # Size the dummy at the maximum allowed audio so memory/encoder + # budgeting reflects the largest clip we accept (MAX_AUDIO_TOKENS). + params = self.info.get_hf_processor().audio_feature_extractor.params + hop = int(round(params.audio_token_duration_s * params.sample_rate)) + audio_len = MAX_AUDIO_TOKENS * hop + audio_overrides = mm_options.get("audio") + mm_data["audio"] = self._get_dummy_audios( + length=audio_len, + num_audios=num_audios, + overrides=cast(AudioDummyOptions | None, audio_overrides), + ) + return mm_data + + +class InklingMultiModalProcessor(BaseMultiModalProcessor[InklingProcessingInfo]): + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + # Inkling is not a standard HF processor (no fused text+mm call), so we run + # the vendored extractors ourselves and tokenize the text separately. + # The MM placeholders in `prompt` are expanded later by the prompt + # updates, so here we emit ONE placeholder id per media item. + processor = self.info.get_hf_processor(**mm_kwargs) + tokenizer = self.info.get_tokenizer() + + images = mm_data.get("images") or [] + audios = mm_data.get("audios") or [] + if not isinstance(images, list): + images = list(cast(Iterable[Any], images)) + if not isinstance(audios, list): + audios = list(cast(Iterable[Any], audios)) + + prompt_ids = self._tokenize_with_placeholders( + prompt, tokenizer, len(images), len(audios) + ) + + data: dict[str, Any] = {"input_ids": [prompt_ids]} + + if images: + img_feat = processor.process_images(images) + data["pixel_values"] = img_feat["vision_patches_bthwc"] + data["num_patches"] = torch.tensor( + img_feat["num_patches"], dtype=torch.int64 + ) + + if audios: + aud_feat = processor.process_audios(audios) + per_clip = aud_feat["dmel_bins"] + num_audio_tokens = aud_feat["num_audio_tokens"] + for i, n in enumerate(num_audio_tokens): + if int(n) > MAX_AUDIO_TOKENS: + raise ValueError( + f"Audio clip {i} produces {int(n)} tokens, exceeding the " + f"maximum of {MAX_AUDIO_TOKENS} (~10 min at 20 tokens/s). " + "Provide a shorter clip." + ) + if per_clip: + input_audio_features = torch.cat( + [torch.as_tensor(c) for c in per_clip], dim=0 + ) + else: + input_audio_features = torch.empty(0) + data["input_audio_features"] = input_audio_features + data["num_audio_tokens"] = torch.tensor(num_audio_tokens, dtype=torch.int64) + + return BatchFeature(data=data, tensor_type=None) + + def _tokenize_with_placeholders( + self, + prompt: str, + tokenizer: Any, + num_images: int, + num_audios: int, + ) -> list[int]: + """Tokenize `prompt`, emitting the block-start marker id per media item. + + Each marker (kept verbatim) is later expanded by ``_get_prompt_updates`` + into `` + * N``. + """ + image_marker = "<|content_image|>" + audio_marker = "<|content_audio_input|>" + + pattern = f"({re.escape(image_marker)}|{re.escape(audio_marker)})" + chunks = re.split(pattern, prompt) + + ids: list[int] = [] + seen_img = seen_aud = 0 + for chunk in chunks: + if chunk == image_marker: + ids.append(IMAGE_MARKER_ID) + seen_img += 1 + elif chunk == audio_marker: + ids.append(AUDIO_MARKER_ID) + seen_aud += 1 + elif chunk: + ids.extend(tokenizer.encode(chunk, add_special_tokens=False)) + + # Reconcile against the declared media counts only when media is + # present. With no media items -- e.g. the base text-only tokenization + # probe (``_apply_hf_processor_text_only``), which calls this via + # ``_call_hf_processor`` with empty ``mm_data`` -- emit the markers + # verbatim; the marker<->item correspondence is enforced later by + # ``_get_prompt_updates`` once the media features are available. + if num_images or num_audios: + # Fail clearly on a placeholder/media-count mismatch instead of + # crashing with an IndexError deep in the per-item replacement logic. + if num_images and seen_img != num_images: + raise ValueError( + f"Prompt contains {seen_img} image placeholder(s), but only " + f"{num_images} image(s) were provided." + ) + if num_audios and seen_aud != num_audios: + raise ValueError( + f"Prompt contains {seen_aud} audio placeholder(s), but only " + f"{num_audios} audio input(s) were provided." + ) + return ids + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + num_patches = hf_inputs.get("num_patches", torch.empty(0, dtype=torch.int64)) + num_audio_tokens = hf_inputs.get( + "num_audio_tokens", torch.empty(0, dtype=torch.int64) + ) + return dict( + # Ragged per-image patches, grouped by num_patches. + pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches), + num_patches=MultiModalFieldConfig.batched("image"), + # Ragged per-audio frames, grouped by num_audio_tokens. + input_audio_features=MultiModalFieldConfig.flat_from_sizes( + "audio", num_audio_tokens + ), + num_audio_tokens=MultiModalFieldConfig.batched("audio"), + ) + + def _get_prompt_updates( + self, + mm_items: Any, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + out_mm_data = out_mm_kwargs.get_data() + num_patches: Any = out_mm_data.get("num_patches") + num_audio_tokens: Any = out_mm_data.get("num_audio_tokens") + + # Keep the block-start marker and append N placeholder tokens after it; + # only the placeholder positions are flagged as embeddings (is_embed), so + # the marker stays a normal text token while the tower features scatter + # into the placeholders. + def image_replacement(item_idx: int) -> PromptUpdateDetails: + n = int(num_patches[item_idx]) + return PromptUpdateDetails.select_token_id( + [IMAGE_MARKER_ID] + [IMAGE_TOKEN_ID] * n, IMAGE_TOKEN_ID + ) + + def audio_replacement(item_idx: int) -> PromptUpdateDetails: + n = int(num_audio_tokens[item_idx]) + return PromptUpdateDetails.select_token_id( + [AUDIO_MARKER_ID] + [AUDIO_TOKEN_ID] * n, AUDIO_TOKEN_ID + ) + + updates: list[PromptUpdate] = [] + if num_patches is not None and len(num_patches) > 0: + updates.append( + PromptReplacement( + modality="image", + target=[IMAGE_MARKER_ID], + replacement=image_replacement, + ) + ) + if num_audio_tokens is not None and len(num_audio_tokens) > 0: + updates.append( + PromptReplacement( + modality="audio", + target=[AUDIO_MARKER_ID], + replacement=audio_replacement, + ) + ) + return updates diff --git a/vllm/models/inkling/common/towers.py b/vllm/models/inkling/common/towers.py new file mode 100644 index 000000000000..5ac3dbc20e6a --- /dev/null +++ b/vllm/models/inkling/common/towers.py @@ -0,0 +1,321 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling Titan vision + audio towers. + +The vision tower (``InklingVision`` / ``HMLPPatchEncoder``) emits one token per +image patch; the audio tower (``InklingAudio``) emits one token per audio frame. +Both use vLLM's standard ``RMSNorm`` (CPU-friendly, with a native fallback). +""" + +from __future__ import annotations + +from typing import cast + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.model_executor.layers.layernorm import RMSNorm + +from ..configs import InklingAudioConfig, InklingVisionConfig + +# =========================================================================== +# Vision tower (HMLPPatchEncoder / InklingVision) +# =========================================================================== + + +def _prime_factors(n: int) -> list[int]: + """Return the prime factors of *n* in ascending order.""" + if n < 1: + raise ValueError("n must be a positive integer") + + factors: list[int] = [] + while n % 2 == 0: + factors.append(2) + n //= 2 + p = 3 + while p * p <= n: + while n % p == 0: + factors.append(p) + n //= p + p += 2 + if n > 1: + factors.append(n) + return factors + + +def plan_out_scales( + temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int = 3 +) -> list[tuple[int, int, int, int]]: + """Plan the (t, h, w, c) dimensions for each HMLP layer. + + Spatial dims expand first, then temporal; channel counts round up to + multiples of 64. + """ + if patch_size <= 1: + raise ValueError( + "patch_size must be greater than 1, otherwise this doesn't make sense" + ) + + def _round_up(x: int) -> int: + return int(np.ceil(x / 64)) * 64 + + last_h_scale = 1 + scales: list[tuple[int, int, int, int]] = [(1, 1, 1, n_channels)] + for pscale in _prime_factors(patch_size)[::-1]: + last_h_scale *= pscale + scales.append( + ( + 1, + last_h_scale, + last_h_scale, + _round_up((last_h_scale**2) * n_channels), + ) + ) + last_t_scale = 1 + for tscale in _prime_factors(temporal_patch_size)[::-1]: + last_t_scale *= tscale + scales.append( + ( + last_t_scale, + last_h_scale, + last_h_scale, + _round_up((last_h_scale**2) * n_channels * last_t_scale), + ) + ) + + size_reduction = np.prod(np.array(scales)[:, :-1], 1) + + log_ideal_scales = np.linspace( + 0, + np.log(patch_size * patch_size * temporal_patch_size * n_channels), + n_layers + 1, + ) + cost_matrix = np.abs(log_ideal_scales[:, None] - np.log(size_reduction)[None]) + + if n_layers >= len(scales): + idxs = np.argmin(cost_matrix, axis=1) + else: + from scipy.optimize import linear_sum_assignment + + idxs = linear_sum_assignment(cost_matrix)[1] + + assert len(idxs) >= 2 + idxs[0] = 0 + idxs[-1] = len(scales) - 1 + + return [scales[i] for i in idxs] + + +def fold_timespace_to_depth( + vision_patches_bthwc: torch.Tensor, t_fold: int, hw_fold: int +) -> torch.Tensor: + """(B, T, H, W, C) -> (B, T//t, H//hw, W//hw, C*(t*hw**2)).""" + B, T, H, W, C = vision_patches_bthwc.shape + + assert T % t_fold == 0, f"Temporal dimension {T} must be divisible by {t_fold}" + assert H % hw_fold == 0, f"Height dimension {H} must be divisible by {hw_fold}" + assert W % hw_fold == 0, f"Width dimension {W} must be divisible by {hw_fold}" + + t_new = T // t_fold + h_new = H // hw_fold + w_new = W // hw_fold + + x = vision_patches_bthwc.reshape( + B, t_new, t_fold, h_new, hw_fold, w_new, hw_fold, C + ) + x = x.permute(0, 1, 3, 5, 2, 4, 6, 7) + x = x.reshape(B, t_new, h_new, w_new, t_fold * hw_fold * hw_fold * C) + return x + + +class HMLPPatchEncoder(nn.Module): + def __init__(self, config: InklingVisionConfig): + super().__init__() + self.decoder_dmodel = config.decoder_dmodel + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.n_channels = config.n_channels + self.n_layers = config.n_layers + self.use_vision_norm = config.use_vision_norm + + self.scales: list[tuple[int, int, int, int]] = plan_out_scales( + self.temporal_patch_size, self.patch_size, self.n_layers, self.n_channels + ) + self.layers: nn.ModuleDict = nn.ModuleDict() + for i, (start_scale, end_scale) in enumerate( + zip(self.scales[:-1], self.scales[1:]) + ): + shuffle_mult = ( + (end_scale[0] // start_scale[0]) + * (end_scale[1] // start_scale[1]) + * (end_scale[2] // start_scale[2]) + ) + if i == self.n_layers - 1: + self.layers[f"linear_{i}"] = nn.Linear( + start_scale[3] * shuffle_mult, self.decoder_dmodel, bias=False + ) + else: + self.layers[f"linear_{i}"] = nn.Linear( + start_scale[3] * shuffle_mult, end_scale[3], bias=False + ) + self.layers[f"norm_{i}"] = RMSNorm(end_scale[3]) + + self.final_norm: RMSNorm | None = None + if self.use_vision_norm: + assert self.decoder_dmodel is not None + self.final_norm = RMSNorm(self.decoder_dmodel) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Fused norm+gelu on CUDA (same fp32-accum/bf16-rounding structure as + # the generic path below; differs only by reduction order). + fused = None + if x.is_cuda and x.dtype == torch.bfloat16: + from vllm.models.inkling.nvidia.ops.mm_towers import rmsnorm_gelu + + fused = rmsnorm_gelu + + num_patches, T, H, W, C = x.shape + prefolded = False + for i, (start_scale, end_scale) in enumerate( + zip(self.scales[:-1], self.scales[1:]) + ): + t_fold = end_scale[0] // start_scale[0] + hw_fold = end_scale[1] // start_scale[1] + if (hw_fold > 1 or t_fold > 1) and not prefolded: + x = fold_timespace_to_depth(x, t_fold, hw_fold) + prefolded = False + assert x.shape[1:-1] == ( + T // end_scale[0], + H // end_scale[1], + W // end_scale[2], + ) + x = self.layers[f"linear_{i}"](x) + if i < self.n_layers - 1: + norm = cast(RMSNorm, self.layers[f"norm_{i}"]) + if fused is not None: + # If the NEXT layer starts with a copying fold (spatial + # dims still > 1 after folding), store this layer's + # output directly in the folded layout instead. + nxt = self.scales[i + 2] + ntf = nxt[0] // end_scale[0] + nhf = nxt[1] // end_scale[1] + copy_fold = (ntf > 1 or nhf > 1) and ( + x.shape[2] // nhf > 1 or x.shape[3] // nhf > 1 + ) + x = fused( + x, + norm.weight, + norm.variance_epsilon, + gelu=True, + fold=(ntf, nhf) if copy_fold else None, + ) + prefolded = copy_fold + else: + # rms_norm kernel only supports rank 2-4; x is 5-D here. + orig = x.shape + x = norm(x.reshape(-1, x.shape[-1])).reshape(orig) + x = F.gelu(x) + + if self.final_norm is not None: + if fused is not None: + x = fused( + x, + self.final_norm.weight, + self.final_norm.variance_epsilon, + gelu=False, + ) + else: + orig = x.shape + x = self.final_norm(x.reshape(-1, x.shape[-1])).reshape(orig) + + x = x.reshape(num_patches, -1) + return x + + +class InklingVision(nn.Module): + def __init__(self, config: InklingVisionConfig, prefix: str = ""): + del prefix + super().__init__() + assert config.vision_encoder_type == "hmlp" + self.vision_encoder = HMLPPatchEncoder(config) + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + def forward(self, vision_features: torch.Tensor) -> torch.Tensor: + return self.vision_encoder(vision_features) + + +# =========================================================================== +# Audio tower (InklingAudio) +# =========================================================================== + + +class InklingAudio(nn.Module): + def __init__(self, config: InklingAudioConfig, prefix: str = ""): + del prefix + super().__init__() + assert config.audio_mode == "dmel" + self.n_mel_bins = config.n_mel_bins + self.mel_vocab_size = config.mel_vocab_size + self.use_audio_norm = config.use_audio_norm + self.encoder = nn.Embedding( + config.n_mel_bins * config.mel_vocab_size, config.decoder_dmodel + ) + self.final_norm: RMSNorm | None = None + if self.use_audio_norm: + assert config.decoder_dmodel is not None + self.final_norm = RMSNorm(config.decoder_dmodel, eps=1e-6) + + @property + def dtype(self) -> torch.dtype: + return self.encoder.weight.dtype + + @property + def device(self) -> torch.device: + return self.encoder.weight.device + + def forward(self, audio_features: torch.Tensor) -> torch.Tensor: + assert audio_features.shape[1] == self.n_mel_bins + + # dMel bins are integer indices; cast once to int32 on the right device + # (no float round-trip). + audio_features = audio_features.to( + device=self.encoder.weight.device, dtype=torch.int32 + ) + + weight = self.encoder.weight + if audio_features.is_cuda and weight.dtype == torch.bfloat16: + # One kernel: per-bin offset + embedding gather + fp32 sum + norm. + # Skips the [T, n_mel_bins, D] intermediate entirely (bit-exact). + from vllm.models.inkling.nvidia.ops.mm_towers import dmel_embed_sum_norm + + return dmel_embed_sum_norm( + audio_features.contiguous(), + weight, + self.final_norm.weight if self.final_norm is not None else None, + self.final_norm.variance_epsilon if self.final_norm else 0.0, + ) + + embedding_indices = ( + torch.arange(self.n_mel_bins, device=audio_features.device) + * self.mel_vocab_size + ).unsqueeze(0) + audio_features + + hidden_states = ( + self.encoder(embedding_indices.reshape(-1)) + .reshape(audio_features.shape[0], audio_features.shape[1], -1) + .sum(axis=1) + ) + + if self.final_norm is not None: + hidden_states = self.final_norm(hidden_states) + + return hidden_states diff --git a/vllm/models/inkling/configs.py b/vllm/models/inkling/configs.py new file mode 100644 index 000000000000..aed70b143023 --- /dev/null +++ b/vllm/models/inkling/configs.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling model configs for the text backbone and audio/vision towers.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, ClassVar, Literal, cast + +import torch +from transformers.configuration_utils import PretrainedConfig + + +class InklingModelConfig(PretrainedConfig): + model_type = "inkling_model" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + *, + vocab_size: int = 201024, + hidden_size: int = 1536, + intermediate_size: int = 768, + dense_intermediate_size: int | None = None, + num_hidden_layers: int = 16, + num_attention_heads: int = 12, + num_key_value_heads: int = 4, + head_dim: int | None = None, + v_head_dim: int | None = None, + d_rel: int = 16, + rel_extent: int = 1024, + local_layer_ids: list[int] | None = None, + sliding_window_size: int = 512, + swa_num_attention_heads: int | None = None, + swa_num_key_value_heads: int | None = None, + swa_head_dim: int | None = None, + swa_v_head_dim: int | None = None, + rms_norm_eps: float = 1e-6, + hidden_act: str = "silu", + q_bias: bool = False, + o_bias: bool = False, + use_embed_norm: bool = False, + use_sconv: bool = False, + sconv_kernel_size: int = 4, + dense_mlp_idx: int = 0, + n_routed_experts: int = 0, + n_shared_experts: int = 0, + num_experts_per_tok: int = 1, + route_scale: float = 1.0, + use_gate_bias: bool = False, + use_global_scale: bool = False, + norm_after_topk: bool = True, + gate_activation: Literal["sigmoid", "softmax"] = "sigmoid", + shared_expert_sink: bool = False, + shared_experts_size: int = 1, + inference_moe_w13_interleaved: bool = True, + log_scaling_n_floor: int | None = None, + log_scaling_alpha: float = 0.1, + unpadded_vocab_size: int | None = None, + padded_vocab_size: int | None = None, + logits_mup_width_multiplier: float | None = None, + final_logit_softcapping: float | None = None, + tie_word_embeddings: bool = False, + num_nextn_predict_layers: int = 0, + chain_hidden_post_norm: bool = False, + **kwargs: Any, + ) -> None: + if head_dim is None: + head_dim = hidden_size // num_attention_heads + if v_head_dim is None: + v_head_dim = head_dim + if swa_num_attention_heads is None: + swa_num_attention_heads = num_attention_heads + if swa_num_key_value_heads is None: + swa_num_key_value_heads = num_key_value_heads + if swa_head_dim is None: + swa_head_dim = head_dim + if swa_v_head_dim is None: + swa_v_head_dim = swa_head_dim + if dense_intermediate_size is None: + dense_intermediate_size = intermediate_size + if local_layer_ids is None: + local_layer_ids = [] + + if padded_vocab_size is None: + padded_vocab_size = vocab_size + vocab_size = ( + unpadded_vocab_size + if ( + unpadded_vocab_size is not None + and unpadded_vocab_size < padded_vocab_size + ) + else vocab_size + ) + + self.vocab_size = vocab_size + self.padded_vocab_size = padded_vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.dense_intermediate_size = dense_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.v_head_dim = v_head_dim + self.d_rel = d_rel + self.rel_extent = rel_extent + self.local_layer_ids = local_layer_ids + self.sliding_window_size = sliding_window_size + self.swa_num_attention_heads = swa_num_attention_heads + self.swa_num_key_value_heads = swa_num_key_value_heads + self.swa_head_dim = swa_head_dim + self.swa_v_head_dim = swa_v_head_dim + self.rms_norm_eps = rms_norm_eps + self.hidden_act = hidden_act + self.q_bias = q_bias + self.o_bias = o_bias + self.use_embed_norm = use_embed_norm + self.use_sconv = use_sconv + self.sconv_kernel_size = sconv_kernel_size + self.dense_mlp_idx = dense_mlp_idx + self.n_routed_experts = n_routed_experts + self.num_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.num_shared_experts = n_shared_experts + self.num_experts_per_tok = num_experts_per_tok + self.route_scale = route_scale + self.use_gate_bias = use_gate_bias + self.use_global_scale = use_global_scale + self.norm_after_topk = norm_after_topk + self.gate_activation = gate_activation + self.shared_expert_sink = shared_expert_sink + self.shared_experts_size = shared_experts_size + self.inference_moe_w13_interleaved = inference_moe_w13_interleaved + self.log_scaling_n_floor = log_scaling_n_floor + self.log_scaling_alpha = log_scaling_alpha + self.unpadded_vocab_size = self.vocab_size + self.logits_mup_width_multiplier = logits_mup_width_multiplier + self.final_logit_softcapping = final_logit_softcapping + # MTP (multi-token prediction) draft head: number of depth layers in the + # checkpoint (0 if absent); chain_norm applied after each depth. + self.num_nextn_predict_layers = num_nextn_predict_layers + self.chain_hidden_post_norm = chain_hidden_post_norm + + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + @property + def conv_layer_ids(self) -> list[int]: + return list(range(self.num_hidden_layers)) + + @property + def linear_layer_ids(self) -> list[int]: + return self.conv_layer_ids + + @property + def full_attention_layer_ids(self) -> list[int]: + return list(range(self.num_hidden_layers)) + + @property + def mamba_chunk_size(self) -> int: + # Floor at 64: mamba_cache_chunk_size = max(mamba_chunk_size, page_size), + # and a floor of 1 lets the radix tree adopt another request's KV at + # tiny shared prefixes, whose different kernel-rounding perturbs decode + # logits. + return 64 + + @property + def mamba2_cache_params(self) -> TMLConvCacheParams | None: + try: + from vllm.distributed import get_tensor_model_parallel_world_size + + tp_size = get_tensor_model_parallel_world_size() + except (AssertionError, RuntimeError): + tp_size = 1 + + def tp_local_kv_conv_dim(num_kv_heads: int, head_dim: int) -> int: + return max(1, num_kv_heads // tp_size) * head_dim + + full_kv_conv_dim = tp_local_kv_conv_dim(self.num_key_value_heads, self.head_dim) + local_kv_conv_dim = tp_local_kv_conv_dim( + self.swa_num_key_value_heads, self.swa_head_dim + ) + stream_dim = self.hidden_size + conv_len = self.sconv_kernel_size - 1 + shape = TMLConvStateShape( + conv=[ + (conv_len, full_kv_conv_dim), + (conv_len, full_kv_conv_dim), + (conv_len, local_kv_conv_dim), + (conv_len, local_kv_conv_dim), + (conv_len, stream_dim), + (conv_len, stream_dim), + ], + temporal=(0, 0, 0), + ) + dtype = TMLStateDType(conv=torch.bfloat16, temporal=torch.bfloat16) + return TMLConvCacheParams(shape=shape, layers=self.conv_layer_ids, dtype=dtype) + + +class InklingAudioConfig(PretrainedConfig): + model_type = "inkling_audio_model" + + def __init__( + self, + *, + decoder_dmodel: int | None = None, + n_mel_bins: int | None = None, + mel_vocab_size: int | None = None, + dmel_min_value: float | None = None, + dmel_max_value: float | None = None, + use_audio_norm: bool | None = None, + audio_mode: Literal["dmel", "flow"] | None = None, + **kwargs: Any, + ) -> None: + values = { + "n_mel_bins": n_mel_bins, + "mel_vocab_size": mel_vocab_size, + "dmel_min_value": dmel_min_value, + "dmel_max_value": dmel_max_value, + "use_audio_norm": use_audio_norm, + "audio_mode": audio_mode, + } + if decoder_dmodel is not None and ( + missing := [name for name, value in values.items() if value is None] + ): + raise ValueError( + "Enabled Inkling audio tower is missing config fields: " + + ", ".join(missing) + ) + self.decoder_dmodel = decoder_dmodel + self.n_mel_bins = cast(int, n_mel_bins) + self.mel_vocab_size = cast(int, mel_vocab_size) + self.dmel_min_value = cast(float, dmel_min_value) + self.dmel_max_value = cast(float, dmel_max_value) + self.use_audio_norm = cast(bool, use_audio_norm) + self.audio_mode = cast(Literal["dmel", "flow"], audio_mode) + super().__init__(**kwargs) + + +class InklingVisionConfig(PretrainedConfig): + model_type = "inkling_vision_model" + + def __init__( + self, + *, + vision_encoder_type: Literal["linear", "hmlp"] | None = None, + decoder_dmodel: int | None = None, + patch_size: int | None = None, + temporal_patch_size: int | None = None, + n_channels: int | None = None, + n_layers: int | None = None, + use_vision_norm: bool | None = None, + **kwargs: Any, + ) -> None: + values = { + "vision_encoder_type": vision_encoder_type, + "patch_size": patch_size, + "temporal_patch_size": temporal_patch_size, + "n_channels": n_channels, + "n_layers": n_layers, + "use_vision_norm": use_vision_norm, + } + if decoder_dmodel is not None and ( + missing := [name for name, value in values.items() if value is None] + ): + raise ValueError( + "Enabled Inkling vision tower is missing config fields: " + + ", ".join(missing) + ) + self.vision_encoder_type = cast(Literal["linear", "hmlp"], vision_encoder_type) + self.decoder_dmodel = decoder_dmodel + self.patch_size = cast(int, patch_size) + self.temporal_patch_size = cast(int, temporal_patch_size) + self.n_channels = cast(int, n_channels) + self.n_layers = cast(int, n_layers) + self.use_vision_norm = cast(bool, use_vision_norm) + super().__init__(**kwargs) + + +class InklingMMConfig(PretrainedConfig): + model_type = "inkling_mm_model" + keys_to_ignore_at_inference = ["past_key_values"] + sub_configs: ClassVar[dict[str, type[PretrainedConfig]]] = { + "text_config": InklingModelConfig, + "audio_config": InklingAudioConfig, + "vision_config": InklingVisionConfig, + } + + def __init__( + self, + *, + text_config: dict[str, Any] | InklingModelConfig | None = None, + audio_config: dict[str, Any] | InklingAudioConfig | None = None, + vision_config: dict[str, Any] | InklingVisionConfig | None = None, + tie_word_embeddings: bool = False, + **kwargs: Any, + ) -> None: + self.text_config = ( + text_config + if isinstance(text_config, InklingModelConfig) + else InklingModelConfig(**(text_config or {})) + ) + self.audio_config = ( + audio_config + if isinstance(audio_config, InklingAudioConfig) + else InklingAudioConfig(**(audio_config or {})) + ) + self.vision_config = ( + vision_config + if isinstance(vision_config, InklingVisionConfig) + else InklingVisionConfig(**(vision_config or {})) + ) + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + def get_text_config(self, *args: Any, **kwargs: Any) -> InklingModelConfig: + return self.text_config + + @property + def vocab_size(self) -> int: + return self.text_config.vocab_size + + @property + def hidden_size(self) -> int: + return self.text_config.hidden_size + + @property + def num_hidden_layers(self) -> int: + return self.text_config.num_hidden_layers + + @property + def num_attention_heads(self) -> int: + return self.text_config.num_attention_heads + + @property + def num_key_value_heads(self) -> int: + return self.text_config.num_key_value_heads + + @property + def head_dim(self) -> int: + return self.text_config.head_dim + + @property + def full_attention_layer_ids(self) -> list[int]: + return self.text_config.full_attention_layer_ids + + @property + def linear_layer_ids(self) -> list[int]: + return self.text_config.linear_layer_ids + + @property + def conv_layer_ids(self) -> list[int]: + return self.text_config.conv_layer_ids + + @property + def mamba_chunk_size(self) -> int: + return self.text_config.mamba_chunk_size + + @property + def mamba2_cache_params(self) -> TMLConvCacheParams | None: + return self.text_config.mamba2_cache_params + + +@dataclass(kw_only=True, frozen=True) +class TMLConvStateShape: + conv: list[tuple[int, int]] + temporal: tuple[int, int, int] + + +@dataclass(kw_only=True, frozen=True) +class TMLStateDType: + conv: torch.dtype = torch.bfloat16 + temporal: torch.dtype = torch.bfloat16 + + +@dataclass(kw_only=True, frozen=True) +class TMLConvCacheParams: + shape: TMLConvStateShape + layers: list[int] + dtype: TMLStateDType = field(default_factory=TMLStateDType) diff --git a/vllm/models/inkling/nvfp4.py b/vllm/models/inkling/nvfp4.py new file mode 100644 index 000000000000..129ab82aa14e --- /dev/null +++ b/vllm/models/inkling/nvfp4.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVFP4 (ModelOpt) support for the Inkling mixture-of-experts. + +Only the routed MoE experts are quantized in the Inkling checkpoint; +attention, the dense MLP, and the shared "sink" experts stay bf16 (they are +in the checkpoint ``exclude_modules``). The routed experts are served by +vLLM's standard ModelOpt NVFP4 fused-MoE stack (see ``moe.py``); this module +keeps the checkpoint detection. +""" + +from __future__ import annotations + +FLOAT8_E4M3_MAX = 448.0 +FLOAT4_E2M1_MAX = 6.0 + + +class InklingNvfp4Config: + """Lightweight NVFP4 descriptor parsed from the checkpoint quant config. + + Holds the (mapped) ``exclude_modules`` so the model can decide, per MoE + layer and per expert group, whether the weights are NVFP4 or plain bf16. + """ + + def __init__(self, group_size: int, exclude_modules: list[str]) -> None: + self.group_size = group_size + self.exclude_modules = set(exclude_modules) + + @staticmethod + def _is_nvfp4(quant_cfg: dict) -> bool: + wq = quant_cfg["modelopt_quant_config"]["quant_cfg"]["*weight_quantizer"] + return tuple(wq["num_bits"]) == (2, 1) and tuple( + wq["block_sizes"].get("scale_bits", []) + ) == (4, 3) + + @classmethod + def from_hf_config(cls, hf_config) -> InklingNvfp4Config | None: + quant_cfg = getattr(hf_config, "quantization_config", None) + text_config = getattr(hf_config, "text_config", None) + if quant_cfg is None and text_config is not None: + quant_cfg = getattr(text_config, "quantization_config", None) + if quant_cfg is None: + return None + # ModelOpt <=0.29 nests everything under "quantization". + if "quantization" in quant_cfg: + quant_cfg = quant_cfg["quantization"] + if not cls._is_nvfp4(quant_cfg): + return None + group_size = quant_cfg.get("group_size", 16) + if group_size != 16: + raise ValueError("Inkling NVFP4 only supports group size 16") + exclude = list(quant_cfg.get("exclude_modules", []) or []) + return cls(group_size=group_size, exclude_modules=exclude) + + def experts_quantized(self, layer_id: int) -> bool: + """Whether the routed experts of ``layer_id`` are NVFP4 (vs excluded).""" + return f"model.llm.layers.{layer_id}.mlp.experts" not in self.exclude_modules + + def shared_experts_quantized(self, layer_id: int) -> bool: + """Whether the shared sink experts of ``layer_id`` are NVFP4.""" + return ( + f"model.llm.layers.{layer_id}.mlp.shared_experts" + not in self.exclude_modules + ) diff --git a/vllm/models/inkling/nvidia/__init__.py b/vllm/models/inkling/nvidia/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/inkling/nvidia/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/inkling/nvidia/attention.py b/vllm/models/inkling/nvidia/attention.py new file mode 100644 index 000000000000..fc617b89ad6a --- /dev/null +++ b/vllm/models/inkling/nvidia/attention.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from typing import cast + +import torch +from torch import nn + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.utils.torch_utils import ( + canonicalize_singleton_dim_strides, + kv_cache_dtype_str_to_dtype, +) +from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.attention.backends.flash_attn import ( + FlashAttentionBackend, + FlashAttentionMetadata, +) +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheSpec, + SlidingWindowSpec, +) + +from ..configs import InklingModelConfig +from .layernorm import InklingRMSNorm +from .ops.fa4_rel_attention import ( + bucket_max_seqlen_q, + inkling_fa4_num_splits, + inkling_fa4_rel_attention, +) +from .ops.fa4_warmup import InklingFA4WarmupConfig, register_fa4_warmup +from .ops.qkvr_prep import fused_qkvr_prep +from .sconv_swa_attn import _K, _V, InklingConvState, InklingSconvMetadata +from .short_conv import InklingShortConv + + +def compute_log_scaling_tau( + positions: torch.Tensor, n_floor: int, alpha: float +) -> torch.Tensor: + effective_n = (positions + 1).to(torch.float32) + return 1.0 + alpha * torch.log(torch.clamp(effective_n / float(n_floor), min=1.0)) + + +class RelLogitsProj(nn.Module): + """Project the per-head relative branch ``r`` to per-distance logits.""" + + def __init__(self, d_rel: int, rel_extent: int) -> None: + super().__init__() + self.d_rel = d_rel + self.rel_extent = rel_extent + self.proj = nn.Parameter(torch.empty(d_rel, rel_extent), requires_grad=False) + + def forward(self, r_out: torch.Tensor) -> torch.Tensor: + # r_out: (T, num_heads, d_rel) -> (T, num_heads, rel_extent) + return torch.einsum("thd,de->the", r_out, self.proj) + + +class InklingAttention(nn.Module, AttentionLayerBase): + def __init__( + self, + config: InklingModelConfig, + *, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rel_extent: int, + local_extent: int, + is_local: bool, + prefix: str, + quant_config: QuantizationConfig | None = None, + conv_owner: InklingConvState, + ) -> None: + super().__init__() + self.prefix = prefix + self.is_local = is_local + self.hidden_size = config.hidden_size + self.head_dim = head_dim + self.d_rel = config.d_rel + self.log_scaling_n_floor = config.log_scaling_n_floor + self.log_scaling_alpha = config.log_scaling_alpha + # q/k are per-head RMS-normed (unit norm), so Inkling scales by 1/head_dim. + self.scaling = 1.0 / head_dim + + tp_size = get_tensor_model_parallel_world_size() + self.num_total_heads = num_heads + self.num_total_kv_heads = num_kv_heads + assert self.num_total_heads % tp_size == 0 + self.num_heads = self.num_total_heads // tp_size + if self.num_total_kv_heads >= tp_size: + assert self.num_total_kv_heads % tp_size == 0 + else: + assert tp_size % self.num_total_kv_heads == 0 + self.num_kv_heads = max(1, self.num_total_kv_heads // tp_size) + # When tp_size > num_kv_heads the K/V projections are padded up to + # tp_size heads so each rank gets at least one (GQA replication). + kv_total_for_sizing = max(self.num_total_kv_heads, tp_size) + + self.qkvr = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[ + head_dim * self.num_total_heads, + head_dim * kv_total_for_sizing, + head_dim * kv_total_for_sizing, + self.d_rel * self.num_total_heads, + ], + bias=config.q_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkvr", + ) + self.wo_ud = RowParallelLinear( + input_size=head_dim * self.num_total_heads, + output_size=config.hidden_size, + bias=config.o_bias, + quant_config=quant_config, + # reduce_results=False: the partial output is all-reduced below + # (one-shot custom AR) so the attention-output sconv can run on the + # full hidden width fused with the residual add + rmsnorm. + reduce_results=False, + prefix=f"{prefix}.wo_ud", + ) + self.rel_extent = local_extent if is_local else rel_extent + self.local_extent = local_extent if is_local else None + self.rel_logits_proj = RelLogitsProj(self.d_rel, self.rel_extent) + self.q_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps) + self.k_norm = InklingRMSNorm(head_dim, eps=config.rms_norm_eps) + + # Short convolution on the K/V streams (per-head-width, TP sharded), + # applied after the qkvr projection and before q/k norm. + kv_conv_dim = self.num_kv_heads * head_dim + self.conv_owner = conv_owner + self.k_sconv = InklingShortConv( + kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_K + ) + self.v_sconv = InklingShortConv( + kv_conv_dim, config.sconv_kernel_size, owner=conv_owner, stream_idx=_V + ) + + # FA4 left/right window; right=0 keeps it causal. local_extent-1 mirrors + # the source (sliding_window_size - 1). + self.window_size: tuple[int, int] = ( + (local_extent - 1, 0) if is_local else (-1, -1) + ) + # Static per-layer-type KV length bound for the split heuristic: local + # layers never see more than the sliding window. + vllm_config = get_current_vllm_config() + self._max_kv_len = ( + local_extent if is_local else vllm_config.model_config.max_model_len + ) + + # ---- KV-cache wiring (reuse FlashAttentionBackend for metadata) ---- + cache_config = vllm_config.cache_config + self.kv_cache_dtype = ( + cache_config.cache_dtype if cache_config is not None else "auto" + ) + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype( + self.kv_cache_dtype, vllm_config.model_config + ) + self.register_buffer("k_scale", torch.ones((), dtype=torch.float32)) + self.register_buffer("v_scale", torch.ones((), dtype=torch.float32)) + + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) # replaced by bind_kv_cache + + register_fa4_warmup( + InklingFA4WarmupConfig( + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + rel_extent=self.rel_extent, + window_size=self.window_size, + is_local=self.is_local, + max_kv_len=self._max_kv_len, + dtype=vllm_config.model_config.dtype, + kv_dtype=self.kv_cache_torch_dtype, + block_size=vllm_config.cache_config.block_size, + max_num_reqs=vllm_config.scheduler_config.max_num_seqs, + max_num_batched_tokens=( + vllm_config.scheduler_config.max_num_batched_tokens + ), + ) + ) + + def get_attn_backend(self) -> type[AttentionBackend]: + return FlashAttentionBackend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + block_size = vllm_config.cache_config.block_size + if self.is_local: + assert self.local_extent is not None + return SlidingWindowSpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + dtype=self.kv_cache_torch_dtype, + sliding_window=self.local_extent, + ) + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + dtype=self.kv_cache_torch_dtype, + ) + + def _split_kv_cache(self) -> tuple[torch.Tensor, torch.Tensor]: + key_cache, value_cache = self.kv_cache.transpose(1, 2).split( + self.head_dim, dim=-1 + ) + return ( + canonicalize_singleton_dim_strides(key_cache), + canonicalize_singleton_dim_strides(value_cache), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + log_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + num_tokens = hidden_states.shape[0] + qkvr, _ = self.qkvr(hidden_states) + + attn_metadata = get_forward_context().attn_metadata + attn_output = torch.empty( + (num_tokens, self.num_heads, self.head_dim), + dtype=qkvr.dtype, + device=qkvr.device, + ) + if not isinstance(attn_metadata, dict): + attn_output.zero_() + else: + conv_meta = attn_metadata[self.conv_owner.prefix] + md = attn_metadata[self.prefix] + assert isinstance(conv_meta, InklingSconvMetadata) + fa_md = cast(FlashAttentionMetadata, md) + assert self.kv_cache.numel() > 0 + assert self.conv_owner.kv_cache.numel() > 0 + # One launch: K/V sconv (conv-cache insert + conv + residual), + # Q/K per-head rmsnorm, and the attention KV-cache write. K/V are + # consumed via the KV cache; only normed q is materialized. + key_cache, value_cache = self._split_kv_cache() + off_k, _ = self.conv_owner.stream_ranges[_K] + off_v, _ = self.conv_owner.stream_ranges[_V] + q, rel_logits = fused_qkvr_prep( + qkvr, + self.k_sconv.weight.squeeze(1), + self.v_sconv.weight.squeeze(1), + self.q_norm.weight, + self.k_norm.weight, + self.rel_logits_proj.proj, + self.q_norm.variance_epsilon, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.d_rel, + self.conv_owner.kv_cache, + key_cache, + value_cache, + positions, + conv_meta.block_table, + conv_meta.seq_idx, + conv_meta.slot_mapping, + conv_meta.query_start, + fa_md.slot_mapping, + off_k, + off_v, + self.conv_owner.block_size, + log_scaling if not self.is_local else None, + ) + q = q.view(num_tokens, self.num_heads, self.head_dim) + self._attention(q, rel_logits, attn_output) + + flat = attn_output.view(num_tokens, -1) + output, _ = self.wo_ud(flat) + return output + + @eager_break_during_capture + def _attention( + self, + q: torch.Tensor, + rel_logits: torch.Tensor, + output: torch.Tensor, + ) -> None: + attn_metadata = get_forward_context().attn_metadata + assert isinstance(attn_metadata, dict) + md = cast(FlashAttentionMetadata, attn_metadata[self.prefix]) + + nt = md.num_actual_tokens + key_cache, value_cache = self._split_kv_cache() + max_seqlen_q = bucket_max_seqlen_q(md.max_query_len) + num_splits = inkling_fa4_num_splits( + is_local=self.is_local, + batch_size=md.seq_lens.shape[0], + max_query_len=max_seqlen_q, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + max_kv_len=self._max_kv_len, + ) + inkling_fa4_rel_attention( + q[:nt], + key_cache, + value_cache, + block_table=md.block_table, + cache_seqlens=md.seq_lens, + cu_seqlens_q=md.query_start_loc, + max_seqlen_q=max_seqlen_q, + softmax_scale=self.scaling, + causal=True, + window_size=self.window_size, + rel_extent=self.rel_extent, + rel_logits=rel_logits[:nt], + num_splits=num_splits, + out=output[:nt], + ) diff --git a/vllm/models/inkling/nvidia/layernorm.py b/vllm/models/inkling/nvidia/layernorm.py new file mode 100644 index 000000000000..18f9951f65d4 --- /dev/null +++ b/vllm/models/inkling/nvidia/layernorm.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling RMSNorm (no bias, weight-scaled), backed by the vendored Triton kernel.""" + +from __future__ import annotations + +import torch +from torch import nn + +from .ops import rmsnorm + + +class InklingRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.numel() == 0: + return x + original_shape = x.shape + x_2d = x.contiguous().view(-1, self.hidden_size) + y = rmsnorm(x_2d, self.weight, self.variance_epsilon) + return y.view(original_shape) diff --git a/vllm/models/inkling/nvidia/logits_processor.py b/vllm/models/inkling/nvidia/logits_processor.py new file mode 100644 index 000000000000..2e142c5961ef --- /dev/null +++ b/vllm/models/inkling/nvidia/logits_processor.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling logits processor (muP + LoRA aware). + +Inkling divides the final logits by a muP width multiplier +(``logits_mup_width_multiplier``). This applies it two ways, depending on +whether an lm_head LoRA is attached: + +* No LoRA: fold ``1/mup`` into the lm_head GEMM alpha (fp32 epilogue) -- no + separate elementwise kernel, no extra rounding, no weight mutation. +* LoRA attached: the LoRA manager wraps this layer in + ``LogitsProcessorWithLoRA``, whose ``forward`` calls + ``type(base_layer).forward(self=wrapper)`` -- so this ``forward`` runs with + ``self`` bound to the wrapper. We detect that via ``base_layer`` and take the + LoRA path: run the wrapper's ``_get_logits`` (base logits + the lm_head LoRA + delta), then divide the full logits by the multiplier so the delta is scaled + too. muP thus composes with the LoRA delta, with the dispatch as the only + model-side branching. +""" + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding + + +class InklingLogitsProcessor(LogitsProcessor): + """``LogitsProcessor`` that applies Inkling's muP logits width multiplier. + + Args: + vocab_size: Padded vocabulary size. + org_vocab_size: Unpadded vocabulary size (defaults to ``vocab_size``). + scale: Base logits scale (kept ``1.0`` for the served checkpoint). + logits_as_input: Whether the input is already logits. + soft_cap: Optional logit soft cap (``None`` for the served checkpoint). + logits_mup_width_multiplier: muP width divisor for the final logits; + ``None`` or ``0`` disables it. + """ + + def __init__( + self, + vocab_size: int, + org_vocab_size: int | None = None, + scale: float = 1.0, + logits_as_input: bool = False, + soft_cap: float | None = None, + logits_mup_width_multiplier: float | None = None, + ) -> None: + super().__init__( + vocab_size=vocab_size, + org_vocab_size=org_vocab_size, + scale=scale, + logits_as_input=logits_as_input, + soft_cap=soft_cap, + ) + self.logits_mup_width_multiplier = logits_mup_width_multiplier + self._logits_zero: torch.Tensor | None = None + + def forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + # ``base_layer`` exists only on the LogitsProcessorWithLoRA wrapper, + # which calls this forward with ``self`` bound to the wrapper. The + # wrapper is not an ``InklingLogitsProcessor`` instance, so dispatch + # ``_lora_forward`` explicitly through the base_layer's class (it + # provides ``_get_logits``/``logits_as_input``; only ``_lora_forward`` + # lives on this class). + if hasattr(self, "base_layer"): + return type(self.base_layer)._lora_forward( + self, lm_head, hidden_states, embedding_bias + ) + return self._base_forward(lm_head, hidden_states, embedding_bias) + + def _lora_forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + # ``self`` is the LogitsProcessorWithLoRA wrapper here: ``_get_logits`` + # returns the base logits plus the lm_head LoRA delta. Apply the muP + # divisor on the full logits so the LoRA delta is scaled too. + mup_multiplier = self.base_layer.logits_mup_width_multiplier + mup = 1.0 / mup_multiplier if mup_multiplier else None + if self.logits_as_input: + logits = hidden_states + else: + logits = self._get_logits(hidden_states, lm_head, embedding_bias) + # TODO: fuse this multiplication + if logits is not None and mup: + assert self.base_layer.soft_cap is None + assert self.base_layer.scale == 1.0 + logits = logits * mup + return logits + + def _base_forward( + self, + lm_head: VocabParallelEmbedding, + hidden_states: torch.Tensor, + embedding_bias: torch.Tensor | None = None, + ) -> torch.Tensor | None: + mup = self.logits_mup_width_multiplier + if not mup: + return super().forward(lm_head, hidden_states, embedding_bias) + # Fold the muP width divisor into the lm_head GEMM alpha (fp32 epilogue): + # no separate elementwise kernel, no bf16 rounding of scaled logits, and + # no weight mutation. Overfit to the served checkpoint: bf16 lm_head, no + # soft cap, unit logits scale. + assert self.soft_cap is None + assert self.scale == 1.0 + w = lm_head.weight + if self._logits_zero is None: + self._logits_zero = w.new_zeros(1) + logits = torch.addmm( + self._logits_zero, + hidden_states, + w.t(), + beta=0.0, + alpha=1.0 / mup, + ) + logits = self._gather_logits(logits) + if logits is not None: + logits = logits[..., : self.org_vocab_size] + return logits diff --git a/vllm/models/inkling/nvidia/mlp.py b/vllm/models/inkling/nvidia/mlp.py new file mode 100644 index 000000000000..ab7dec0b709c --- /dev/null +++ b/vllm/models/inkling/nvidia/mlp.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling dense SwiGLU MLP (also used as the MoE shared expert). + +The checkpoint stores the gate/up projection as a single fused, *interleaved* +weight (``[gate0, up0, gate1, up1, ...]``), so we use a plain +``ColumnParallelLinear`` whose contiguous row-sharding keeps each gate/up pair +together, and an interleaved SwiGLU activation. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig + + +class InklingDenseMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + *, + use_global_scale: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = ColumnParallelLinear( + hidden_size, + 2 * intermediate_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.down_proj", + ) + if use_global_scale: + self.global_scale = nn.Parameter(torch.empty(1), requires_grad=False) + else: + self.global_scale = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from .ops import silu_and_mul_triton + + gate_up, _ = self.gate_up_proj(x) + x = silu_and_mul_triton(gate_up) + x, _ = self.down_proj(x) + if self.global_scale is not None: + x = x * self.global_scale + # TP-partial output: the layer's reduce-scatter fallback consumes it. + return x diff --git a/vllm/models/inkling/nvidia/model.py b/vllm/models/inkling/nvidia/model.py new file mode 100644 index 000000000000..bf7f700dee66 --- /dev/null +++ b/vllm/models/inkling/nvidia/model.py @@ -0,0 +1,697 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling model implementation for NVIDIA GPUs.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import regex as re +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, + tensor_model_parallel_reduce_scatter, +) +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.models.inkling.common.mm_preprocess import ( + InklingDummyInputsBuilder, + InklingMultiModalProcessor, + InklingProcessingInfo, + inkling_audio_enabled, + inkling_vision_enabled, +) +from vllm.models.inkling.common.towers import InklingAudio, InklingVision +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors + +from ..configs import InklingMMConfig, InklingModelConfig +from ..nvfp4 import InklingNvfp4Config +from .attention import InklingAttention, compute_log_scaling_tau +from .layernorm import InklingRMSNorm +from .logits_processor import InklingLogitsProcessor +from .mlp import InklingDenseMLP +from .moe import InklingMoE +from .ops.lamport import get_lamport_rs_conv, initialize_lamport_rs_conv +from .ops.norm import add_rmsnorm, embed_rmsnorm +from .sconv_swa_attn import _ATTN, _MLP, InklingConvState, InklingSconvMetadata +from .short_conv import InklingShortConv + + +def _layer_id(name: str) -> int | None: + m = re.search(r"\.layers\.(\d+)\.", name) + return int(m.group(1)) if m else None + + +def _sconv_add_norm( + delta: torch.Tensor, + hidden: torch.Tensor, + sconv: InklingShortConv, + norm: InklingRMSNorm | None, + positions: torch.Tensor, +) -> tuple[torch.Tensor | None, torch.Tensor]: + """``h = hidden + sconv(TP-sum(delta)); y = rmsnorm(h)``. + + The Lamport path performs reduce-scatter + shard sconv + all-gather + + residual add + norm. The NCCL path handles unsupported configurations.""" + attn_metadata = get_forward_context().attn_metadata + m = ( + attn_metadata.get(sconv.owner.prefix) + if isinstance(attn_metadata, dict) + else None + ) + cache = sconv.owner.kv_cache + off_s, ws = sconv.owner.stream_ranges[sconv.stream_idx] + norm_w = norm.weight if norm is not None else None + eps = norm.variance_epsilon if norm is not None else 0.0 + + mm = get_lamport_rs_conv(hidden.shape[-1], sconv.kernel_size) + if mm is not None and mm.usable(delta.shape[0]) and m is not None: + assert cache.numel() > 0 + assert isinstance(m, InklingSconvMetadata) + return mm.rs_sconv_ag_add_norm( + delta, + hidden, + sconv.weight.squeeze(1), + norm_w, + eps, + cache, + positions, + m.block_table, + m.seq_idx, + m.slot_mapping, + off_s, + ws, + sconv.owner.block_size, + ) + + # Fallback: NCCL RS -> shard sconv -> AG -> fused add(+rmsnorm). + shard = tensor_model_parallel_reduce_scatter(delta, dim=-1) + shard = sconv(shard.contiguous(), positions) + full = tensor_model_parallel_all_gather(shard, dim=-1) + if norm is None: + return None, hidden + full + return add_rmsnorm(hidden, full, norm_w, eps) + + +class InklingDecoderLayer(nn.Module): + def __init__( + self, + config: InklingModelConfig, + layer_id: int, + is_local: bool, + quant_config: QuantizationConfig | None, + prefix: str, + nvfp4_config: InklingNvfp4Config | None = None, + force_dense_mlp: bool = False, + ) -> None: + super().__init__() + # Per-layer owner of the conv state as a paged SWA cache. The 4 sconv + # streams (K/V/attn/mlp) are packed head-major into one block and share + # it. Built first so the attention layer can wire its K/V sconv to it. + self.conv_state = InklingConvState( + num_kv_heads=( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ), + head_dim=config.swa_head_dim if is_local else config.head_dim, + hidden_size=config.hidden_size, + kernel_size=config.sconv_kernel_size, + prefix=f"{prefix}.conv_state", + ) + self.attn_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.attn = InklingAttention( + config, + num_heads=( + config.swa_num_attention_heads + if is_local + else config.num_attention_heads + ), + num_kv_heads=( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ), + head_dim=config.swa_head_dim if is_local else config.head_dim, + rel_extent=config.rel_extent, + local_extent=config.sliding_window_size, + is_local=is_local, + prefix=f"{prefix}.attn", + quant_config=quant_config, + conv_owner=self.conv_state, + ) + self.mlp_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if force_dense_mlp or layer_id < config.dense_mlp_idx: + self.mlp: nn.Module = InklingDenseMLP( + hidden_size=config.hidden_size, + intermediate_size=config.dense_intermediate_size, + use_global_scale=config.use_global_scale, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + else: + # InklingMoE decides per layer (from the checkpoint exclude list) + # whether the routed experts are NVFP4 or bf16; the shared sink + # experts are always bf16. + self.mlp = InklingMoE( + config, + layer_id, + prefix=f"{prefix}.mlp", + nvfp4_config=nvfp4_config, + ) + + # Short convolution on the attention-output and MLP-output residual + # streams, hidden-sharded: the sublayer outputs are reduce-scattered + # to [T, H/tp], the sconv runs on the shard, and an all-gather + # restores the full residual — all fused with the residual add + next + # rmsnorm via the Lamport P2P kernels for decode-sized batches. + tp_size = get_tensor_model_parallel_world_size() + sconv_dim = config.hidden_size // tp_size + self.attn_sconv = InklingShortConv( + sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_ATTN + ) + self.mlp_sconv = InklingShortConv( + sconv_dim, config.sconv_kernel_size, owner=self.conv_state, stream_idx=_MLP + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + pending: tuple[torch.Tensor | None, InklingShortConv] | None = None, + defer_mlp_add: bool = False, + attn_in: torch.Tensor | None = None, + log_scaling: torch.Tensor | None = None, + ) -> ( + torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor | None, InklingShortConv]] + ): + # The previous sublayer's (pre-reduce, pre-sconv) delta is folded in + # fused with its RS/sconv/AG and this layer's pre-attention rmsnorm. + # A None delta means the partials sit in the NVLS symm buffer. + if pending is None: + if attn_in is None: + # First layer; on the text path attn_norm comes fused with + # the embedding gather (chain_weight in embed_rmsnorm). + attn_in = self.attn_norm(hidden_states) + else: + attn_in, hidden_states = _sconv_add_norm( + pending[0], hidden_states, pending[1], self.attn_norm, positions + ) + attn_output = self.attn(positions, attn_in, log_scaling) + mlp_in, hidden_states = _sconv_add_norm( + attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions + ) + mlp_output = self.mlp(mlp_in) + if defer_mlp_add: + # Caller folds mlp_output (pre-reduce, pre-sconv) into the next + # fused sconv+add+rmsnorm. + return hidden_states, (mlp_output, self.mlp_sconv) + return _sconv_add_norm( + mlp_output, hidden_states, self.mlp_sconv, None, positions + )[1] + + +class InklingReplicatedEmbedding(nn.Module): + """Full-vocab embedding table replicated on every TP rank. + + Trades the full table per rank (~2.3 GiB at V=201k / H=6144 bf16, vs a + 1/tp shard) for no masked lookup or per-lookup TP all-reduce, and keeps the + full table on-rank for the fused gather+norm kernel. Bit-exact vs + vocab-parallel: the all-reduce there only ever summed one real row against + exact zeros. The LM head stays vocab-sharded. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int) -> None: + super().__init__() + self.weight = nn.Parameter( + torch.empty(num_embeddings, embedding_dim, dtype=torch.get_default_dtype()), + requires_grad=False, + ) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return embed_rmsnorm(input_ids, self.weight, None, 0.0) + + +class InklingModel(nn.Module): + def __init__( + self, + *, + config: InklingModelConfig, + quant_config: QuantizationConfig | None, + prefix: str, + nvfp4_config: InklingNvfp4Config | None = None, + ) -> None: + super().__init__() + self.config = config + self.embed_tokens = InklingReplicatedEmbedding( + config.padded_vocab_size, config.hidden_size + ) + self.embed_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.use_embed_norm + else None + ) + local_ids = set(config.local_layer_ids) + + def get_layer(prefix: str) -> InklingDecoderLayer: + idx = _layer_id(prefix + ".") or int(prefix.split(".")[-1]) + return InklingDecoderLayer( + config, idx, idx in local_ids, quant_config, prefix, nvfp4_config + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers" + ) + self.norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + # Row gather + embed_norm in one launch. + norm = self.embed_norm + return embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + norm.variance_epsilon if norm is not None else 0.0, + ) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + attn_in0: torch.Tensor | None = None + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + # embed_norm was already applied when producing inputs_embeds. + hidden_states = inputs_embeds + else: + # Gather + embed_norm + the first layer's attn_norm, one launch. + norm = self.embed_norm + hidden_states, attn_in0 = embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + self.config.rms_norm_eps, + chain_weight=self.layers[self.start_layer].attn_norm.weight, + ) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + log_scaling = None + if self.config.log_scaling_n_floor is not None: + log_scaling = compute_log_scaling_tau( + positions, + self.config.log_scaling_n_floor, + self.config.log_scaling_alpha, + ) + + pending: tuple[torch.Tensor | None, InklingShortConv] | None = None + for layer in self.layers[self.start_layer : self.end_layer]: + hidden_states, pending = layer( + positions, + hidden_states, + pending=pending, + defer_mlp_add=True, + attn_in=attn_in0, + log_scaling=log_scaling, + ) + attn_in0 = None + + if not get_pp_group().is_last_rank: + if pending is not None: + hidden_states = _sconv_add_norm( + pending[0], hidden_states, pending[1], None, positions + )[1] + return IntermediateTensors({"hidden_states": hidden_states}) + if pending is not None: + # Final RS/sconv/AG + residual add fused with the final rmsnorm. + norm_out = _sconv_add_norm( + pending[0], hidden_states, pending[1], self.norm, positions + )[0] + assert norm_out is not None + return norm_out + return self.norm(hidden_states) + + +class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA): + """Shared text-backbone causal-LM scaffolding for both entry classes.""" + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_substr={ + ".w13_dn": ".gate_up_proj", + ".w2_md": ".down_proj", + }, + orig_to_new_stacked={ + ".attn.wq_du.": (".attn.qkvr.", 0), + ".attn.wk_dv.": (".attn.qkvr.", 1), + ".attn.wv_dv.": (".attn.qkvr.", 2), + ".attn.wr_du.": (".attn.qkvr.", 3), + }, + orig_to_new_prefix={ + "model.llm.layers.": "model.layers.", + "model.llm.embed_norm": "model.embed_norm", + "model.llm.embed": "model.embed_tokens", + "model.llm.norm": "model.norm", + "model.llm.unembed": "lm_head", + "language_model.layers.": "model.layers.", + "language_model.lm_head.": "lm_head.", + }, + orig_to_new_suffix={ + # NVFP4 scale + ".w13_weight.scale": ".w13_weight_scale", + ".w13_weight.scale2": ".w13_weight_scale_2", + ".w2_weight.scale": ".w2_weight_scale", + ".w2_weight.scale2": ".w2_weight_scale_2", + }, + ) + packed_modules_mapping = { + "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"], + "w13": ["w1", "w3"], + } + embedding_modules = { + "lm_head": "output_embeddings", + } + + def _build( + self, + vllm_config: VllmConfig, + text_config: InklingModelConfig, + prefix: str, + ) -> None: + quant_config = vllm_config.quant_config + self.config = text_config + # NVFP4 experts are detected directly from the checkpoint quant config; + # only the MoE experts are quantized (attention/dense MLP stay bf16). + self.nvfp4_config = InklingNvfp4Config.from_hf_config( + vllm_config.model_config.hf_config + ) + # Read by the MRV2 runner to publish per-request short-conv metadata. + # Short convolution is intrinsic to Inkling, so this is always set. + self.uses_sconv = True + self.model = InklingModel( + config=text_config, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "model"), + nvfp4_config=self.nvfp4_config, + ) + initialize_lamport_rs_conv( + text_config.hidden_size, + text_config.sconv_kernel_size, + vllm_config.scheduler_config.max_num_batched_tokens, + ) + self.lm_head = ParallelLMHead( + text_config.padded_vocab_size, + text_config.hidden_size, + org_num_embeddings=text_config.padded_vocab_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = InklingLogitsProcessor( + text_config.padded_vocab_size, + org_vocab_size=text_config.vocab_size, + soft_cap=text_config.final_logit_softcapping, + logits_mup_width_multiplier=text_config.logits_mup_width_multiplier, + ) + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return _load_inkling_weights(self, weights, self.config) + + +class InklingForCausalLM(_TmlForCausalLMBase): + """Text-only entry point (``inkling_model`` checkpoints).""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self._build(vllm_config, vllm_config.model_config.hf_config, prefix) + + +@MULTIMODAL_REGISTRY.register_processor( + InklingMultiModalProcessor, + info=InklingProcessingInfo, + dummy_inputs=InklingDummyInputsBuilder, +) +class InklingForConditionalGeneration(_TmlForCausalLMBase, SupportsMultiModal): + """Top-level (multimodal) entry point. + + Builds the vision + audio towers on top of the shared text backbone. Inkling has + NO cross-modal fusion (the vision tower emits one token per patch, the audio + tower one token per frame), so generation reuses the inherited backbone + ``forward`` / ``compute_logits`` (the latter already applies muP) and this + class only adds multimodal embedding + merge. + """ + + hf_to_vllm_mapper = _TmlForCausalLMBase.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={ + "model.audio.": "audio.", + "model.visual.": "visual.vision_encoder.", + }, + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return "<|content_image|>" + if modality.startswith("audio"): + return "<|content_audio_input|>" + raise ValueError("Only image or audio modality is supported") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config: InklingMMConfig = vllm_config.model_config.hf_config + + self.visual = ( + InklingVision(config.vision_config, prefix=maybe_prefix(prefix, "visual")) + if inkling_vision_enabled(config) + else None + ) + self.audio = ( + InklingAudio(config.audio_config, prefix=maybe_prefix(prefix, "audio")) + if inkling_audio_enabled(config) + else None + ) + + self._build(vllm_config, config.text_config, prefix) + + # -- multimodal embedding ------------------------------------------- + + def _process_image_input( + self, pixel_values: Any, num_patches: Any + ) -> tuple[torch.Tensor, ...]: + assert self.visual is not None + # pixel_values is a list (per item) of [P_i, 2, P, P, 3] tensors, + # or a single concatenated tensor. Normalize to a flat batch, run the + # tower once, then split back per item. + if isinstance(pixel_values, (list, tuple)): + if not pixel_values: + return () + sizes = [int(p.shape[0]) for p in pixel_values] + patches = torch.cat(list(pixel_values), dim=0) + else: + patches = pixel_values + sizes = self._sizes_from(num_patches, patches.shape[0]) + + patches = patches.to(device=self.visual.device, dtype=self.visual.dtype) + embeds = self.visual(patches) # [total_patches, D] + return tuple(embeds.split(sizes)) + + def _process_audio_input( + self, input_audio_features: Any, num_audio_tokens: Any + ) -> tuple[torch.Tensor, ...]: + assert self.audio is not None + if isinstance(input_audio_features, (list, tuple)): + if not input_audio_features: + return () + sizes = [int(d.shape[0]) for d in input_audio_features] + dmel = torch.cat(list(input_audio_features), dim=0) + else: + dmel = input_audio_features + sizes = self._sizes_from(num_audio_tokens, dmel.shape[0]) + + dmel = dmel.to(device=self.audio.device) + embeds = self.audio(dmel) # [total_frames, D] + return tuple(embeds.split(sizes)) + + @staticmethod + def _sizes_from(counts: Any, total: int) -> list[int]: + if counts is None: + return [total] + if isinstance(counts, torch.Tensor): + return [int(c) for c in counts.flatten().tolist()] + if isinstance(counts, (list, tuple)): + flat: list[int] = [] + for c in counts: + flat.append(int(c.item()) if isinstance(c, torch.Tensor) else int(c)) + return flat + return [int(counts)] + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + # Iterate modalities in a stable order so the returned per-item tensors + # line up with their appearance order; the positional merge in + # embed_input_ids handles actual placement. + pixel_values = kwargs.get("pixel_values") + num_patches = kwargs.get("num_patches") + input_audio_features = kwargs.get("input_audio_features") + num_audio_tokens = kwargs.get("num_audio_tokens") + + embeddings: tuple[torch.Tensor, ...] = () + if pixel_values is not None and self.visual is not None: + embeddings += self._process_image_input(pixel_values, num_patches) + if input_audio_features is not None and self.audio is not None: + embeddings += self._process_audio_input( + input_audio_features, num_audio_tokens + ) + return embeddings + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + # Override the base's 1-arg embed_input_ids: the runner calls this 3-arg + # signature for multimodal models. Text embeddings come from the shared + # backbone (which applies embed_norm); MM embeddings are scattered in. + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + # Placeholder ids use unused vocabulary slots and these positions are + # overwritten by MM embeds below. + inputs_embeds = self.model.embed_input_ids(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def get_language_model(self) -> nn.Module: + # This class IS the causal LM (the towers are side branches), so the + # language model is self — callers expect a module exposing ``.model`` + # and ``.lm_head``. + return self + + +# =========================================================================== +# Weight loading +# =========================================================================== + + +_MOE_EXPERT_WEIGHT_RE = re.compile( + r"^(?P.*\.mlp)\.(?P(?:shared_)?experts\..+)$" +) + + +def _load_inkling_weights( + module: nn.Module, + weights: Iterable[tuple[str, torch.Tensor]], + config: InklingModelConfig, +) -> set[str]: + moe_modules = { + name: mod for name, mod in module.named_modules() if isinstance(mod, InklingMoE) + } + loaded: set[str] = set() + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + local_ids = set(config.local_layer_ids) + + def _iter_loadable_weights() -> Iterable[tuple[str, torch.Tensor]]: + for name, weight in module.hf_to_vllm_mapper.apply(weights): + shard_id = getattr(weight, "shard_id", None) + # Replicate K/V conv-free GQA heads when tp_size > num_kv_heads. + if ( + shard_id in (1, 2) + and name.endswith(".attn.qkvr.weight") + and weight.shape[0] > 0 + ): + lid = _layer_id(name) + if lid is not None: + is_local = lid in local_ids + n_kv = ( + config.swa_num_key_value_heads + if is_local + else config.num_key_value_heads + ) + head_dim = config.swa_head_dim if is_local else config.head_dim + if tp_size > n_kv and weight.shape[0] == n_kv * head_dim: + kv_idx = (tp_rank * n_kv) // tp_size + weight = weight.narrow(0, kv_idx * head_dim, head_dim) + weight.shard_id = shard_id + + # MoE expert tensors (fused stacked, routed + shared sink): translate + # the checkpoint layout to per-expert FusedMoE loads. + moe_match = _MOE_EXPERT_WEIGHT_RE.match(name) + if moe_match is not None and moe_match.group("mlp") in moe_modules: + moe = moe_modules[moe_match.group("mlp")] + for rel in moe.load_expert_weight(moe_match.group("rest"), weight): + loaded.add(f"{moe_match.group('mlp')}.{rel}") + continue + + yield name, weight + + # The release checkpoint also carries auxiliary prediction-head weights; + # they are not part of the causal LM served by this implementation. + loader = AutoWeightsLoader(module, skip_prefixes=["model.mtp."]) + loaded |= loader.load_weights(_iter_loadable_weights()) + + # Post-load MoE fixups (default input scales, zeroed EP-padding experts). + for moe_name, moe in moe_modules.items(): + for rel in moe.finalize_load(): + loaded.add(f"{moe_name}.{rel}") + return loaded + + +EntryClass = [InklingForCausalLM, InklingForConditionalGeneration] diff --git a/vllm/models/inkling/nvidia/moe.py b/vllm/models/inkling/nvidia/moe.py new file mode 100644 index 000000000000..255c9328c8aa --- /dev/null +++ b/vllm/models/inkling/nvidia/moe.py @@ -0,0 +1,648 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling mixture-of-experts on vLLM's FusedMoE abstraction. + +Overfit to the served checkpoint: sigmoid gate (+ selection bias) top-k over +the routed experts, log-sigmoid renormalization over the k routed + S shared +"sink" logits, scaled by route_scale * global_scale. The routed top-k goes +through vLLM's FusedMoE (which handles TP/EP); the sink experts run in +:class:`InklingSinkExperts` -- replicated across EP ranks (every token +activates every sink) and always bf16 (the checkpoint excludes every +``shared_experts`` from quantization). + +NVFP4 routed experts reuse vLLM's ModelOpt NVFP4 fused-MoE method; excluded +(bf16) layers fall back to the unquantized method. The checkpoint's fused +stacked tensors (interleaved gate/up rows, ``.scale`` / ``.scale2`` / +``.input_amax`` aux tensors) are translated to the standard per-expert loads +in :meth:`InklingMoE.load_expert_weight`. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch +from torch import nn +from torch.nn.parameter import Parameter + +import vllm.envs as envs +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_dp_group, + get_pcp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16 +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform +from vllm.triton_utils import tl, tldevice, triton +from vllm.utils.multi_stream_utils import maybe_execute_in_parallel +from vllm.utils.torch_utils import aux_stream + +from ..configs import InklingModelConfig +from ..nvfp4 import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import ( + RoutedExperts, + ) + + from ..nvfp4 import InklingNvfp4Config + +# --------------------------------------------------------------------------- +# Gate / expert selection +# --------------------------------------------------------------------------- + +_INKLING_LL_BF16_MAX_TOKENS = 64 + + +def _linear_with_fp32_out(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + leading = list(x.shape[:-1]) + flat = x.flatten(0, -2) + if ( + flat.shape[0] <= _INKLING_LL_BF16_MAX_TOKENS + and flat.dtype == torch.bfloat16 + and weight.dtype == torch.bfloat16 + and flat.is_cuda + and flat.is_contiguous() + and weight.is_contiguous() + and flat.shape[1] % 8 == 0 + and current_platform.has_device_capability(90) + and ll_bf16.is_available() + ): + out = ll_bf16.ll_bf16_gemm(flat, weight) + else: + out = torch.mm(flat, weight.T, out_dtype=torch.float32) + return out.view(*leading, weight.shape[0]) + + +@triton.jit(do_not_specialize=["T", "route_scale"]) +def _inkling_gate_select_kernel( + logits_ptr, # [T, G] fp32 gate logits (stride_logits_0 may include pad) + bias_ptr, # [R] fp32 selection bias (or 0 ptr if HAS_BIAS=False) + global_scale_ptr, # [1] fp32 (or unused if HAS_GSCALE=False) + ids_ptr, # [T, K + S] int32 out: selected expert ids + weights_ptr, # [T, K + S] fp32 out: renormalized weights + route_scale, + T, + G: tl.constexpr, # total gate experts (routed + shared) + stride_logits_0, + R: tl.constexpr, # routed experts + K: tl.constexpr, # top-k routed + S: tl.constexpr, # shared (sink) experts + HAS_BIAS: tl.constexpr, + HAS_GSCALE: tl.constexpr, + BLOCK_G: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + if pid >= T: + return + offs = tl.arange(0, BLOCK_G) + mask_r = offs < R + logits = tl.load( + logits_ptr + pid * stride_logits_0 + offs, + mask=offs < G, + other=float("-inf"), + ).to(tl.float32) + + # Selection scores: sigmoid(routed logits) (+ bias), non-routed lanes -inf. + sel = tl.where(mask_r, tl.sigmoid(logits), float("-inf")) + if HAS_BIAS: + bias = tl.load(bias_ptr + offs, mask=mask_r, other=0.0).to(tl.float32) + sel = tl.where(mask_r, sel + bias, float("-inf")) + + scale = route_scale + if HAS_GSCALE: + scale = scale * tl.load(global_scale_ptr).to(tl.float32) + + # Iterative top-K (K is small); argmax tie-breaks to the lowest index + # (stable ordering). + A: tl.constexpr = K + S + offs_a = tl.arange(0, A) + top_ids = tl.zeros([A], dtype=tl.int32) + active = tl.zeros([A], dtype=tl.float32) + for kk in tl.static_range(K): + idx = tl.argmax(sel, axis=0).to(tl.int32) + raw = tl.max(tl.where(offs == idx, logits, float("-inf")), axis=0) + top_ids = tl.where(offs_a == kk, idx, top_ids) + active = tl.where(offs_a == kk, raw, active) + sel = tl.where(offs == idx, float("-inf"), sel) + if S > 0: + # Shared sink logits sit at the tail of the gate output; their expert + # ids continue after the routed range (R + j). + for jj in tl.static_range(S): + raw = tl.max(tl.where(offs == R + jj, logits, float("-inf")), axis=0) + top_ids = tl.where(offs_a == K + jj, tl.full([], R + jj, tl.int32), top_ids) + active = tl.where(offs_a == K + jj, raw, active) + + # Log-sigmoid renormalization over the K + S active logits. + abs_l = tl.abs(active) + min_l = tl.minimum(active, 0.0) + log_probs = min_l - tldevice.log1p(tldevice.exp(-abs_l)) + max_lp = tl.max(log_probs, axis=0) + exp_shifted = tldevice.exp(log_probs - max_lp) + sum_exp = tl.sum(exp_shifted, axis=0) + weights = exp_shifted / sum_exp * scale + + tl.store(ids_ptr + pid * A + offs_a, top_ids) + tl.store(weights_ptr + pid * A + offs_a, weights) + + +def inkling_gate_select( + logits: torch.Tensor, # [T, >=G] fp32 (rows may carry GEMM padding) + n_gate_experts: int, + n_routed_experts: int, + topk: int, + n_shared_experts: int, + bias: torch.Tensor | None, + route_scale: float, + global_scale: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sigmoid + bias + top-k + log-sigmoid renorm; returns (weights, ids).""" + assert logits.dtype == torch.float32 + tokens = logits.shape[0] + active = topk + n_shared_experts + topk_ids = torch.empty((tokens, active), dtype=torch.int32, device=logits.device) + topk_weights = torch.empty( + (tokens, active), dtype=torch.float32, device=logits.device + ) + if tokens == 0: + return topk_weights, topk_ids + _inkling_gate_select_kernel[(tokens,)]( + logits, + bias if bias is not None else logits, + global_scale if global_scale is not None else logits, + topk_ids, + topk_weights, + route_scale, + tokens, + n_gate_experts, + logits.stride(0), + n_routed_experts, + topk, + n_shared_experts, + HAS_BIAS=bias is not None, + HAS_GSCALE=global_scale is not None, + BLOCK_G=triton.next_power_of_2(n_gate_experts), + ) + return topk_weights, topk_ids + + +class InklingGate(nn.Module): + """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and + global scale (the served checkpoint's only configuration).""" + + def __init__( + self, + d_model: int, + n_routed_experts: int, + n_shared_experts: int, + experts_per_token: int, + route_scale: float, + *, + use_global_scale: bool = False, + use_gate_bias: bool = False, + ) -> None: + super().__init__() + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.n_total_experts = n_routed_experts + n_shared_experts + self.topk = experts_per_token + self.route_scale = route_scale + + padded_experts = self.n_total_experts + (-self.n_total_experts) % 8 + self.weight = Parameter( + torch.empty(padded_experts, d_model), requires_grad=False + ) + set_weight_attrs(self.weight, {"weight_loader": self._load_weight}) + if use_global_scale: + self.global_scale = Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + else: + self.global_scale = None + if use_gate_bias: + self.bias = Parameter( + torch.empty(n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + else: + self.bias = None + + @staticmethod + def _load_weight(param: Parameter, loaded_weight: torch.Tensor) -> None: + param.data.zero_() + param.data[: loaded_weight.shape[0]].copy_(loaded_weight) + + def compute_logits(self, x: torch.Tensor) -> torch.Tensor: + """fp32 gate logits [T, n_total_experts + pad] (pad columns are junk).""" + return _linear_with_fp32_out(x, self.weight) + + def select_experts( + self, gating_output: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Full selection: (weights, ids) of [T, K + S]. The first K entries + are the routed top-k; the S trailing entries are the sink gammas.""" + return inkling_gate_select( + gating_output, + self.n_total_experts, + self.n_routed_experts, + self.topk, + self.n_shared_experts, + self.bias, + self.route_scale, + self.global_scale, + ) + + +# --------------------------------------------------------------------------- +# MoE layer +# --------------------------------------------------------------------------- + + +def _inkling_moe_ep_size() -> int: + """EP size the FusedMoE layer will run with (mirrors + FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when + expert parallelism is enabled).""" + parallel_config = get_current_vllm_config().parallel_config + if not parallel_config.enable_expert_parallel: + return 1 + world = ( + get_tensor_model_parallel_world_size() + * get_dp_group().world_size + * get_pcp_group().world_size + ) + return world if world > 1 else 1 + + +class InklingSinkExperts(nn.Module): + """Shared "sink" experts with per-token gammas, in bf16. + + Replicated across EP ranks (every token activates every sink, so + EP-sharding them would hotspot the owning rank) and TP-sharded on the + intermediate dim so the output remains a TP-partial sum like the routed + output. The sinks are always bf16 (the checkpoint excludes every + ``shared_experts`` from quantization): the experts concatenate into two + plain dense GEMMs with the fused sink epilogue between them. + """ + + def __init__( + self, n_experts: int, d_model: int, d_mlp: int, *, prefix: str = "" + ) -> None: + super().__init__() + self.n_experts = n_experts + tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = get_tensor_model_parallel_rank() + intermediate_pp = d_mlp // tp_size + self.w13_weight = Parameter( + torch.empty(n_experts, 2 * intermediate_pp, d_model), + requires_grad=False, + ) + self.w2_weight = Parameter( + torch.empty(d_model, n_experts * intermediate_pp), + requires_grad=False, + ) + self._unit: torch.Tensor | None = None + + def load_weight(self, key: str, weight: torch.Tensor) -> list[str]: + """Load one checkpoint sink tensor (stacked over the S experts).""" + if key == "w13_weight": + if weight.shape != self.w13_weight.shape: + shard = self.w13_weight.shape[1] + weight = weight.narrow(1, self.tp_rank * shard, shard) + self.w13_weight.data.copy_(weight) + return [key] + + assert key == "w2_weight" + shard = self.w2_weight.shape[1] // self.n_experts + shard_start = 0 if weight.shape[2] == shard else self.tp_rank * shard + for expert_idx, expert_weight in enumerate(weight): + local_weight = expert_weight.narrow(1, shard_start, shard) + start = expert_idx * shard + self.w2_weight.data[:, start : start + shard].copy_(local_weight) + return [key] + + def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: + """``sum_e gammas[:, e] * MLP_e(x)`` (TP-partial along d_mlp).""" + from .ops import sink_silu_mul_epilogue + + # One GEMM over the experts' stacked w13 (a view), fused epilogue, + # then one GEMM whose K-reduction over the K-concatenated w2 performs + # the expert sum. + if self._unit is None or self._unit.device != x.device: + self._unit = torch.ones( + self.n_experts, dtype=torch.float32, device=x.device + ) + raw = x @ self.w13_weight.view(-1, x.shape[-1]).T # (T, S*2F) + h = sink_silu_mul_epilogue( + raw, self._unit, gammas, self._unit, self.n_experts, x.dtype + ) + return h @ self.w2_weight.T # (T, D) + + +class InklingSinkExpertsLinear(nn.Module): + """LoRA-capable implementation of the Inkling sink experts.""" + + def __init__( + self, + n_experts: int, + d_model: int, + d_mlp: int, + *, + prefix: str = "", + ) -> None: + super().__init__() + from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, + ) + + self.n_experts = n_experts + self.d_mlp = d_mlp + total = n_experts * d_mlp + self.w13 = MergedColumnParallelLinear( + input_size=d_model, + output_sizes=[total, total], + bias=False, + prefix=f"{prefix}.w13", + ) + self.w2 = RowParallelLinear( + input_size=total, + output_size=d_model, + bias=False, + reduce_results=False, + prefix=f"{prefix}.w2", + ) + self._w2_input_pp = self.w2.input_size_per_partition + self._col_expert: torch.Tensor | None = None + + def _gamma_expand(self, gammas: torch.Tensor) -> torch.Tensor: + if self._col_expert is None or self._col_expert.device != gammas.device: + local = self._w2_input_pp + start = get_tensor_model_parallel_rank() * local + cols = torch.arange(start, start + local, device=gammas.device) + self._col_expert = (cols // self.d_mlp).long() + return gammas[:, self._col_expert] + + def load_weight(self, key: str, weight: torch.Tensor) -> list[str]: + if key == "w13_weight": + d_model = weight.shape[-1] + gate = weight[:, 0::2, :].reshape(-1, d_model).contiguous() + up = weight[:, 1::2, :].reshape(-1, d_model).contiguous() + self.w13.weight_loader(self.w13.weight, gate, 0) + self.w13.weight_loader(self.w13.weight, up, 1) + return ["w13.weight"] + w = weight.permute(1, 0, 2).reshape(weight.shape[1], -1).contiguous() + self.w2.weight_loader(self.w2.weight, w) + return ["w2.weight"] + + def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.w13(x) + gate, up = gate_up.chunk(2, dim=-1) + hidden_states = torch.nn.functional.silu(gate) * up + hidden_states = (hidden_states * self._gamma_expand(gammas)).to(x.dtype) + output, _ = self.w2(hidden_states) + return output + + +class InklingMoE(nn.Module): + def __init__( + self, + config: InklingModelConfig, + layer_id: int, + *, + prefix: str = "", + nvfp4_config: InklingNvfp4Config | None = None, + ) -> None: + super().__init__() + # Overfit to the served checkpoint: sigmoid gate renormalized after + # top-k, shared sink experts, interleaved gate/up checkpoint rows. + assert config.gate_activation == "sigmoid" and config.norm_after_topk + assert config.n_shared_experts > 0 and config.shared_expert_sink + assert config.inference_moe_w13_interleaved + n_routed = config.n_routed_experts + n_shared = config.n_shared_experts + self.n_routed_experts = n_routed + self.gate = InklingGate( + d_model=config.hidden_size, + n_routed_experts=n_routed, + n_shared_experts=n_shared, + experts_per_token=config.num_experts_per_tok, + route_scale=config.route_scale, + use_global_scale=config.use_global_scale, + use_gate_bias=config.use_gate_bias, + ) + + moe_quant_config = None + if nvfp4_config is not None and nvfp4_config.experts_quantized(layer_id): + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4Config, + ) + + # The Inkling checkpoint is ModelOpt NVFP4; exclusion is decided per + # layer right here, so no exclude list is needed. + moe_quant_config = ModelOptNvFp4Config( + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=True, + kv_cache_quant_algo=None, + exclude_modules=[], + group_size=nvfp4_config.group_size, + ) + + # TRTLLM MoE kernels assume equal, contiguous per-rank expert slabs + # (local_expert_offset = ep_rank * local_num_experts), so pad the + # expert count to a multiple of the EP size. A no-op for the usual + # power-of-two EP sizes (n_routed is a power of two). + num_experts = n_routed + (-n_routed) % _inkling_moe_ep_size() + + self.experts = FusedMoE( + num_experts=num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + renormalize=False, + quant_config=moe_quant_config, + prefix=f"{prefix}.experts", + custom_routing_function=self._select_routed, + router_logits_dtype=torch.float32, + activation="silu", + ) + # The decoder layer reduce-scatters the MoE delta into the sconv + # stream itself (RS -> shard sconv -> AG); the runner must return the + # per-rank partial sum instead of all-reducing. + self.experts.moe_config.skip_final_all_reduce = True + + self._routed_sel: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + # The sinks are always bf16; fail loudly on a checkpoint that + # quantizes them instead of silently misloading. + assert nvfp4_config is None or not nvfp4_config.shared_experts_quantized( + layer_id + ), f"layer {layer_id}: NVFP4 shared experts are not supported" + + sink_experts_cls = ( + InklingSinkExpertsLinear + if get_current_vllm_config().lora_config is not None + else InklingSinkExperts + ) + self.sink_experts = sink_experts_cls( + n_experts=n_shared, + d_model=config.hidden_size, + d_mlp=config.intermediate_size, + prefix=f"{prefix}.shared_experts", + ) + + # Sink chain overlaps the routed MoE call on the aux stream for + # decode-sized batches (same pattern as the runner's SharedExperts + # multi-stream overlap). The routed GEMM runs on the default stream and + # the sink chain on the aux stream, joined via these two events by + # ``maybe_execute_in_parallel``. + self._sink_stream: torch.cuda.Stream | None = aux_stream() + self._sink_events = (torch.cuda.Event(), torch.cuda.Event()) + + def _select_routed( + self, + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + """FusedMoE ``custom_routing_function``: the routed top-k slice of the + full (routed + sink) selection. + + forward() stashes its selection (keyed by logits identity) so the + gate select runs once per layer; the fallback covers paths where the + runner re-derives the logits (e.g. naive DP dispatch). + """ + del hidden_states, renormalize + assert topk == self.gate.topk + cached = self._routed_sel + self._routed_sel = None + if cached is not None and cached[0] is gating_output: + return cached[1], cached[2] + weights, ids = self.gate.select_experts(gating_output) + return weights[:, :topk].contiguous(), ids[:, :topk].contiguous() + + def forward(self, x: torch.Tensor) -> torch.Tensor | None: + router_logits = self.gate.compute_logits(x) + num_tokens = x.shape[0] + # One gate select per layer: the routed slice is stashed for the + # routing function inside the FusedMoE op; the sink gammas are the + # trailing columns. + k = self.gate.topk + weights, ids = self.gate.select_experts(router_logits) + self._routed_sel = ( + router_logits, + weights[:, :k].contiguous(), + ids[:, :k].contiguous(), + ) + gammas = weights[:, k:] + + out, sink_out = maybe_execute_in_parallel( + lambda: self.experts(hidden_states=x, router_logits=router_logits), + lambda: self.sink_experts(x, gammas), + self._sink_events[0], + self._sink_events[1], + self._sink_stream + if num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD + else None, + ) + self._routed_sel = None + + return out + sink_out + + # -- weight loading ---------------------------------------------------- + + def _local_expert_slots(self) -> dict[int, int]: + """Global expert id -> local slot for this rank's expert partition.""" + manager = self.experts.routed_experts.expert_map_manager + if manager.expert_map is None: + return {g: g for g in range(manager.global_num_experts)} + emap = manager.expert_map.tolist() + return {g: slot for g, slot in enumerate(emap) if slot >= 0} + + def load_expert_weight(self, name: str, weight: torch.Tensor) -> list[str]: + """Load one checkpoint expert tensor. + + ``name`` is relative to the mlp module: ``experts.`` (routed + stack) or ``shared_experts.shared_`` (sink experts). Returns the + loaded param names (relative to this module). + """ + if name.startswith("shared_experts."): + key = name.split(".", 1)[1].replace("shared_", "", 1) + return [ + f"sink_experts.{p}" for p in self.sink_experts.load_weight(key, weight) + ] + + experts: RoutedExperts = self.experts.routed_experts + key = name.split(".", 1)[1] + + # original_shape is unused by the vLLM serving layout. + if key.endswith(".original_shape"): + return [] + if key.endswith(".input_amax"): + projection = "w13" if key.startswith("w13") else "w2" + amax = float(weight.max()) + assert math.isfinite(amax) and amax > 0, ( + f"bad {projection} input_amax: {amax}" + ) + input_scale = getattr(experts, f"{projection}_input_scale") + input_scale.data.fill_(amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)) + return [f"experts.routed_experts.{projection}_input_scale"] + + param = getattr(experts, key) + slots = self._local_expert_slots() + gids = sorted(slots) + lids = [slots[g] for g in gids] + tp_rank = experts.moe_config.moe_parallel_config.tp_rank + + if key.endswith("_scale_2"): + # Per-expert scalars, vectorized over the local experts. The + # fused w13 param carries one slot per gate/up half. + vals = weight[gids].float().to(param.device) + param.data[lids] = vals[:, None] if param.data.ndim == 2 else vals + elif key.startswith("w13"): + # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the + # fused param layout is [w1(gate); w3(up)]. The TP-local rows form + # one contiguous slab of the interleaved tensor, so upload just + # that slab (a single bounded synchronous H2D; pre-uploading whole + # untrimmed tensors pins the mmap pages of the entire checkpoint + # and OOMs the host) and de-interleave on device. + half = param.shape[1] // 2 + for gid, lid in slots.items(): + slab = weight[gid].narrow(0, tp_rank * 2 * half, 2 * half) + slab = slab.to(param.device) + param.data[lid, :half].copy_(slab[0::2]) + param.data[lid, half:].copy_(slab[1::2]) + else: + # w2: shard the packed intermediate (last) dim. + shard = param.shape[2] + for gid, lid in slots.items(): + param.data[lid].copy_(weight[gid].narrow(1, tp_rank * shard, shard)) + return [f"experts.routed_experts.{key}"] + + def finalize_load(self) -> list[str]: + """Post-load fixups for zeroed padding experts.""" + experts = self.experts.routed_experts + out: list[str] = [] + # Zero the EP-alignment padding experts (if any) so their + # (never-routed) slots hold defined values. + slots = self._local_expert_slots() + for gid in range(self.n_routed_experts, experts.global_num_experts): + lid = slots.get(gid) + if lid is None: + continue + for pname in ( + "w13_weight", + "w2_weight", + "w13_weight_scale", + "w2_weight_scale", + "w13_weight_scale_2", + "w2_weight_scale_2", + ): + p = getattr(experts, pname, None) + if p is not None: + p.data[lid].zero_() + return out diff --git a/vllm/models/inkling/nvidia/mtp.py b/vllm/models/inkling/nvidia/mtp.py new file mode 100644 index 000000000000..868ede3bed64 --- /dev/null +++ b/vllm/models/inkling/nvidia/mtp.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling MTP (Multi-Token Prediction) draft model (NVIDIA). + +Implements the first MTP depth from the reference ``mtp_model.py`` shipped with +the checkpoint. It owns ``hidden_norm`` / ``embed_norm`` RMSNorms, a ``2H -> H`` +input projection, and a full Inkling transformer block with a dense bf16 MLP. + +The draft shares the target's token embedding table and LM head +(``load_eagle_model`` wires those references) and applies the backbone +``embed_norm`` on top: the depth layers were trained on the same normed +embeddings the backbone consumes (their own ``embed_norm`` weights are +near-identity trims, unlike the backbone's whitening ``embed_norm``). +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import regex as re +import torch +from torch import nn + +from vllm.config import VllmConfig +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.utils import maybe_prefix +from vllm.sequence import IntermediateTensors + +from ..configs import InklingModelConfig +from .layernorm import InklingRMSNorm +from .model import InklingDecoderLayer, InklingReplicatedEmbedding +from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm + +# Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr. +# Mirrors the backbone's hf_to_vllm_mapper.orig_to_new_stacked; kept as a +# local (pname, wname, shard) list since the MTP loader remaps by hand. +_ATTENTION_PARAMS_MAPPING = [ + ("qkvr", "wq_du", 0), + ("qkvr", "wk_dv", 1), + ("qkvr", "wv_dv", 2), + ("qkvr", "wr_du", 3), +] + + +def _mtp_depth_from_name(name: str) -> int | None: + m = re.search(r"\.mtp\.layers\.(\d+)\.", name) + return int(m.group(1)) if m else None + + +class InklingMTPDepthLayer(nn.Module): + """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" + + def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None: + super().__init__() + self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_proj = ReplicatedLinear( + config.hidden_size * 2, + config.hidden_size, + bias=False, + return_bias=False, + prefix=f"{prefix}.input_proj", + ) + # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs + # full attention (the swa_* head config and sliding_window window) to + # match this depth's checkpoint transformer_block weights. + self.transformer_block = InklingDecoderLayer( + config, + layer_id=0, + is_local=is_local, + quant_config=None, + prefix=f"{prefix}.transformer_block", + nvfp4_config=None, + force_dense_mlp=True, + ) + + def forward(self, combined: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # ``combined`` is the fused-normed [rmsnorm(hidden) | embed_norm(emb)] + # input, built by InklingMultiTokenPredictor.fused_input_cat in one launch. + hidden = self.input_proj(combined) + # The short conv self-fetches its paged SWA-cache metadata from the + # forward context (via its conv_owner prefix); no conv_meta to thread. + return self.transformer_block(positions, hidden) + + +class InklingMultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config: InklingModelConfig = ( + vllm_config.speculative_config.draft_model_config.hf_config + ) + self.config = config + if vllm_config.speculative_config.num_speculative_tokens != 1: + raise ValueError( + "Inkling MTP currently supports exactly one speculative token" + ) + self.chain_hidden_post_norm = config.chain_hidden_post_norm + local_ids = set(config.local_layer_ids) + self.layers = nn.ModuleDict( + {"0": InklingMTPDepthLayer(config, f"{prefix}.layers.0", 0 in local_ids)} + ) + self.chain_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if self.chain_hidden_post_norm + else None + ) + # The target's raw token embedding (pre embed_norm), attached by + # load_eagle_model. Never materialized here: building our own + # replicated copy would transiently double the 2.3 GiB table. + self.embed_tokens: InklingReplicatedEmbedding = None # type: ignore[assignment] + # The depth layers consume the *backbone-normed* embedding + # (embed_norm(embed(ids))), not the raw one: mtp embed_norm weights + # are near-identity (trained on already-normalized inputs), and + # feeding raw embeddings drops MTP1 acceptance from ~0.85 to ~0.70. + # Weight loaded from the target's embed_norm.weight; gated like the + # target's InklingModel.embed_norm. + self.backbone_embed_norm = ( + InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if config.use_embed_norm + else None + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: object | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + """Draft-prefill embedding: fused gather + backbone embed_norm, then + the target's tower embeddings scattered in unnormed (the backbone + convention — MM embeds are merged after embed_norm).""" + norm = self.backbone_embed_norm + embeds = embed_rmsnorm( + input_ids, + self.embed_tokens.weight, + norm.weight if norm is not None else None, + norm.variance_epsilon if norm is not None else 0.0, + ) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: # type: ignore[arg-type] + return embeds + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def fused_input_cat( + self, + layer: InklingMTPDepthLayer, + previous_hidden: torch.Tensor, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor | None, + ) -> torch.Tensor: + """The depth layer's [rmsnorm(hidden) | embed_norm(embed)] input in one + launch: embedding row gather + the backbone embed_norm + the depth + embed_norm chain on one side, hidden_norm on the other, written + straight into the cat buffer.""" + hidden_w = layer.hidden_norm.weight + embed_w = layer.embed_norm.weight + eps = layer.hidden_norm.variance_epsilon + if inputs_embeds is not None: + # Draft prefill with target-merged MM embeddings (already + # backbone-normed via embed_input_ids); only the depth embed_norm + # remains. + return embed_dual_rmsnorm_cat( + previous_hidden, hidden_w, embed_w, eps, embeds=inputs_embeds + ) + return embed_dual_rmsnorm_cat( + previous_hidden, + hidden_w, + embed_w, + eps, + input_ids=input_ids, + embed_table=self.embed_tokens.weight, + pre_norm_weight=( + self.backbone_embed_norm.weight + if self.backbone_embed_norm is not None + else None + ), + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + # The draft's short conv is a paged SWA-cache layer (its conv_owner is + # auto-enumerated as a draft attention layer); its per-token metadata is + # built by the speculator's build_attn_metadata and read from the + # forward context, so nothing extra is threaded here. + if spec_step_idx != 0: + raise ValueError("Inkling MTP only supports spec_step_idx=0") + layer = self.layers["0"] + combined = self.fused_input_cat( + layer, previous_hidden_states, input_ids, inputs_embeds + ) + hidden = layer(combined, positions) + if self.chain_norm is not None: + hidden = self.chain_norm(hidden) + return hidden + + +class InklingMTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config: InklingModelConfig = ( + vllm_config.speculative_config.draft_model_config.hf_config + ) + self.config = config + self.model = InklingMultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + # The target's (vocab-sharded) LM head, attached by load_eagle_model; + # never materialized here (same reasoning as model.embed_tokens). + self.lm_head: ParallelLMHead = None # type: ignore[assignment] + self.logits_processor = LogitsProcessor( + config.padded_vocab_size, + org_vocab_size=config.vocab_size, + soft_cap=config.final_logit_softcapping, + ) + self._logits_zero: torch.Tensor | None = None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: object | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model.embed_input_ids( + input_ids, multimodal_embeddings, is_multimodal=is_multimodal + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, + positions, + hidden_states, + inputs_embeds, + spec_step_idx, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + # The MTP shares the base model's LM head, which is trained on + # ``hidden / mup``-scaled inputs, so apply the same mup scaling here + # for a matching logit scale — folded into the lm_head GEMM alpha + # (fp32 epilogue) like the target's compute_logits. (Argmax-invariant + # for greedy draft sampling, but it matters for the gumbel sampling + # distribution at temperature > 0.) + mup = self.config.logits_mup_width_multiplier + if not mup: + return self.logits_processor(self.lm_head, hidden_states) + assert self.logits_processor.soft_cap is None + assert self.logits_processor.scale == 1.0 + w = self.lm_head.weight + if self._logits_zero is None: + self._logits_zero = w.new_zeros(1) + logits = torch.addmm( + self._logits_zero, + hidden_states, + w.t(), + beta=0.0, + alpha=1.0 / mup, + ) + logits = self.logits_processor._gather_logits(logits) + if logits is not None: + logits = logits[..., : self.logits_processor.org_vocab_size] + return logits + + def get_top_tokens(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Greedy draft tokens via rank-local argmax + tiny (value, index) + reduction — no full-vocab logits all-gather. The muP divisor is a + positive scalar, so the argmax is invariant and the scaling is + skipped entirely.""" + return self.logits_processor.get_top_tokens(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return _load_inkling_mtp_weights(self, weights) + + +def _load_inkling_mtp_weights( + module: InklingMTP, + weights: Iterable[tuple[str, torch.Tensor]], +) -> set[str]: + """Load ``model.mtp.*`` weights into the MTP module. + + Checkpoint keys look like ``model.mtp.chain_norm.weight`` and + ``model.mtp.layers.{i}.{...}``. The transformer block reuses the backbone + layer's fused-projection layout, so we apply the same qkvr / gate_up / down + remapping as ``_load_inkling_weights``. Token embedding and LM head are shared + (provided by ``load_eagle_model``) and are not present in mtp.safetensors. + """ + # Per-depth attention is full or sliding-window (config.local_layer_ids); + # each depth's qkvr MergedColumnParallelLinear is built with the matching + # (swa_)num_key_value_heads, and its weight_loader handles the TP sharding. + # The sconv SWA cache pins tp_size <= num_key_value_heads, so tp never + # exceeds a layer's kv-head count and no GQA K/V replication is needed here. + params = dict(module.named_parameters()) + loaded: set[str] = set() + + def _load(name: str, weight: torch.Tensor, shard_id: object = None) -> bool: + param = params.get(name) + if param is None: + return False + loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: + if loader is default_weight_loader or param.shape == weight.shape: + default_weight_loader(param, weight) + else: + loader(param, weight) + else: + loader(param, weight, shard_id) # type: ignore[call-arg] + loaded.add(name) + return True + + for name, weight in weights: + depth = _mtp_depth_from_name(name) + # Token embedding and LM head are never materialized on the draft + # (no params to load into); load_eagle_model attaches the target's. + if name in ("model.llm.embed.weight", "model.llm.unembed.weight"): + continue + # The backbone embed_norm, applied to the shared embedding before the + # depth layers (see InklingMultiTokenPredictor.embed_input_ids). The + # per-depth mtp.layers.{i}.embed_norm keys carry ".mtp." and are loaded + # below. Only the shared backbone key routes here. + if name == "model.llm.embed_norm.weight": + _load("model.backbone_embed_norm.weight", weight) + continue + # Only consume the MTP weights; everything else belongs to the target. + if ".mtp." not in name: + continue + # Only the first checkpoint depth is used for MTP=1. + if depth is not None and depth != 0: + continue + # model.mtp.chain_norm.weight -> model.chain_norm.weight + # model.mtp.layers.{i}.X -> model.layers.{i}.X + original_name = name + name = name.replace(".mtp.layers.", ".layers.").replace( + ".mtp.chain_norm.", ".chain_norm." + ) + + if ".chain_norm." in name and module.model.chain_norm is None: + raise ValueError( + "Inkling checkpoint contains chain_norm weights but " + "chain_hidden_post_norm is disabled." + ) + + # Fused attention qkvr (wq_du/wk_dv/wv_dv/wr_du -> qkvr). + matched = False + for pname, wname, shard in _ATTENTION_PARAMS_MAPPING: + if f".attn.{wname}." in name: + mapped_name = name.replace(f".{wname}.", f".{pname}.") + if not _load(mapped_name, weight, shard): + raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") + matched = True + break + if matched: + continue + + # Dense MLP fused gate/up + down. + if ".mlp.w13_dn.weight" in name: + loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight) + elif ".mlp.w2_md.weight" in name: + loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight) + else: + if name.endswith(".bias") and name not in params: + continue + loaded_weight = _load(name, weight) + if not loaded_weight: + raise ValueError(f"Unexpected Inkling MTP weight: {original_name}") + required = { + name + for name in params + if name.startswith("model.layers.") or name.startswith("model.chain_norm.") + } + if missing := sorted(required - loaded): + raise ValueError( + "Inkling MTP checkpoint is missing required parameters: " + + ", ".join(missing) + ) + return loaded + + +EntryClass = [InklingMTP] diff --git a/vllm/models/inkling/nvidia/ops/__init__.py b/vllm/models/inkling/nvidia/ops/__init__.py new file mode 100644 index 000000000000..2d364fc6149a --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/__init__.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling kernels (NVIDIA). + +``rmsnorm`` / ``sconv`` import eagerly. The SwiGLU kernels and the FA4 +relative-attention wrapper are exposed lazily to keep this package's +import path lightweight. +""" + +from typing import TYPE_CHECKING + +from .norm import add_rmsnorm, rmsnorm +from .sconv import fused_sconv + +_LAZY_EXPORTS = { + "silu_and_mul_triton": "silu_and_mul", + "sink_silu_mul_epilogue": "silu_and_mul", + "inkling_fa4_rel_attention": "fa4_rel_attention", +} + +if TYPE_CHECKING: + from .fa4_rel_attention import inkling_fa4_rel_attention # noqa: F401 + from .silu_and_mul import ( # noqa: F401 + silu_and_mul_triton, + sink_silu_mul_epilogue, + ) + +__all__ = [ + "add_rmsnorm", + "rmsnorm", + "fused_sconv", + *sorted(_LAZY_EXPORTS), +] + + +def __getattr__(name: str): + module = _LAZY_EXPORTS.get(name) + if module is not None: + import importlib + + mod = importlib.import_module(f".{module}", __name__) + return getattr(mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py new file mode 100644 index 000000000000..f547f6506291 --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/fa4_rel_attention.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from collections.abc import Callable +from functools import cache +from typing import Any + +import torch + +from vllm.platforms import current_platform + + +def bucket_max_seqlen_q(max_seqlen_q: int) -> int: + """Round the FA4 scheduling bound up to a power of two.""" + return 1 << max(0, max_seqlen_q - 1).bit_length() + + +@cache +def _use_sheared_bias() -> bool: + capability = current_platform.get_device_capability() + return capability is not None and capability.major in (10, 11) + + +@cache +def _get_score_mod(rel_extent: int) -> Callable: + """Return the score modification that adds Inkling relative bias.""" + import cutlass.cute as cute + from cutlass.cute import Float32 + + from vllm.vllm_flash_attn.cute.seqlen_info import SeqlenInfoQK + + @cute.jit + def score_mod_rel_bias( + scores: cute.TensorSSA, + b_idx: cute.TensorSSA, + h_idx: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info: SeqlenInfoQK, + aux_tensors: list[cute.Tensor], + ) -> cute.TensorSSA: + rel_logits = aux_tensors[0] + + seqlen_local_offset = seqlen_info.seqlen_k - seqlen_info.seqlen_q + rel_dist = (q_idx + seqlen_local_offset) - kv_idx + global_q_idx = seqlen_info.offset_q + q_idx + + rel_dist_0 = rel_dist[0] + rel_idx = rel_dist_0 if rel_dist_0 >= 0 else 0 + rel_idx = rel_idx if rel_idx < rel_extent else (rel_extent - 1) + + rel_bias = rel_logits[global_q_idx[0], h_idx[0], rel_idx] + rel_bias = Float32(rel_bias) if rel_dist_0 == rel_idx else Float32(0.0) + return scores + rel_bias + + return score_mod_rel_bias + + +def inkling_fa4_num_splits( + *, + is_local: bool, + batch_size: int, + max_query_len: int, + num_heads: int, + num_kv_heads: int, + max_kv_len: int, +) -> int: + """Return the split-KV cap for Inkling relative attention.""" + capability = current_platform.get_device_capability() + if capability is not None and capability.major == 9: + return 1 + if is_local: + return 1 + + q_rows = max_query_len * (num_heads // num_kv_heads) + q_tiles = (q_rows + 255) // 256 + base_ctas = batch_size * num_kv_heads * q_tiles + # Shearing makes split/combine overhead more visible. Multi-tile causal + # prefill saturates around 64 CTAs. Batch-1 decode at very long context is + # memory-bound and uses a TP-specific cap measured through 1M KV tokens. + target_ctas = ( + 256 if q_tiles == 1 and batch_size == 1 else (128 if q_tiles == 1 else 64) + ) + max_splits = 128 + if q_tiles == 1 and batch_size == 1: + if num_kv_heads == 8: + max_splits = 16 + elif num_kv_heads == 4 or max_kv_len <= 8192: + max_splits = 32 + elif max_kv_len <= 65536: + max_splits = 64 + else: + max_splits = 128 + return max( + 1, + min(target_ctas // base_ctas, max_splits, (max_kv_len + 127) // 128), + ) + + +def inkling_fa4_rel_attention( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, + softmax_scale: float, + causal: bool, + window_size: tuple[int, int], + rel_extent: int, + rel_logits: torch.Tensor, + num_splits: int = 32, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Paged varlen FA4 over the bound K/V cache with the Inkling relative bias. + + ``q`` is ``(num_tokens, num_heads, head_dim)``; ``key_cache`` / ``value_cache`` + are the paged caches ``(num_blocks, block_size, num_kv_heads, head_dim)``; + ``block_table`` is the per-request page table and ``cache_seqlens`` the + per-request KV lengths (``seqused_k``). ``rel_logits`` is + ``(num_tokens, num_heads, rel_extent)``. + + Hopper uses standard FA4's score-mod gather. Blackwell uses tml-fa4's + sheared relative-bias layout. + """ + # cute uses (None, None) to mean "no window". + cute_window = (None, None) if window_size == (-1, -1) else window_size + + rel_logits = rel_logits.contiguous() + if _use_sheared_bias(): + from vllm.third_party.tml_fa4 import flash_attn_varlen_func + + bias_kwargs: dict[str, Any] = {"rel_bias": rel_logits} + else: + from vllm.vllm_flash_attn.cute import flash_attn_varlen_func + + bias_kwargs = { + "score_mod": _get_score_mod(rel_extent), + "aux_tensors": [rel_logits], + } + + ret = flash_attn_varlen_func( + q=q, + k=key_cache, + v=value_cache, + cu_seqlens_q=cu_seqlens_q, + seqused_k=cache_seqlens, + max_seqlen_q=max_seqlen_q, + page_table=block_table, + softmax_scale=softmax_scale, + causal=causal, + window_size=cute_window, + num_splits=num_splits, + return_lse=False, + out=out, + **bias_kwargs, + ) + if isinstance(ret, tuple): + return ret[0] + return ret diff --git a/vllm/models/inkling/nvidia/ops/fa4_warmup.py b/vllm/models/inkling/nvidia/ops/fa4_warmup.py new file mode 100644 index 000000000000..ff8cdb029362 --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/fa4_warmup.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Startup compilation of the FA4 kernels used by Inkling.""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +from functools import partial + +import torch + +from vllm.model_executor.warmup.cutedsl_warmup import ( + CuTeDSLCompileUnit, + register_cutedsl_warmup_provider, +) + +from .fa4_rel_attention import ( + bucket_max_seqlen_q, + inkling_fa4_num_splits, + inkling_fa4_rel_attention, +) + + +@dataclass(frozen=True) +class InklingFA4WarmupConfig: + num_heads: int + num_kv_heads: int + head_dim: int + rel_extent: int + window_size: tuple[int, int] + is_local: bool + max_kv_len: int + dtype: torch.dtype + kv_dtype: torch.dtype + block_size: int + max_num_reqs: int + max_num_batched_tokens: int + + +def _num_warps_bucket(num_reqs: int) -> int: + num_warps = min((num_reqs + 30) // 31, 32) + return 1 << (num_warps - 1).bit_length() + + +def _compile(config: InklingFA4WarmupConfig, max_seqlen_q: int, num_reqs: int) -> None: + from torch._subclasses.fake_tensor import FakeTensorMode + + num_splits = inkling_fa4_num_splits( + is_local=config.is_local, + batch_size=num_reqs, + max_query_len=max_seqlen_q, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + max_kv_len=config.max_kv_len, + ) + with FakeTensorMode(): + device = torch.accelerator.current_accelerator() + total_q = max_seqlen_q + num_reqs - 1 + q = torch.empty( + total_q, + config.num_heads, + config.head_dim, + dtype=config.dtype, + device=device, + ) + kv = torch.empty( + 1, + 2, + config.block_size, + config.num_kv_heads, + config.head_dim, + dtype=config.kv_dtype, + device=device, + ) + key_cache, value_cache = kv.unbind(1) + inkling_fa4_rel_attention( + q, + key_cache, + value_cache, + block_table=torch.empty(num_reqs, 1, dtype=torch.int32, device=device), + cache_seqlens=torch.empty(num_reqs, dtype=torch.int32, device=device), + cu_seqlens_q=torch.empty(num_reqs + 1, dtype=torch.int32, device=device), + max_seqlen_q=max_seqlen_q, + softmax_scale=1.0 / config.head_dim, + causal=True, + window_size=config.window_size, + rel_extent=config.rel_extent, + rel_logits=torch.empty( + total_q, + config.num_heads, + config.rel_extent, + dtype=config.dtype, + device=device, + ), + num_splits=num_splits, + out=torch.empty_like(q), + ) + + +def _iter_compile_units( + config: InklingFA4WarmupConfig, +) -> Iterator[CuTeDSLCompileUnit]: + max_bucket = bucket_max_seqlen_q(config.max_num_batched_tokens) + max_seqlen_q = 1 + while max_seqlen_q <= max_bucket: + min_query_len = max_seqlen_q // 2 + 1 + max_num_reqs = min( + config.max_num_reqs, + config.max_num_batched_tokens - min_query_len + 1, + ) + seen: set[tuple[int, int, int | None, bool]] = set() + for num_reqs in range(1, max_num_reqs + 1): + num_splits = inkling_fa4_num_splits( + is_local=config.is_local, + batch_size=num_reqs, + max_query_len=max_seqlen_q, + num_heads=config.num_heads, + num_kv_heads=config.num_kv_heads, + max_kv_len=config.max_kv_len, + ) + key = ( + max_seqlen_q, + num_splits, + _num_warps_bucket(num_reqs) if num_splits > 1 else None, + num_reqs > 1024, + ) + if key in seen: + continue + seen.add(key) + yield CuTeDSLCompileUnit( + name="inkling_fa4", + key=("inkling_fa4", config, key), + compile=partial(_compile, config, max_seqlen_q, num_reqs), + ) + max_seqlen_q *= 2 + + +class _WarmupProvider: + def __init__(self) -> None: + self.configs: set[InklingFA4WarmupConfig] = set() + + def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]: + return tuple( + unit for config in self.configs for unit in _iter_compile_units(config) + ) + + +_PROVIDER = _WarmupProvider() + + +def register_fa4_warmup(config: InklingFA4WarmupConfig) -> None: + _PROVIDER.configs.add(config) + register_cutedsl_warmup_provider(_PROVIDER) diff --git a/vllm/models/inkling/nvidia/ops/lamport.py b/vllm/models/inkling/nvidia/ops/lamport.py new file mode 100644 index 000000000000..c02290dc485b --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/lamport.py @@ -0,0 +1,766 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deadlock-free fused RS + short-conv + AG + residual + RMSNorm. + +The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``. + +Liveness +-------- +Large grids are deliberately split at the two communication dependencies: + + 1. ``_publish_input_kernel`` only publishes rank partials. + 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard. + 3. ``_sconv_publish_kernel`` only computes and publishes the local result. + 4. ``_gather_norm_kernel`` waits, gathers, and normalizes. + +Without PDL, CUDA stream order completes a producer before its consumer. With +PDL, every producer CTA posts all peer stores before triggering its dependent, +and the consumer executes ``gdc_wait`` before polling. A consumer therefore +waits only for stores from a producer that is already running (or complete) on +another GPU. It never waits for another block in its own grid. Consequently +spinning consumers cannot occupy resources needed by any producer, and the +wait-for graph has no cycle. The proof is independent of grid size, block +dispatch order, and occupancy. + +For one token, the first three phases use eight independent channel slices to +expose enough CTA parallelism for decode latency. Each slice has exclusive +ownership of its cache and Lamport columns; the gather/RMSNorm phase retains +one CTA per token so no cross-CTA reduction or completion counter is needed. + +Immediate buffer reuse (including replay of a captured CUDA graph) is also +safe. Rank R cannot republish input for call n+1 until its gather for n has +finished; that gather waited for owner O's output, which O publishes only +after consuming R's input for n. Likewise, R cannot republish output for n+1 +until its reduction for n+1 has observed destination D's input; D publishes +that input only after its gather consumed R's output for n. Thus every prior +read happens-before a same-slot rewrite. Three generations reduce incidental +coupling for ordinary launches, but correctness does not depend on rotation. + +The payload itself is the Lamport flag. A 32-bit store publishes two bf16s +atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real +negative zeroes are changed to positive zero before publication. Consumers +use volatile 32-bit loads and restore the sentinel after consuming a slot. +""" + +from __future__ import annotations + +import os + +import torch + +from vllm.distributed import get_tp_group +from vllm.distributed.parallel_state import in_the_same_node_as +from vllm.logger import init_logger +from vllm.triton_utils import HAS_TRITON, tl, triton + +logger = init_logger(__name__) + +_MAX_TOKENS = 16384 +_EMPTY_PAIR = tl.constexpr(0x80008000) if HAS_TRITON else 0x80008000 + + +@triton.jit +def _pack_bf16_pairs(values): + """Pack bf16 values into atomic u32 pairs and reserve negative zero.""" + lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2])) + lo = lo.to(tl.uint16, bitcast=True) + hi = hi.to(tl.uint16, bitcast=True) + lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32) + hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32) + return lo | (hi << 16) + + +@triton.jit +def _unpack_bf16_pairs(values): + lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + return tl.interleave(lo, hi) + + +@triton.jit +def _wait_pairs(ptr, offsets, mask): + values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) + while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0: + values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) + return values + + +@triton.jit +def _publish_input_kernel( + stage_ptr, + peer_ptrs, + peer_offset_u32, + stride_stage_t, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Publish this rank's full partial row into every shard owner's slots.""" + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + elem = tl.arange(0, CS_P2) + elem_mask = elem < CSS + ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) + # Address generation is independent of the preceding stream kernel. The + # acquire remains before stage/peer loads, which may consume its writes. + if USE_PDL: + tl.extra.cuda.gdc_wait() + + for owner in tl.static_range(WORLD): + values = tl.load( + stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem, + mask=elem_mask, + other=0.0, + ) + packed = _pack_bf16_pairs(values) + base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32)) + dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair + tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _reduce_insert_kernel( + input_peer_ptrs, + input_peer_offset_u32, + cache_ptr, + slot_ptr, + stride_cache_block, + stride_cache_head, + stride_cache_token, + stride_cache_dim, + block_size, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + CACHE_OFFSET: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Consume all partials for this rank and insert the reduced cache row.""" + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + source = tl.arange(0, WORLD) + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + offsets = ( + (token * WORLD + source)[:, None] * (CS // 2) + + split * (CSS // 2) + + pair[None, :] + ) + mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :] + input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) + input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32)) + input_u32 += input_peer_offset_u32 + # Slot metadata and the cache destination do not depend on input publish. + slot = tl.load(slot_ptr + token) + valid = slot >= 0 + channel = tl.arange(0, CS_P2) + channel_mask = channel < CSS + global_channel = split * CSS + channel + head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) + dim = CACHE_OFFSET + global_channel % HEAD_SIZE + safe_slot = tl.maximum(slot, 0).to(tl.int64) + dst = ( + cache_ptr + + (safe_slot // block_size) * stride_cache_block + + head * stride_cache_head + + (safe_slot % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + if USE_PDL: + tl.extra.cuda.gdc_wait() + packed = _wait_pairs(input_u32, offsets, mask) + + lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + reduced = tl.interleave( + tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16), + tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16), + ) + tl.store( + input_u32 + offsets, + tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32), + mask=mask, + ) + + tl.store(dst, reduced, mask=valid & channel_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _sconv_publish_kernel( + peer_ptrs, + peer_offset_u32, + residual_ptr, + weight_ptr, + cache_ptr, + position_ptr, + sequence_ptr, + slot_ptr, + block_table_ptr, + stride_residual_t, + stride_cache_block, + stride_cache_head, + stride_cache_token, + stride_cache_dim, + stride_block_table_r, + max_blocks, + block_size, + C: tl.constexpr, + CS: tl.constexpr, + CS_P2: tl.constexpr, + SPLITS: tl.constexpr, + HEAD_SIZE: tl.constexpr, + CACHE_OFFSET: tl.constexpr, + RANK: tl.constexpr, + WORLD: tl.constexpr, + WINDOW: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Compute this rank's shard, then publish it to every rank.""" + tl.static_assert(WINDOW == 4) + token = tl.program_id(0).to(tl.int64) + split = tl.program_id(1).to(tl.int64) + CSS: tl.constexpr = CS // SPLITS + channel = tl.arange(0, CS_P2) + channel_mask = channel < CSS + global_channel = split * CSS + channel + head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) + dim = CACHE_OFFSET + global_channel % HEAD_SIZE + slot = tl.load(slot_ptr + token) + valid = slot >= 0 + position = tl.load(position_ptr + token) + sequence = tl.load(sequence_ptr + token) + safe_slot = tl.maximum(slot, 0).to(tl.int64) + own_ptr = ( + cache_ptr + + (safe_slot // block_size) * stride_cache_block + + head * stride_cache_head + + (safe_slot % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + # These loads were made visible before reduce/insert could signal us and + # are independent of its cache store. Hoisting them gives PDL useful work + # to overlap while retaining the acquire before every cache read. + residual = tl.load( + residual_ptr + token * stride_residual_t + RANK * CS + global_channel, + mask=channel_mask, + other=0.0, + ) + weight0 = tl.load( + weight_ptr + global_channel * WINDOW, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight1 = tl.load( + weight_ptr + global_channel * WINDOW + 1, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight2 = tl.load( + weight_ptr + global_channel * WINDOW + 2, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + weight3 = tl.load( + weight_ptr + global_channel * WINDOW + 3, + mask=channel_mask, + other=0.0, + ).to(tl.float32) + ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) + if USE_PDL: + tl.extra.cuda.gdc_wait() + current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0) + + conv = tl.zeros([CS_P2], tl.float32) + for tap_idx in tl.static_range(WINDOW): + source_position = position - (WINDOW - 1) + tap_idx + take = valid & (source_position >= 0) + if tap_idx == WINDOW - 1: + value = tl.where(take, current.to(tl.float32), 0.0) + else: + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum(safe_position // block_size, max_blocks - 1) + physical_block = tl.load( + block_table_ptr + sequence * stride_block_table_r + logical_block, + mask=take, + other=0, + ).to(tl.int64) + source_ptr = ( + cache_ptr + + physical_block * stride_cache_block + + head * stride_cache_head + + (safe_position % block_size) * stride_cache_token + + dim * stride_cache_dim + ) + cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0) + value = tl.where(take, cached.to(tl.float32), 0.0) + if tap_idx == 0: + weight = weight0 + elif tap_idx == 1: + weight = weight1 + elif tap_idx == 2: + weight = weight2 + else: + weight = weight3 + conv += value * weight + + # Preserve both bf16 rounding points of the original sublayer. + short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16) + output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to( + tl.bfloat16 + ) + + pair = tl.arange(0, CS_P2 // 2) + pair_mask = pair < CSS // 2 + packed = _pack_bf16_pairs(output) + row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair + + for destination in tl.static_range(WORLD): + base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32)) + tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _gather_norm_kernel( + output_peer_ptrs, + output_peer_offset_u32, + norm_weight_ptr, + normed_ptr, + residual_out_ptr, + eps, + stride_output_t, + C: tl.constexpr, + C_P2: tl.constexpr, + RANK: tl.constexpr, + HAS_NORM: tl.constexpr, + USE_PDL: tl.constexpr, + launch_pdl: tl.constexpr, +): + """Consume a complete row; one CTA owns all outputs for one token.""" + token = tl.program_id(0).to(tl.int64) + pair = tl.arange(0, C_P2 // 2) + pair_mask = pair < C // 2 + output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) + output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32)) + output_u32 += output_peer_offset_u32 + offsets = token * (C // 2) + pair + channel = tl.arange(0, C_P2) + channel_mask = channel < C + # Gamma is independent of the preceding sconv publication. Keep the + # acquire immediately before polling the Lamport output slots. + if HAS_NORM: + weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0) + if USE_PDL: + tl.extra.cuda.gdc_wait() + packed = _wait_pairs(output_u32, offsets, pair_mask) + row = _unpack_bf16_pairs(packed) + tl.store( + residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask + ) + if HAS_NORM: + row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0) + inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps) + tl.store( + normed_ptr + token * stride_output_t + channel, + (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16), + mask=channel_mask, + ) + tl.store( + output_u32 + offsets, + tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32), + mask=pair_mask, + ) + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _validate_lamport_init_kernel( + input_peer_ptrs, + output_peer_ptrs, + bad_ptr, + num_pairs, + RANK: tl.constexpr, + BLOCK: tl.constexpr, +): + """Validate both complete local allocations through their fabric pointers.""" + offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < num_pairs + ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) + ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) + local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32)) + local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32)) + value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR) + value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR) + bad = tl.max( + tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0) + ) + if bad != 0: + tl.atomic_max(bad_ptr, 1) + + +class LamportRSConv: + """Persistent symmetric buffers for one TP group.""" + + def _initialize_lamport_buffers(self) -> None: + """Arm every slot and collectively verify the exact sentinel bits. + + Do not replace this with a zero-fill: Lamport readiness distinguishes + the bf16 bit pattern for -0.0 (0x8000) from every published payload. + The validation is deliberately collective so one rank cannot enter the + first polling kernel while another rank still has an unarmed buffer. + """ + if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous(): + raise RuntimeError("Lamport symmetric buffers must be contiguous") + if self.buf_in.numel() % 2 or self.buf_out.numel() % 2: + raise RuntimeError("Lamport buffers must contain whole uint32 pairs") + + # Fill through int16 so each bf16 lane receives the exact -0.0 bits. + # This covers all three generations and all max-token slots, including + # slots that a smaller first invocation does not touch. + self.buf_in.view(torch.int16).fill_(-0x8000) + self.buf_out.view(torch.int16).fill_(-0x8000) + torch.accelerator.synchronize(self.device) + self.tp.barrier() + + # Scan the complete local allocations. Since every rank performs the + # scan and participates in MAX, success proves every symmetric backing + # allocation was armed before any rank is allowed to use generation 0. + bad = torch.logical_or( + self.buf_in.view(torch.int16).ne(-0x8000).any(), + self.buf_out.view(torch.int16).ne(-0x8000).any(), + ).to(dtype=torch.float32) + bad = self.tp.all_reduce(bad) + if int(bad.item()) != 0: + raise RuntimeError("Lamport sentinel initialization failed on a TP rank") + self.tp.barrier() + + def _initialize_mnnvl_buffers(self) -> None: + """Initialize and validate FlashInfer fabric-mapped Lamport storage.""" + self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16) + self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16) + torch.accelerator.synchronize(self.device) + self.tp.barrier() + + bad = torch.zeros((), dtype=torch.int32, device=self.device) + num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2 + _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)]( + self.input_peer_ptrs, + self.output_peer_ptrs, + bad, + num_pairs, + RANK=self.rank, + BLOCK=256, + num_warps=4, + ) + bad = self.tp.all_reduce(bad.to(torch.float32)) + if int(bad.item()) != 0: + raise RuntimeError("MNNVL Lamport sentinel initialization failed") + self.tp.barrier() + + def __init__( + self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS + ) -> None: + import torch.distributed._symmetric_memory as symm_mem + + tp = get_tp_group() + self.tp = tp + self.group = tp.device_group + self.world_size = tp.world_size + self.rank = tp.rank_in_group + self.device = torch.device(tp.device) + if self.world_size not in (2, 4, 8): + raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}") + if hidden_size % (2 * self.world_size) != 0: + raise ValueError("hidden size must produce an even shard on every rank") + if window_size != 4: + raise ValueError(f"short-conv window size must be 4, got {window_size}") + if max_tokens < 1 or max_tokens > _MAX_TOKENS: + raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]") + + is_cross_node = not all(in_the_same_node_as(tp.cpu_group)) + self.hidden_size = hidden_size + self.window_size = window_size + self.max_tokens = max_tokens + self.shard_size = hidden_size // self.world_size + # Three generations follow FlashInfer's Lamport layout. A generation + # is reused only after two intervening collective calls. + self.num_buffers = 3 + self.input_generation_bytes = max_tokens * hidden_size * 2 + self.output_generation_bytes = max_tokens * hidden_size * 2 + self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9 + if is_cross_node: + try: + from flashinfer.comm.mnnvl import ( + McastGPUBuffer, + TorchDistBackend, + is_mnnvl_fabric_supported, + ) + except ImportError as error: + raise RuntimeError( + "cross-node TP requires FlashInfer MNNVL support" + ) from error + + local_supported = int( + is_mnnvl_fabric_supported(torch.accelerator.current_device_index()) + ) + unsupported = torch.tensor( + 1 - local_supported, dtype=torch.float32, device=self.device + ) + unsupported = tp.all_reduce(unsupported) + if int(unsupported.item()) != 0: + raise RuntimeError("cross-node TP is supported only on MNNVL fabric") + + comm_backend = TorchDistBackend(self.group) + allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2 + self._mnnvl_input_handle = McastGPUBuffer( + allocation_bytes, + self.world_size, + self.rank, + self.device, + comm_backend, + ) + self._mnnvl_output_handle = McastGPUBuffer( + allocation_bytes, + self.world_size, + self.rank, + self.device, + comm_backend, + ) + self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev() + self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev() + self._initialize_mnnvl_buffers() + logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers") + else: + self.buf_in = symm_mem.empty( + self.num_buffers, + max_tokens, + self.world_size, + self.shard_size, + dtype=torch.bfloat16, + device=self.device, + ) + self.buf_out = symm_mem.empty( + self.num_buffers, + max_tokens, + hidden_size, + dtype=torch.bfloat16, + device=self.device, + ) + group_name = self.group.group_name + input_handle = symm_mem.rendezvous(self.buf_in, group_name) + output_handle = symm_mem.rendezvous(self.buf_out, group_name) + self.input_peer_ptrs = input_handle.buffer_ptrs_dev + self.output_peer_ptrs = output_handle.buffer_ptrs_dev + self._input_handle = input_handle + self._output_handle = output_handle + self._initialize_lamport_buffers() + self.generation = 0 + + def usable(self, num_tokens: int) -> bool: + return 0 < num_tokens <= self.max_tokens + + def rs_sconv_ag_add_norm( + self, + input_tensor: torch.Tensor, + residual: torch.Tensor, + conv_weight: torch.Tensor, + norm_weight: torch.Tensor | None, + eps: float, + cache: torch.Tensor, + positions: torch.Tensor, + block_table: torch.Tensor, + seq_idx: torch.Tensor, + slot_mapping: torch.Tensor, + off_s: int, + ws: int, + block_size: int, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``.""" + tokens, hidden_size = residual.shape + if not self.usable(tokens): + raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]") + if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16: + raise ValueError("residual must be bf16 [T, 6144]") + if ( + input_tensor.shape != residual.shape + or input_tensor.dtype != torch.bfloat16 + or input_tensor.stride(1) != 1 + ): + raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]") + shard_size = hidden_size // self.world_size + if conv_weight.shape != (shard_size, self.window_size): + raise ValueError( + f"conv_weight must have shape [{shard_size}, {self.window_size}]" + ) + if ( + conv_weight.dtype != torch.bfloat16 + or conv_weight.stride(0) != self.window_size + ): + raise ValueError("conv_weight must be contiguous bf16") + if norm_weight is not None and ( + norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16 + ): + raise ValueError("norm_weight must be bf16 [6144] or None") + if cache.dtype != torch.bfloat16 or cache.ndim != 4: + raise ValueError("cache must be a 4-D bf16 tensor") + if shard_size % ws != 0 or cache.shape[1] != shard_size // ws: + raise ValueError("cache head layout is inconsistent with ws") + if off_s < 0 or off_s + ws > cache.shape[3]: + raise ValueError("cache channel offset is out of bounds") + + index = self.generation + input_offset = index * self.input_generation_bytes // 4 + output_offset = index * self.output_generation_bytes // 4 + normed = torch.empty_like(residual) if norm_weight is not None else None + residual_out = torch.empty_like(residual) + phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1 + phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits) + phase_grid = (tokens, phase_splits) + # Wide CTAs win before the grid saturates; smaller CTAs reduce register + # pressure once high-throughput batches provide enough parallelism. + if 128 <= tokens <= 2048: + phase_warps = 16 + elif tokens > 2048: + phase_warps = 8 + else: + phase_warps = 4 + gather_warps = 4 if tokens >= 256 else 8 + _publish_input_kernel[phase_grid]( + input_tensor, + self.input_peer_ptrs, + input_offset, + input_tensor.stride(0), + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + RANK=self.rank, + WORLD=self.world_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _reduce_insert_kernel[phase_grid]( + self.input_peer_ptrs, + input_offset, + cache, + slot_mapping, + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + block_size, + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + HEAD_SIZE=ws, + CACHE_OFFSET=off_s, + RANK=self.rank, + WORLD=self.world_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _sconv_publish_kernel[phase_grid]( + self.output_peer_ptrs, + output_offset, + residual, + conv_weight, + cache, + positions, + seq_idx, + slot_mapping, + block_table, + residual.stride(0), + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + block_table.stride(0), + block_table.shape[1], + block_size, + C=hidden_size, + CS=shard_size, + CS_P2=phase_tile_p2, + SPLITS=phase_splits, + HEAD_SIZE=ws, + CACHE_OFFSET=off_s, + RANK=self.rank, + WORLD=self.world_size, + WINDOW=self.window_size, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=phase_warps, + ) + _gather_norm_kernel[(tokens,)]( + self.output_peer_ptrs, + output_offset, + norm_weight if norm_weight is not None else residual, + normed if normed is not None else residual_out, + residual_out, + eps, + residual.stride(0), + C=hidden_size, + C_P2=triton.next_power_of_2(hidden_size), + RANK=self.rank, + HAS_NORM=norm_weight is not None, + USE_PDL=self.use_pdl, + launch_pdl=self.use_pdl, + num_warps=gather_warps, + ) + self.generation = (index + 1) % self.num_buffers + return normed, residual_out + + +_STATE: LamportRSConv | None = None +_STATE_FAILED = False + + +def initialize_lamport_rs_conv( + hidden_size: int, window_size: int, max_num_batched_tokens: int +) -> None: + """Collectively initialize the TP-group state during model construction.""" + global _STATE, _STATE_FAILED + if _STATE is not None: + if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size: + raise RuntimeError("all Lamport users must share hidden and window sizes") + return + if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0": + return + try: + max_tokens = min(_MAX_TOKENS, max_num_batched_tokens) + _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens) + except Exception: + _STATE_FAILED = True + logger.exception("fused collective unavailable; use the NCCL fallback") + + +def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None: + """Return the state initialized with the model, or ``None`` for fallback.""" + if _STATE is not None and ( + _STATE.hidden_size != hidden_size or _STATE.window_size != window_size + ): + raise RuntimeError("all Lamport users must share hidden and window sizes") + return _STATE diff --git a/vllm/models/inkling/nvidia/ops/mm_towers.py b/vllm/models/inkling/nvidia/ops/mm_towers.py new file mode 100644 index 000000000000..021cf5274bd2 --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/mm_towers.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused CUDA kernels for the Inkling vision/audio towers. + +Both kernels keep the reference paths' fp32 accumulation and per-op bf16 +rounding points (native ``rms_norm`` / ``F.gelu``); outputs are frequently +bit-identical and otherwise differ by 1-2 bf16 ulps from reduction-order +(real-checkpoint-weight cosine vs reference > 0.9999998). +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, tldevice, triton + +from .norm import _get_num_warps_from_block_size + + +@triton.jit +def _dmel_embed_sum_norm_kernel( + idx_ptr, # [T, NB] int32 dMel bin indices (values in [0, VOCAB)) + w_ptr, # [NB * VOCAB, D] bf16 embedding table + norm_w_ptr, # [D] (unused if HAS_NORM=False) + out_ptr, # [T, D] bf16 + eps, + D, + stride_idx_t, + NB: tl.constexpr, + VOCAB: tl.constexpr, + D_P2: tl.constexpr, + HAS_NORM: tl.constexpr, +): + t = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, D_P2) + mask = offs < D + # One embedding row per mel bin (bin b uses table rows [b*VOCAB, (b+1)*VOCAB)), + # summed in fp32 (matches torch's fp32-accumulated bf16 .sum()). + acc = tl.zeros([D_P2], dtype=tl.float32) + for b in tl.static_range(NB): + v = tl.load(idx_ptr + t * stride_idx_t + b) + row = (b * VOCAB + v).to(tl.int64) + acc += tl.load(w_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) + h = acc.to(tl.bfloat16) + if HAS_NORM: + # Match ir.ops.rms_norm: fp32 variance/normalize, then a single-rounded + # bf16 multiply with the bf16 weight. + x32 = h.to(tl.float32) + var = tl.sum(tl.where(mask, x32 * x32, 0.0), axis=0) / D + xn = (x32 * tl.math.rsqrt(var + eps)).to(tl.bfloat16) + w = tl.load(norm_w_ptr + offs, mask=mask, other=0.0) + h = xn * w + tl.store(out_ptr + t * D + offs, h, mask=mask) + + +def dmel_embed_sum_norm( + dmel_idx: torch.Tensor, # [T, NB] int32 + weight: torch.Tensor, # [NB * VOCAB, D] bf16 + norm_weight: torch.Tensor | None, + eps: float, +) -> torch.Tensor: + """``rmsnorm(sum_b weight[b * VOCAB + idx[:, b]])`` in one launch (no + [T, NB, D] intermediate).""" + T, nb = dmel_idx.shape + D = weight.shape[1] + assert weight.shape[0] % nb == 0 + vocab = weight.shape[0] // nb + out = torch.empty((T, D), dtype=weight.dtype, device=weight.device) + if T == 0: + return out + d_p2 = triton.next_power_of_2(D) + _dmel_embed_sum_norm_kernel[(T,)]( + dmel_idx, + weight, + norm_weight if norm_weight is not None else weight, + out, + eps, + D, + dmel_idx.stride(0), + NB=nb, + VOCAB=vocab, + D_P2=d_p2, + HAS_NORM=norm_weight is not None, + # Swept on GB200: 8 warps beats the block-size heuristic's 16 by ~4% + # at large T (the 80-row gather chain is latency- not lane-bound). + num_warps=8, + ) + return out + + +@triton.jit +def _rmsnorm_gelu_kernel( + x_ptr, # [R, D] bf16 + w_ptr, # [D] + out_ptr, # [R, D] bf16 (or the folded layout when FOLD) + eps, + R, + D, + D_P2: tl.constexpr, + BLOCK_M: tl.constexpr, # rows per block (>1 for small D) + HAS_GELU: tl.constexpr, + FOLD: tl.constexpr, + # fold geometry: input rows index [N, T, H, W]; the store scatters each + # row to (out_row, slot) of fold_timespace_to_depth's output layout. + FT: tl.constexpr, + FH: tl.constexpr, + FW: tl.constexpr, + TF: tl.constexpr, + HF: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) + rmask = rows < R + offs = tl.arange(0, D_P2) + mask = rmask[:, None] & (offs < D)[None, :] + x32 = tl.load(x_ptr + rows[:, None] * D + offs[None, :], mask=mask, other=0.0).to( + tl.float32 + ) + var = tl.sum(x32 * x32, axis=1) / D + xn = (x32 * tl.math.rsqrt(var + eps)[:, None]).to(tl.bfloat16) + w = tl.load(w_ptr + offs, mask=offs < D, other=0.0) + h = xn * w[None, :] # bf16 multiply, matching ir.ops.rms_norm + if HAS_GELU: + # Exact (erf) GELU on the bf16-rounded norm output, fp32 math, + # matching F.gelu's opmath on a bf16 tensor. + g32 = h.to(tl.float32) + h = (0.5 * g32 * (1.0 + tldevice.erf(g32 * 0.7071067811865476))).to(tl.bfloat16) + if FOLD: + # Store directly into the next layer's folded layout (a pure + # permutation — replaces the separate fold copy pass). + t = (rows // (FH * FW)) % FT + hh = (rows // FW) % FH + ww = rows % FW + n = rows // (FT * FH * FW) + slot = ((t % TF) * HF + hh % HF) * HF + ww % HF + out_row = ((n * (FT // TF) + t // TF) * (FH // HF) + hh // HF) * ( + FW // HF + ) + ww // HF + base = (out_row * (TF * HF * HF) + slot) * D + tl.store(out_ptr + base[:, None] + offs[None, :], h, mask=mask) + else: + tl.store(out_ptr + rows[:, None] * D + offs[None, :], h, mask=mask) + + +def rmsnorm_gelu( + x: torch.Tensor, # [..., D] bf16 contiguous + weight: torch.Tensor, + eps: float, + gelu: bool = True, + fold: tuple[int, int] | None = None, # (t_fold, hw_fold) of the NEXT fold +) -> torch.Tensor: + """Fused ``gelu(rmsnorm(x))`` (or plain rmsnorm); multiple rows per block + when D is small. With ``fold``, x must be [N, T, H, W, D] and the output + comes back as ``fold_timespace_to_depth(result, *fold)``.""" + D = x.shape[-1] + flat = x.reshape(-1, D) + assert flat.stride(1) == 1 and flat.stride(0) == D + R = flat.shape[0] + if fold is None: + out = torch.empty_like(flat) + ft = fh = fw = tf = hf = 1 + out_shape = x.shape + else: + tf, hf = fold + N, ft, fh, fw, _ = x.shape + out_shape = (N, ft // tf, fh // hf, fw // hf, tf * hf * hf * D) + out = torch.empty(out_shape, dtype=x.dtype, device=x.device) + if R == 0: + return out.reshape(out_shape) + d_p2 = triton.next_power_of_2(D) + block_m = max(1, 4096 // d_p2) + _rmsnorm_gelu_kernel[(triton.cdiv(R, block_m),)]( + flat, + weight, + out, + eps, + R, + D, + D_P2=d_p2, + BLOCK_M=block_m, + HAS_GELU=gelu, + FOLD=fold is not None, + FT=ft, + FH=fh, + FW=fw, + TF=tf, + HF=hf, + num_warps=_get_num_warps_from_block_size(d_p2 * block_m), + ) + return out.reshape(out_shape) diff --git a/vllm/models/inkling/nvidia/ops/norm.py b/vllm/models/inkling/nvidia/ops/norm.py new file mode 100644 index 000000000000..d34c1c3259b9 --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/norm.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from functools import lru_cache + +import torch + +from vllm.triton_utils import tl, triton + +_MAX_FUSED_SIZE = 65536 + + +def _get_num_warps_from_block_size(block_size: int) -> int: + if block_size >= 32768: + return 32 + if block_size >= 8192: + return 16 + if block_size >= 2048: + return 8 + return 4 + + +def _largest_power_of_2(n: int) -> int: + assert n > 0, f"{n=}" + return 1 << (n.bit_length() - 1) + + +@lru_cache(maxsize=128) +def _get_grid_size_for_mem_bw_kernel(device: torch.device, factor: int = 8) -> int: + num_sms = torch.cuda.get_device_properties(device).multi_processor_count + return _largest_power_of_2(num_sms) * factor + + +@triton.jit +def _rmsnorm_fwd_kernel( + x_ptr, + weight_ptr, + y_ptr, + rstd_ptr, + eps, + x_stride_0, + y_stride_0, + n_cols, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + x = tl.load(x_ptr + pid_m * x_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + row_var = tl.sum(x * x, axis=0) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + tl.store(rstd_ptr + pid_m, rstd) + + y = x * rstd * weight + tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n) + + +@triton.jit(do_not_specialize=["n_rows"]) +def _rmsnorm_fwd_kernel_block_m( + x_ptr, + weight_ptr, + y_ptr, + rstd_ptr, + eps, + x_stride_0, + y_stride_0, + n_rows, + n_cols, + block_size_m: tl.constexpr, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + num_blocks_m = tl.cdiv(n_rows, block_size_m) + blocks_per_pid = tl.cdiv(num_blocks_m, tl.num_programs(0)) + block_id_start = pid_m * blocks_per_pid + block_id_end = min(block_id_start + blocks_per_pid, num_blocks_m) + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + + for block_id in range(block_id_start, block_id_end): + offs_m = block_id * block_size_m + tl.arange(0, block_size_m) + mask_m = offs_m < n_rows + mask_mn = mask_m[:, None] & mask_n[None, :] + x = tl.load( + x_ptr + offs_m[:, None] * x_stride_0 + offs_n[None, :], + mask=mask_mn, + other=0.0, + ).to(tl.float32) + row_var = tl.sum(x * x, axis=1) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + tl.store(rstd_ptr + offs_m, rstd, mask=mask_m) + + y = x * rstd[:, None] * weight + tl.store( + y_ptr + offs_m[:, None] * y_stride_0 + offs_n[None, :], + y, + mask=mask_mn, + ) + + +@triton.jit +def _add_rmsnorm_fwd_kernel( + res_ptr, # [T, N] residual (read) + delta_ptr, # [T, N] delta to add (read) + weight_ptr, + y_ptr, # [T, N] normed output + res_out_ptr, # [T, N] updated residual output + eps, + res_stride_0, + delta_stride_0, + y_stride_0, + res_out_stride_0, + n_cols, + block_size_n: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + + weight = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + r = tl.load(res_ptr + pid_m * res_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + d = tl.load(delta_ptr + pid_m * delta_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + # Round the sum to the residual dtype first (matches the eager + # `residual + delta` then rmsnorm-on-bf16 sequence bit-for-bit). + s = (r + d).to(res_out_ptr.dtype.element_ty) + tl.store(res_out_ptr + pid_m * res_out_stride_0 + offs_n, s, mask=mask_n) + x = s.to(tl.float32) + row_var = tl.sum(x * x, axis=0) / n_cols + rstd = tl.math.rsqrt(row_var + eps) + y = x * rstd * weight + tl.store(y_ptr + pid_m * y_stride_0 + offs_n, y, mask=mask_n) + + +def add_rmsnorm( + residual: torch.Tensor, + delta: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused ``res = residual + delta; y = rmsnorm(res)``. + + Returns ``(y, res)``; both are fresh tensors (cudagraph-friendly, no + in-place update of the inputs). + """ + assert residual.ndim == 2 and delta.ndim == 2, (residual.shape, delta.shape) + n_rows, n_cols = residual.shape + assert weight.shape[0] == n_cols + y = torch.empty_like(residual) + res_out = torch.empty_like(residual) + if n_rows == 0: + return y, res_out + + block_size_n = triton.next_power_of_2(n_cols) + max_block_size_n = _MAX_FUSED_SIZE // residual.element_size() + if max_block_size_n < block_size_n: + raise RuntimeError(f"Large {n_cols=} is not supported") + num_warps = _get_num_warps_from_block_size(block_size_n) + _add_rmsnorm_fwd_kernel[(n_rows,)]( + residual, + delta, + weight, + y, + res_out, + eps, + residual.stride(0), + delta.stride(0), + y.stride(0), + res_out.stride(0), + n_cols, + block_size_n, + num_warps=num_warps, + ) + return y, res_out + + +@triton.jit +def _embed_rmsnorm_kernel( + ids_ptr, # [T] token ids + table_ptr, # [V, N] embedding table + weight_ptr, # [N] (HAS_NORM only) + chain_weight_ptr, # [N] (HAS_CHAIN only) + out_ptr, # [T, N] rmsnorm(table[ids], weight) + chain_out_ptr, # [T, N] rmsnorm(out, chain_weight) (HAS_CHAIN only) + eps, + table_stride_0, + n_cols, + block_size_n: tl.constexpr, + HAS_NORM: tl.constexpr, + HAS_CHAIN: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + row = tl.load(ids_ptr + pid_m).to(tl.int64) + x = tl.load(table_ptr + row * table_stride_0 + offs_n, mask=mask_n, other=0.0) + if HAS_NORM: + xf = x.to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps) + w = tl.load(weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + # Round to the output dtype so the chained norm is bit-exact vs the + # unfused pair (which stores bf16 in between). + x = (xf * rstd * w).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + pid_m * n_cols + offs_n, x, mask=mask_n) + if HAS_CHAIN: + xf = x.to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(xf * xf, axis=0) / n_cols + eps) + w = tl.load(chain_weight_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + tl.store( + chain_out_ptr + pid_m * n_cols + offs_n, + (xf * rstd * w).to(chain_out_ptr.dtype.element_ty), + mask=mask_n, + ) + + +def embed_rmsnorm( + input_ids: torch.Tensor, + embed_table: torch.Tensor, + weight: torch.Tensor | None, + eps: float, + chain_weight: torch.Tensor | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Fused ``rmsnorm(embed_table[input_ids], weight)`` row gather + norm. + + Requires the full vocab on-rank (replicated or tp_size == 1). + ``weight=None`` skips the norm (``use_embed_norm=False``), leaving a pure + embedding-table gather. + ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the + first decoder layer's pre-attention norm) as a second output, still one + launch. Bit-exact vs the unfused module sequence.""" + ids = input_ids.view(-1) + (T,) = ids.shape + n = embed_table.shape[1] + out = torch.empty( + (*input_ids.shape, n), dtype=embed_table.dtype, device=embed_table.device + ) + chain_out = torch.empty_like(out) if chain_weight is not None else None + if T > 0: + block_size_n = triton.next_power_of_2(n) + _embed_rmsnorm_kernel[(T,)]( + ids, + embed_table, + weight if weight is not None else embed_table, + chain_weight if chain_weight is not None else embed_table, + out, + chain_out if chain_out is not None else out, + eps, + embed_table.stride(0), + n, + block_size_n, + HAS_NORM=weight is not None, + HAS_CHAIN=chain_weight is not None, + num_warps=_get_num_warps_from_block_size(block_size_n), + ) + if chain_out is not None: + return out, chain_out + return out + + +@triton.jit +def _embed_dual_rmsnorm_cat_kernel( + hidden_ptr, # [T, N] + emb_ptr, # [T, N] embeddings, or the [V, N] embedding table when GATHER + ids_ptr, # [T] token ids (GATHER only) + w_hidden_ptr, # [N] + w_pre_ptr, # [N] chained pre-norm on the embed side (HAS_PRE_NORM only) + w_embed_ptr, # [N] + out_ptr, # [T, 2N]: [rmsnorm(hidden) | rmsnorm(rmsnorm?(emb))] + eps, + hidden_stride_0, + emb_stride_0, + n_cols, + block_size_n: tl.constexpr, + GATHER: tl.constexpr, + HAS_PRE_NORM: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + which = tl.program_id(1) # 0 -> hidden into cols [0, N); 1 -> emb into [N, 2N) + offs_n = tl.arange(0, block_size_n) + mask_n = offs_n < n_cols + if which == 0: + x = tl.load( + hidden_ptr + pid_m * hidden_stride_0 + offs_n, mask=mask_n, other=0.0 + ).to(tl.float32) + w = tl.load(w_hidden_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + else: + row = tl.load(ids_ptr + pid_m).to(tl.int64) if GATHER else pid_m + x = tl.load(emb_ptr + row * emb_stride_0 + offs_n, mask=mask_n, other=0.0).to( + tl.float32 + ) + if HAS_PRE_NORM: + w_pre = tl.load(w_pre_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) + # Round-trip through the output dtype so the chained norm is + # bit-exact vs the unfused pair (which stores bf16 in between). + x = (x * rstd * w_pre).to(out_ptr.dtype.element_ty).to(tl.float32) + w = tl.load(w_embed_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + rstd = tl.math.rsqrt(tl.sum(x * x, axis=0) / n_cols + eps) + tl.store( + out_ptr + pid_m * (2 * n_cols) + which * n_cols + offs_n, + (x * rstd * w).to(out_ptr.dtype.element_ty), + mask=mask_n, + ) + + +def embed_dual_rmsnorm_cat( + hidden: torch.Tensor, + hidden_weight: torch.Tensor, + embed_weight: torch.Tensor, + eps: float, + *, + embeds: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + embed_table: torch.Tensor | None = None, + pre_norm_weight: torch.Tensor | None = None, +) -> torch.Tensor: + """The MTP depth-layer input in one launch: + ``cat([rmsnorm(hidden, w_h), rmsnorm(pre?(emb), w_e)], -1)``. + + The embed side is either a fused row gather ``embed_table[input_ids]`` + (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged + multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the + backbone embed_norm in front of the depth embed_norm (bit-exact vs the + unfused sequence). The concat copies collapse into direct writes.""" + T, n = hidden.shape + if embeds is not None: + assert embeds.shape == hidden.shape + src, ids, src_stride = embeds, embeds, embeds.stride(0) + gather = False + else: + assert input_ids is not None and embed_table is not None + assert input_ids.shape == (T,) and embed_table.shape[1] == n + src, ids, src_stride = embed_table, input_ids, embed_table.stride(0) + gather = True + out = torch.empty((T, 2 * n), dtype=hidden.dtype, device=hidden.device) + if T == 0: + return out + block_size_n = triton.next_power_of_2(n) + _embed_dual_rmsnorm_cat_kernel[(T, 2)]( + hidden, + src, + ids, + hidden_weight, + pre_norm_weight if pre_norm_weight is not None else embed_weight, + embed_weight, + out, + eps, + hidden.stride(0), + src_stride, + n, + block_size_n, + GATHER=gather, + HAS_PRE_NORM=pre_norm_weight is not None, + num_warps=_get_num_warps_from_block_size(block_size_n), + ) + return out + + +def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + assert x.ndim == 2, f"{x.shape=}" + assert weight.ndim == 1, f"{weight.shape=}" + n_rows, n_cols = x.shape + assert weight.shape[0] == n_cols, f"{weight.shape=} {x.shape=}" + y = torch.empty_like(x) + rstd = torch.empty((n_rows,), dtype=torch.float32, device=x.device) + + block_size_n = triton.next_power_of_2(n_cols) + max_block_size_n = _MAX_FUSED_SIZE // x.element_size() + if max_block_size_n < block_size_n: + raise RuntimeError(f"Large {n_cols=} is not supported") + block_size_m = max(1, 4096 // block_size_n) + num_warps = _get_num_warps_from_block_size(block_size_n) + + if block_size_m == 1: + _rmsnorm_fwd_kernel[(n_rows,)]( + x, + weight, + y, + rstd, + eps, + x.stride(0), + y.stride(0), + n_cols, + block_size_n, + num_warps=num_warps, + ) + else: + grid_size = _get_grid_size_for_mem_bw_kernel(x.device) + _rmsnorm_fwd_kernel_block_m[(grid_size,)]( + x, + weight, + y, + rstd, + eps, + x.stride(0), + y.stride(0), + n_rows, + n_cols, + block_size_m, + block_size_n, + num_warps=num_warps, + ) + return y diff --git a/vllm/models/inkling/nvidia/ops/qkvr_prep.py b/vllm/models/inkling/nvidia/ops/qkvr_prep.py new file mode 100644 index 000000000000..a910e72170d1 --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/qkvr_prep.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import aux_stream + +LOW_BLOCK_M = 32 +LOW_BLOCK_N = 64 +LOW_NUM_WARPS = 4 +THROUGHPUT_BLOCK_M = 32 +THROUGHPUT_BLOCK_N = 128 +THROUGHPUT_GROUP_M = 2 +THROUGHPUT_NUM_WARPS = 4 +SMALL_TOKEN_THRESHOLD = 128 +SMALL_NUM_WARPS = 2 +Q_BLOCK_ROWS = 8 +Q_NUM_WARPS = 2 +KV_BLOCK_ROWS = 4 +KV_NUM_WARPS = 2 + + +@triton.jit(do_not_specialize=["rows"]) +def _rel_proj_low_latency_kernel( + qkvr_ptr, + rel_proj_ptr, + rel_out_ptr, + log_scaling_ptr, + rows, + stride_x_t, + R_OFFSET: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + REL_EXTENT: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + row = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + inner = tl.arange(0, 16) + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + relative = tl.load( + qkvr_ptr + + token[:, None] * stride_x_t + + R_OFFSET + + head[:, None] * 16 + + inner[None, :], + mask=row[:, None] < rows, + other=0.0, + ) + projection = tl.load( + rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :], + mask=col[None, :] < REL_EXTENT, + other=0.0, + ) + values = tl.dot(relative, projection, out_dtype=tl.float32).to( + rel_out_ptr.dtype.element_ty + ) + values = values.to(tl.float32) + if APPLY_LOG_SCALING: + values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[:, None] + tl.store( + rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :], + values.to(rel_out_ptr.dtype.element_ty), + mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT), + ) + + +@triton.jit(do_not_specialize=["rows"]) +def _rel_proj_throughput_kernel( + qkvr_ptr, + rel_proj_ptr, + rel_out_ptr, + log_scaling_ptr, + rows, + stride_x_t, + R_OFFSET: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + REL_EXTENT: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + GROUP_M: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + row_group = tl.program_id(0) + col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + inner = tl.arange(0, 16) + projection = tl.load( + rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :], + mask=col[None, :] < REL_EXTENT, + other=0.0, + ) + row_offsets = tl.arange(0, BLOCK_M) + for group_offset in tl.static_range(GROUP_M): + row = (row_group * GROUP_M + group_offset) * BLOCK_M + row_offsets + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + relative = tl.load( + qkvr_ptr + + token[:, None] * stride_x_t + + R_OFFSET + + head[:, None] * 16 + + inner[None, :], + mask=row[:, None] < rows, + other=0.0, + ) + values = tl.dot(relative, projection, out_dtype=tl.float32).to( + rel_out_ptr.dtype.element_ty + ) + values = values.to(tl.float32) + if APPLY_LOG_SCALING: + values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[ + :, None + ] + tl.store( + rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :], + values.to(rel_out_ptr.dtype.element_ty), + mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT), + ) + + +def use_rel_proj_throughput(rows: int, rel_extent: int) -> bool: + min_rows = 8192 if rel_extent == 512 else 2048 + return rows >= min_rows + + +def qkvr_rel_proj( + qkvr: torch.Tensor, + rel_proj: torch.Tensor, + rel_out: torch.Tensor, + log_scaling: torch.Tensor | None, + *, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + d_rel: int, +) -> None: + rows = qkvr.shape[0] * num_q_heads + rel_extent = rel_proj.shape[1] + assert d_rel == 16 and rel_proj.shape[0] == 16 + r_offset = num_q_heads * head_dim + 2 * num_kv_heads * head_dim + log_scaling_ptr = log_scaling if log_scaling is not None else qkvr + common = dict( + R_OFFSET=r_offset, + NUM_Q_HEADS=num_q_heads, + REL_EXTENT=rel_extent, + APPLY_LOG_SCALING=log_scaling is not None, + ) + + if use_rel_proj_throughput(rows, rel_extent): + grid = ( + triton.cdiv(rows, THROUGHPUT_BLOCK_M * THROUGHPUT_GROUP_M), + triton.cdiv(rel_extent, THROUGHPUT_BLOCK_N), + ) + _rel_proj_throughput_kernel[grid]( + qkvr, + rel_proj, + rel_out, + log_scaling_ptr, + rows, + qkvr.stride(0), + BLOCK_M=THROUGHPUT_BLOCK_M, + BLOCK_N=THROUGHPUT_BLOCK_N, + GROUP_M=THROUGHPUT_GROUP_M, + num_warps=THROUGHPUT_NUM_WARPS, + **common, + ) + return + + grid = ( + triton.cdiv(rows, LOW_BLOCK_M), + triton.cdiv(rel_extent, LOW_BLOCK_N), + ) + _rel_proj_low_latency_kernel[grid]( + qkvr, + rel_proj, + rel_out, + log_scaling_ptr, + rows, + qkvr.stride(0), + BLOCK_M=LOW_BLOCK_M, + BLOCK_N=LOW_BLOCK_N, + num_warps=LOW_NUM_WARPS, + **common, + ) + + +@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"]) +def _qkvr_qkv_kernel( + qkvr_ptr, + q_norm_weight_ptr, + q_out_ptr, + rel_proj_ptr, + rel_out_ptr, + k_weight_ptr, + v_weight_ptr, + k_norm_weight_ptr, + conv_cache_ptr, + key_cache_ptr, + value_cache_ptr, + positions_ptr, + seq_idx_ptr, + conv_slot_mapping_ptr, + conv_block_table_ptr, + query_start_ptr, + attention_slot_mapping_ptr, + log_scaling_ptr, + tokens, + eps, + stride_x_t, + stride_cc_block, + stride_cc_head, + stride_cc_token, + stride_cc_dim, + stride_kc_block, + stride_kc_token, + stride_kc_head, + stride_vc_block, + stride_vc_token, + stride_vc_head, + stride_block_table_req, + max_blocks, + conv_block_size, + attention_page_size, + Q_WIDTH: tl.constexpr, + KV_WIDTH: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + OFF_K: tl.constexpr, + OFF_V: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, + D_REL: tl.constexpr, + REL_EXTENT: tl.constexpr, + REL_EXTENT_PADDED: tl.constexpr, +): + block = tl.program_id(0) + num_q_rows = tokens * NUM_Q_HEADS + dims = tl.arange(0, HEAD_DIM) + + if block < num_q_rows: + row = block + token = row // NUM_Q_HEADS + head = row % NUM_Q_HEADS + values = tl.load( + qkvr_ptr + token * stride_x_t + head * HEAD_DIM + dims, + ).to(tl.float32) + weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(values * values, axis=0) / HEAD_DIM + eps) + normalized = values * rstd * weight + if APPLY_LOG_SCALING: + normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32) + normalized *= tl.load(log_scaling_ptr + token) + tl.store( + q_out_ptr + row * HEAD_DIM + dims, + normalized.to(q_out_ptr.dtype.element_ty), + ) + rel_cols = tl.arange(0, REL_EXTENT_PADDED) + rel_mask = rel_cols < REL_EXTENT + projected = tl.zeros([REL_EXTENT_PADDED], dtype=tl.float32) + rel_offset = Q_WIDTH + 2 * KV_WIDTH + head * D_REL + for rel_dim in tl.static_range(D_REL): + rel_value = tl.load( + qkvr_ptr + token * stride_x_t + rel_offset + rel_dim + ).to(tl.float32) + proj = tl.load( + rel_proj_ptr + rel_dim * REL_EXTENT + rel_cols, + mask=rel_mask, + other=0.0, + ).to(tl.float32) + projected += rel_value * proj + projected = projected.to(rel_out_ptr.dtype.element_ty).to(tl.float32) + if APPLY_LOG_SCALING: + projected *= tl.load(log_scaling_ptr + token) + tl.store( + rel_out_ptr + row * REL_EXTENT + rel_cols, + projected.to(rel_out_ptr.dtype.element_ty), + mask=rel_mask, + ) + else: + row = block - num_q_rows + if row < tokens * NUM_KV_HEADS: + token = row // NUM_KV_HEADS + head = row % NUM_KV_HEADS + position = tl.load(positions_ptr + token) + request = tl.load(seq_idx_ptr + token) + conv_slot = tl.load(conv_slot_mapping_ptr + token) + query_start = tl.load(query_start_ptr + token) + attention_slot = tl.load(attention_slot_mapping_ptr + token) + valid = conv_slot >= 0 + + k_col = Q_WIDTH + head * HEAD_DIM + v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM + k_value = tl.load(qkvr_ptr + token * stride_x_t + k_col + dims) + v_value = tl.load(qkvr_ptr + token * stride_x_t + v_col + dims) + + safe_slot = tl.maximum(conv_slot, 0) + cache_base = ( + conv_cache_ptr + + (safe_slot // conv_block_size) * stride_cc_block + + head * stride_cc_head + + (safe_slot % conv_block_size) * stride_cc_token + ) + tl.store( + cache_base + (OFF_K + dims) * stride_cc_dim, + k_value, + mask=valid, + ) + tl.store( + cache_base + (OFF_V + dims) * stride_cc_dim, + v_value, + mask=valid, + ) + + acc_k = tl.zeros([HEAD_DIM], dtype=tl.float32) + acc_v = tl.zeros([HEAD_DIM], dtype=tl.float32) + for tap in tl.static_range(WINDOW_SIZE): + source_position = position - (WINDOW_SIZE - 1) + tap + source_row = token - (WINDOW_SIZE - 1) + tap + in_window = valid & (source_position >= 0) + intra = in_window & (source_row >= query_start) + cached = in_window & (source_row < query_start) + safe_row = tl.maximum(source_row, 0) + source_k = tl.load( + qkvr_ptr + safe_row * stride_x_t + k_col + dims, + mask=intra, + other=0.0, + ).to(tl.float32) + source_v = tl.load( + qkvr_ptr + safe_row * stride_x_t + v_col + dims, + mask=intra, + other=0.0, + ).to(tl.float32) + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum( + safe_position // conv_block_size, max_blocks - 1 + ) + physical_block = tl.load( + conv_block_table_ptr + + request * stride_block_table_req + + logical_block, + mask=cached, + other=0, + ).to(tl.int64) + tap_base = ( + conv_cache_ptr + + physical_block * stride_cc_block + + head * stride_cc_head + + (safe_position % conv_block_size) * stride_cc_token + ) + source_k += tl.load( + tap_base + (OFF_K + dims) * stride_cc_dim, + mask=cached, + other=0.0, + ).to(tl.float32) + source_v += tl.load( + tap_base + (OFF_V + dims) * stride_cc_dim, + mask=cached, + other=0.0, + ).to(tl.float32) + k_weight = tl.load( + k_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + v_weight = tl.load( + v_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + acc_k += source_k * k_weight + acc_v += source_v * v_weight + + k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + k_float = k_rounded.to(tl.float32) + k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=0) / HEAD_DIM + eps) + k_normalized = (k_float * rstd * k_norm_weight).to( + qkvr_ptr.dtype.element_ty + ) + + safe_attention_slot = tl.maximum(attention_slot, 0) + attention_block = safe_attention_slot // attention_page_size + attention_offset = safe_attention_slot % attention_page_size + tl.store( + key_cache_ptr + + attention_block * stride_kc_block + + attention_offset * stride_kc_token + + head * stride_kc_head + + dims, + k_normalized, + mask=attention_slot >= 0, + ) + tl.store( + value_cache_ptr + + attention_block * stride_vc_block + + attention_offset * stride_vc_token + + head * stride_vc_head + + dims, + v_rounded, + mask=attention_slot >= 0, + ) + + +@triton.jit(do_not_specialize=["num_rows"]) +def _q_kernel( + qkvr_ptr, + q_norm_weight_ptr, + q_out_ptr, + log_scaling_ptr, + num_rows, + stride_x_t, + eps, + NUM_Q_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_ROWS: tl.constexpr, + APPLY_LOG_SCALING: tl.constexpr, +): + rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dims = tl.arange(0, HEAD_DIM) + row_mask = rows < num_rows + tokens = rows // NUM_Q_HEADS + heads = rows % NUM_Q_HEADS + values = tl.load( + qkvr_ptr + + tokens[:, None] * stride_x_t + + heads[:, None] * HEAD_DIM + + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ).to(tl.float32) + weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(values * values, axis=1) / HEAD_DIM + eps) + normalized = values * rstd[:, None] * weight[None, :] + if APPLY_LOG_SCALING: + normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32) + tau = tl.load(log_scaling_ptr + tokens, mask=row_mask, other=1.0) + normalized *= tau[:, None] + tl.store( + q_out_ptr + rows[:, None] * HEAD_DIM + dims[None, :], + normalized.to(q_out_ptr.dtype.element_ty), + mask=row_mask[:, None], + ) + + +@triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"]) +def _kv_kernel( + qkvr_ptr, + k_weight_ptr, + v_weight_ptr, + k_norm_weight_ptr, + conv_cache_ptr, + key_cache_ptr, + value_cache_ptr, + positions_ptr, + seq_idx_ptr, + conv_slot_mapping_ptr, + conv_block_table_ptr, + query_start_ptr, + attention_slot_mapping_ptr, + tokens, + eps, + stride_x_t, + stride_cc_block, + stride_cc_head, + stride_cc_token, + stride_cc_dim, + stride_kc_block, + stride_kc_token, + stride_kc_head, + stride_vc_block, + stride_vc_token, + stride_vc_head, + stride_block_table_req, + max_blocks, + conv_block_size, + attention_page_size, + Q_WIDTH: tl.constexpr, + KV_WIDTH: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + WINDOW_SIZE: tl.constexpr, + OFF_K: tl.constexpr, + OFF_V: tl.constexpr, + BLOCK_ROWS: tl.constexpr, +): + token = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + dims = tl.arange(0, HEAD_DIM) + row_mask = token < tokens + head_id = tl.program_id(1) + head = tl.full([BLOCK_ROWS], head_id, tl.int64) + position = tl.load(positions_ptr + token, mask=row_mask, other=0) + request = tl.load(seq_idx_ptr + token, mask=row_mask, other=0) + conv_slot = tl.load(conv_slot_mapping_ptr + token, mask=row_mask, other=-1) + query_start = tl.load(query_start_ptr + token, mask=row_mask, other=0) + attention_slot = tl.load( + attention_slot_mapping_ptr + token, mask=row_mask, other=-1 + ) + valid = row_mask & (conv_slot >= 0) + + k_col = Q_WIDTH + head * HEAD_DIM + v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM + k_value = tl.load( + qkvr_ptr + token[:, None] * stride_x_t + k_col[:, None] + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ) + v_value = tl.load( + qkvr_ptr + token[:, None] * stride_x_t + v_col[:, None] + dims[None, :], + mask=row_mask[:, None], + other=0.0, + ) + safe_slot = tl.maximum(conv_slot, 0) + cache_base = ( + conv_cache_ptr + + (safe_slot // conv_block_size) * stride_cc_block + + head * stride_cc_head + + (safe_slot % conv_block_size) * stride_cc_token + ) + tl.store( + cache_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim, + k_value, + mask=valid[:, None], + ) + tl.store( + cache_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim, + v_value, + mask=valid[:, None], + ) + + acc_k = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32) + acc_v = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32) + for tap in tl.static_range(WINDOW_SIZE): + source_position = position - (WINDOW_SIZE - 1) + tap + source_row = token - (WINDOW_SIZE - 1) + tap + in_window = valid & (source_position >= 0) + intra = in_window & (source_row >= query_start) + cached = in_window & (source_row < query_start) + safe_row = tl.maximum(source_row, 0) + source_k = tl.load( + qkvr_ptr + safe_row[:, None] * stride_x_t + k_col[:, None] + dims[None, :], + mask=intra[:, None], + other=0.0, + ).to(tl.float32) + source_v = tl.load( + qkvr_ptr + safe_row[:, None] * stride_x_t + v_col[:, None] + dims[None, :], + mask=intra[:, None], + other=0.0, + ).to(tl.float32) + safe_position = tl.maximum(source_position, 0) + logical_block = tl.minimum(safe_position // conv_block_size, max_blocks - 1) + physical_block = tl.load( + conv_block_table_ptr + request * stride_block_table_req + logical_block, + mask=cached, + other=0, + ).to(tl.int64) + tap_base = ( + conv_cache_ptr + + physical_block * stride_cc_block + + head * stride_cc_head + + (safe_position % conv_block_size) * stride_cc_token + ) + source_k += tl.load( + tap_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim, + mask=cached[:, None], + other=0.0, + ).to(tl.float32) + source_v += tl.load( + tap_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim, + mask=cached[:, None], + other=0.0, + ).to(tl.float32) + k_weight = tl.load( + k_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + v_weight = tl.load( + v_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap + ).to(tl.float32) + acc_k += source_k * k_weight[None, :] + acc_v += source_v * v_weight[None, :] + + k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty) + k_float = k_rounded.to(tl.float32) + k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32) + rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=1) / HEAD_DIM + eps) + k_normalized = (k_float * rstd[:, None] * k_norm_weight[None, :]).to( + qkvr_ptr.dtype.element_ty + ) + + safe_attention_slot = tl.maximum(attention_slot, 0) + attention_block = safe_attention_slot // attention_page_size + attention_offset = safe_attention_slot % attention_page_size + attention_mask = row_mask & (attention_slot >= 0) + tl.store( + key_cache_ptr + + attention_block[:, None] * stride_kc_block + + attention_offset[:, None] * stride_kc_token + + head[:, None] * stride_kc_head + + dims[None, :], + k_normalized, + mask=attention_mask[:, None], + ) + tl.store( + value_cache_ptr + + attention_block[:, None] * stride_vc_block + + attention_offset[:, None] * stride_vc_token + + head[:, None] * stride_vc_head + + dims[None, :], + v_rounded, + mask=attention_mask[:, None], + ) + + +def _run_tiled_q( + qkvr: torch.Tensor, + q_norm_weight: torch.Tensor, + q_out: torch.Tensor, + positions: torch.Tensor, + *, + eps: float, + num_q_heads: int, + head_dim: int, + log_scaling: torch.Tensor | None, +) -> None: + num_rows = qkvr.shape[0] * num_q_heads + _q_kernel[(triton.cdiv(num_rows, Q_BLOCK_ROWS),)]( + qkvr, + q_norm_weight, + q_out, + log_scaling if log_scaling is not None else positions, + num_rows, + qkvr.stride(0), + eps, + NUM_Q_HEADS=num_q_heads, + HEAD_DIM=head_dim, + BLOCK_ROWS=Q_BLOCK_ROWS, + APPLY_LOG_SCALING=log_scaling is not None, + num_warps=Q_NUM_WARPS, + ) + + +def _run_tiled_kv( + qkvr: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + conv_block_table: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + *, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + off_k: int, + off_v: int, + conv_block_size: int, +) -> None: + tokens = qkvr.shape[0] + _kv_kernel[(triton.cdiv(tokens, KV_BLOCK_ROWS), num_kv_heads)]( + qkvr, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + tokens, + eps, + qkvr.stride(0), + conv_cache.stride(0), + conv_cache.stride(1), + conv_cache.stride(2), + conv_cache.stride(3), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + conv_block_table.stride(0), + conv_block_table.shape[1], + conv_block_size, + key_cache.shape[1], + Q_WIDTH=num_q_heads * head_dim, + KV_WIDTH=num_kv_heads * head_dim, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + WINDOW_SIZE=k_weight.shape[1], + OFF_K=off_k, + OFF_V=off_v, + BLOCK_ROWS=KV_BLOCK_ROWS, + num_warps=KV_NUM_WARPS, + ) + + +def _run_fused_small( + qkvr: torch.Tensor, + q_norm_weight: torch.Tensor, + q_out: torch.Tensor, + rel_proj: torch.Tensor, + rel_out: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + conv_block_table: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + *, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + off_k: int, + off_v: int, + conv_block_size: int, + log_scaling: torch.Tensor | None, +) -> None: + tokens = qkvr.shape[0] + + num_q_rows = tokens * num_q_heads + grid = (num_q_rows + tokens * num_kv_heads,) + _qkvr_qkv_kernel[grid]( + qkvr, + q_norm_weight, + q_out, + rel_proj, + rel_out, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + log_scaling if log_scaling is not None else positions, + tokens, + eps, + qkvr.stride(0), + conv_cache.stride(0), + conv_cache.stride(1), + conv_cache.stride(2), + conv_cache.stride(3), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + conv_block_table.stride(0), + conv_block_table.shape[1], + conv_block_size, + key_cache.shape[1], + Q_WIDTH=num_q_heads * head_dim, + KV_WIDTH=num_kv_heads * head_dim, + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + WINDOW_SIZE=k_weight.shape[1], + OFF_K=off_k, + OFF_V=off_v, + APPLY_LOG_SCALING=log_scaling is not None, + D_REL=16, + REL_EXTENT=rel_proj.shape[1], + REL_EXTENT_PADDED=triton.next_power_of_2(rel_proj.shape[1]), + num_warps=SMALL_NUM_WARPS, + ) + + +def fused_qkvr_prep( + qkvr: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + rel_proj: torch.Tensor, + eps: float, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + d_rel: int, + conv_cache: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + positions: torch.Tensor, + conv_block_table: torch.Tensor, + seq_idx: torch.Tensor, + conv_slot_mapping: torch.Tensor, + query_start: torch.Tensor, + attention_slot_mapping: torch.Tensor, + off_k: int, + off_v: int, + conv_block_size: int, + log_scaling: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + assert d_rel == 16 and rel_proj.shape[0] == 16 + assert head_dim == 128 + assert qkvr.is_contiguous() + assert k_weight.stride() == (k_weight.shape[1], 1) + assert v_weight.stride() == (v_weight.shape[1], 1) + assert rel_proj.stride() == (rel_proj.shape[1], 1) + assert conv_cache.stride(3) == 1 + assert key_cache.stride(3) == 1 and value_cache.stride(3) == 1 + tokens = qkvr.shape[0] + q_out = torch.empty( + (tokens, num_q_heads * head_dim), dtype=qkvr.dtype, device=qkvr.device + ) + rel_out = torch.empty( + (tokens, num_q_heads, rel_proj.shape[1]), + dtype=qkvr.dtype, + device=qkvr.device, + ) + if tokens == 0: + return q_out, rel_out + + if tokens < SMALL_TOKEN_THRESHOLD: + _run_fused_small( + qkvr, + q_norm_weight, + q_out, + rel_proj, + rel_out, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + eps=eps, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + off_k=off_k, + off_v=off_v, + conv_block_size=conv_block_size, + log_scaling=log_scaling, + ) + return q_out, rel_out + + kv_stream = aux_stream() + assert kv_stream is not None + current_stream = torch.cuda.current_stream() + kv_stream.wait_stream(current_stream) + with torch.cuda.stream(kv_stream): + _run_tiled_kv( + qkvr, + k_weight, + v_weight, + k_norm_weight, + conv_cache, + key_cache, + value_cache, + positions, + seq_idx, + conv_slot_mapping, + conv_block_table, + query_start, + attention_slot_mapping, + eps=eps, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + off_k=off_k, + off_v=off_v, + conv_block_size=conv_block_size, + ) + _run_tiled_q( + qkvr, + q_norm_weight, + q_out, + positions, + eps=eps, + num_q_heads=num_q_heads, + head_dim=head_dim, + log_scaling=log_scaling, + ) + qkvr_rel_proj( + qkvr, + rel_proj, + rel_out, + log_scaling, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + d_rel=d_rel, + ) + current_stream.wait_stream(kv_stream) + return q_out, rel_out diff --git a/vllm/models/inkling/nvidia/ops/sconv.py b/vllm/models/inkling/nvidia/ops/sconv.py new file mode 100644 index 000000000000..7ebb2a8a0a87 --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/sconv.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short-convolution kernels backed by a paged sliding-window state cache. + +Each layer's 4 conv streams (K, V, attn-output, mlp-output) share one paged KV +cache ``[num_blocks, H, N, D]`` (head-major; see ``sconv_swa_attn.py``). A stream +occupies the contiguous D-sub-range ``[off_s, off_s + ws)`` across all ``H`` +heads, so its flat per-token width is ``H * ws`` and that is the conv channel +dim. The cache stores the conv *input* at every absolute position. + +``fused_sconv`` is the single-launch path used by the model: per token it writes +the current input to its slot and convolves the ``W`` taps ending at its +absolute position. A tap landing inside the current forward is read from the +immutable input ``x`` (row ``src - pos + pid_t``); only pre-forward taps are read +from the paged cache (window position ``src`` -> physical block via +``block_table[req, src // N]``). The just-written slot is never read back this +step, so there is no write/read hazard within or across programs -- which is +why this needs no decode-vs-prefill split and is valid for prefill / decode / +spec alike. + +All kernels address the cache purely by ``(slot, absolute_position)`` and +allocate nothing inside the captured region; their grids depend only on the +token count (``fused_sconv`` on a fixed token/channel tiling), so the same +path replays correctly under eager, breakable PIECEWISE, and FULL cudagraphs +without any data-dependent shape or branch. +""" + +from __future__ import annotations + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _fused_sconv_kernel( + x_ptr, # [T, H*WS] head-major current-token inputs (also residual) + cache_ptr, # [num_blocks, H, N, D] paged (page-strided view) + weight_ptr, # [H*WS, W] + out_ptr, # [T, H*WS] + pos_ptr, # [T] int64 absolute position per token + seq_idx_ptr, # [T] int32 token -> batch request + slot_ptr, # [T] int64 flat slot (block*N + blk_off); < 0 => PAD + block_table_ptr, # [num_reqs, max_blocks] int32 block_table + qstart_ptr, # [T] int32 first x-row of the token's request + T, # num tokens + stride_x_t, + stride_c_blk, + stride_c_h, + stride_c_n, + stride_c_d, + stride_w_d, + stride_w_w, + stride_bt_r, + MAX_BLOCKS, + N, # block_size + W: tl.constexpr, + USE_SILU: tl.constexpr, + USE_RESIDUAL: tl.constexpr, + OFF_S: tl.constexpr, + WS: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + BLOCK_C: tl.constexpr, +): + # Each program owns a [BT tokens, BLOCK_C channels] tile. The flat channel + # index packs all H heads head-major (head = c // WS, in-stream offset = + # c % WS), so one program spans heads -- no per-head launch and no + # next_power_of_2(WS) lane waste. + pid_t = tl.program_id(0) + pid_c = tl.program_id(1) + toff = pid_t * BT + tl.arange(0, BT) # [BT] token rows + coff = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) # [BLOCK_C] flat channels + C = H * WS + t_mask = toff < T + c_mask = coff < C + head = tl.minimum(coff // WS, H - 1) # clamp keeps masked lanes in-buffer + cd = OFF_S + coff % WS # cache D-index of the channel's stream slot + + slot = tl.load(slot_ptr + toff, mask=t_mask, other=-1) # [BT] + valid = slot >= 0 + pos = tl.load(pos_ptr + toff, mask=t_mask, other=0) + req = tl.load(seq_idx_ptr + toff, mask=t_mask, other=0) + qstart = tl.load(qstart_ptr + toff, mask=t_mask, other=0) + + tc_mask = t_mask[:, None] & c_mask[None, :] + + # 1) Insert each token's input into its paged slot (skip PAD rows). + xv = tl.load(x_ptr + toff[:, None] * stride_x_t + coff[None, :], mask=tc_mask) + safe_slot = tl.maximum(slot, 0) + dst = ( + cache_ptr + + (safe_slot // N)[:, None] * stride_c_blk + + head[None, :] * stride_c_h + + (safe_slot % N)[:, None] * stride_c_n + + cd[None, :] * stride_c_d + ) + tl.store(dst, xv, mask=tc_mask & valid[:, None]) + + # 2) Convolve the W taps ending at each token's `pos`. Each tap is read from + # `x` if it falls inside this forward (row >= the request's first row), else + # from the paged cache. Exactly one source is unmasked per tap, so we sum both. + acc = tl.zeros([BT, BLOCK_C], dtype=tl.float32) + for iw in tl.static_range(W): + src = pos - (W - 1) + iw # [BT] absolute window position + row = toff - (W - 1) + iw # [BT] x-row of `src` (== src - pos + token) + in_win = valid & (src >= 0) + intra = in_win & (row >= qstart) + cached = in_win & (row < qstart) + # intra-forward tap: read the immutable input x (never the slot we just + # wrote), so there is no write/read hazard. + safe_row = tl.maximum(row, 0) + xt = tl.load( + x_ptr + safe_row[:, None] * stride_x_t + coff[None, :], + mask=c_mask[None, :] & intra[:, None], + other=0.0, + ).to(tl.float32) + # pre-forward tap: read from the paged cache via the block table. Clamp + # addressing terms; the load is masked off when out of window. + safe_src = tl.maximum(src, 0) + safe_lblk = tl.minimum(safe_src // N, MAX_BLOCKS - 1) + blk = tl.load( + block_table_ptr + req * stride_bt_r + safe_lblk, mask=cached, other=0 + ).to(tl.int64) + cbase = ( + cache_ptr + + blk[:, None] * stride_c_blk + + head[None, :] * stride_c_h + + (safe_src % N)[:, None] * stride_c_n + + cd[None, :] * stride_c_d + ) + cv = tl.load(cbase, mask=c_mask[None, :] & cached[:, None], other=0.0).to( + tl.float32 + ) + wv = tl.load( + weight_ptr + coff * stride_w_d + iw * stride_w_w, mask=c_mask, other=0.0 + ).to(tl.float32) + acc += (xt + cv) * wv[None, :] + + if USE_SILU: + acc = acc * tl.sigmoid(acc) + if USE_RESIDUAL: + acc += xv.to(tl.float32) + + tl.store( + out_ptr + toff[:, None] * stride_x_t + coff[None, :], + acc.to(out_ptr.dtype.element_ty), + mask=tc_mask, + ) + + +def fused_sconv( + x: torch.Tensor, # [T, H*ws] head-major current-token inputs + weight: torch.Tensor, # [H*ws, W] + cache: torch.Tensor, # [num_blocks, H, N, D] paged + positions: torch.Tensor, # [T] int64 absolute position per token + block_table: torch.Tensor, # [num_reqs, max_blocks] int32 + seq_idx: torch.Tensor, # [T] int32 token -> batch request + slot_mapping: torch.Tensor, # [T] int64 flat slot (PAD = -1 => skip) + query_start: torch.Tensor, # [T] int32 first x-row of the token's request + off_s: int, + ws: int, + block_size: int, + activation: str | None = None, + use_residual: bool = True, +) -> torch.Tensor: + """Single-launch insert + depthwise causal conv1d over the paged cache. + + Reads same-forward taps from ``x`` and pre-forward taps from the cache, so + it is race-free in one launch for prefill / decode / spec and cudagraph-safe + under eager / piecewise / full capture. + """ + T = x.shape[0] + out = torch.empty_like(x) + if T == 0: + return out + assert x.is_contiguous() + assert cache.stride(3) == 1, "cache D-dim must be contiguous" + H = cache.shape[1] + W = weight.shape[1] + C = H * ws # flat conv channel dim (all heads, head-major) + # Tile BT tokens x BLOCK_C channels per program: enough work per CTA to + # amortize the per-token addressing, while keeping the grid large for + # prefill. BLOCK_C spans heads so there is no per-head launch. A ~2K-element + # tile at 4 warps measured best on Blackwell; larger tiles spill registers. + BLOCK_C = min(triton.next_power_of_2(C), 256) + BT = 8 + grid = (triton.cdiv(T, BT), triton.cdiv(C, BLOCK_C)) + _fused_sconv_kernel[grid]( + x, + cache, + weight, + out, + positions, + seq_idx, + slot_mapping, + block_table, + query_start, + T, + x.stride(0), + cache.stride(0), + cache.stride(1), + cache.stride(2), + cache.stride(3), + weight.stride(0), + weight.stride(1), + block_table.stride(0), + block_table.shape[1], + block_size, + W=W, + USE_SILU=activation in ("silu", "swish"), + USE_RESIDUAL=use_residual, + OFF_S=off_s, + WS=ws, + H=H, + BT=BT, + BLOCK_C=BLOCK_C, + num_warps=4, + ) + return out + + +@triton.jit +def _seq_metadata_kernel( + qsl_ptr, # [num_reqs + 1] int32 cumulative query start rows + seq_idx_ptr, # [T] int32 out: token -> owning request + query_start_ptr, # [T] int32 out: first x-row of the token's request + num_reqs, + num_actual_tokens, + num_padded_tokens, + n_iters, # ceil(log2(num_reqs)): binary-search depth + BLOCK: tl.constexpr, +): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + tok = offs.to(tl.int32) + # Largest j in [0, num_reqs) with qsl[j] <= tok. + lo = tl.zeros([BLOCK], tl.int32) + hi = tl.full([BLOCK], num_reqs - 1, tl.int32) + for _ in range(n_iters): + mid = (lo + hi + 1) // 2 + below = tl.load(qsl_ptr + mid) <= tok + lo = tl.where(below, mid, lo) + hi = tl.where(below, hi, mid - 1) + actual = offs < num_actual_tokens + padded = offs < num_padded_tokens + query_start = tl.load(qsl_ptr + lo) + tl.store(seq_idx_ptr + offs, tl.where(actual, lo, 0), mask=padded) + tl.store( + query_start_ptr + offs, + tl.where(actual, query_start, 0), + mask=padded, + ) + + +def sconv_seq_metadata( + query_start_loc: torch.Tensor, + num_reqs: int, + num_actual_tokens: int, + seq_idx_out: torch.Tensor, + query_start_out: torch.Tensor, + num_padded_tokens: int | None = None, +) -> None: + """Fill static per-token seq_idx / query_start buffers in one launch. + + Replaces the arange + searchsorted + clamp + gather + 2x copy chain of the + sconv metadata build with a single kernel writing both persistent buffers. + Padded rows are filled with zero and must have ``slot_mapping == -1``. + """ + if num_padded_tokens is None: + num_padded_tokens = num_actual_tokens + if num_padded_tokens < num_actual_tokens: + raise ValueError("num_padded_tokens must cover all actual tokens") + if num_padded_tokens > seq_idx_out.shape[0]: + raise ValueError("seq_idx_out is too small for the padded token count") + if num_padded_tokens > query_start_out.shape[0]: + raise ValueError("query_start_out is too small for the padded token count") + + BLOCK = 256 + n_iters = (num_reqs - 1).bit_length() + grid = (triton.cdiv(num_padded_tokens, BLOCK),) + _seq_metadata_kernel[grid]( + query_start_loc, + seq_idx_out, + query_start_out, + num_reqs, + num_actual_tokens, + num_padded_tokens, + n_iters, + BLOCK=BLOCK, + ) diff --git a/vllm/models/inkling/nvidia/ops/silu_and_mul.py b/vllm/models/inkling/nvidia/ops/silu_and_mul.py new file mode 100644 index 000000000000..42ec25f63235 --- /dev/null +++ b/vllm/models/inkling/nvidia/ops/silu_and_mul.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SwiGLU kernels for the Inkling MLP layers. + +``silu_and_mul_triton``: SiLU-and-mul over the checkpoint's interleaved +fused gate/up layout (dense MLP). ``sink_silu_mul_epilogue``: the sink-expert +variant with the per-expert dequant scale and per-token gamma fused in. +""" + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit(do_not_specialize=["M"]) +def _silu_and_mul_triton_kernel( + gateup_out_ptr, + down_inp_ptr, + M, + N: tl.constexpr, + GRID_SIZE: tl.constexpr, + NUM_STAGES: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + EVEN_N: tl.constexpr, + INT64_INDEX: tl.constexpr, +): + start_pid = tl.program_id(0) + if INT64_INDEX: + start_pid = start_pid.to(tl.int64) + M = M.to(tl.int64) + + NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N) + num_blocks_mn = tl.cdiv(M, BLOCK_SIZE_M) * NUM_BLOCKS_N + + for pid in tl.range(start_pid, num_blocks_mn, GRID_SIZE, num_stages=NUM_STAGES): + pid_m = pid // NUM_BLOCKS_N + pid_n = pid % NUM_BLOCKS_N + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + mask_m = offs_m < M + mask_n = offs_n < N + + # Interleaved fused gate/up: [g0, u0, g1, u1, ...]. + mask_offs_2n = pid_n * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) // 2 + tl.static_assert(BLOCK_SIZE_N % 8 == 0, f"{BLOCK_SIZE_N=}") + mask_2n = mask_offs_2n < N + mask_2n = tl.max_constancy(mask_2n, [16]) + + offs_2n = pid_n * 2 * BLOCK_SIZE_N + tl.arange(0, 2 * BLOCK_SIZE_N) + offs_m2n = offs_m[:, None] * N * 2 + offs_2n[None, :] + + if EVEN_N or pid_n * BLOCK_SIZE_N + BLOCK_SIZE_N <= N: + gateup_out = tl.load( + gateup_out_ptr + offs_m2n, mask=mask_m[:, None], other=0.0 + ) + else: + mask_m2n = mask_m[:, None] & mask_2n[None, :] + gateup_out = tl.load(gateup_out_ptr + offs_m2n, mask=mask_m2n, other=0.0) + + gate_out, up_out = tl.split( + tl.reshape(gateup_out, (BLOCK_SIZE_M, BLOCK_SIZE_N, 2)) + ) + gate_out = gate_out.to(tl.float32) + up_out = up_out.to(tl.float32) + + down_inp = gate_out * tl.sigmoid(gate_out) * up_out + + mask_mn = mask_m[:, None] if EVEN_N else mask_m[:, None] & mask_n[None, :] + offs_mn = offs_m[:, None] * N + offs_n[None, :] + tl.store(down_inp_ptr + offs_mn, down_inp, mask=mask_mn) + + +def silu_and_mul_triton(gateup_output: torch.Tensor) -> torch.Tensor: + """SiLU-and-mul for the interleaved fused gate/up layout. + + Adapted from ``inkling_kernels.activation.silu_and_mul_fwd`` (without MXFP). + """ + assert gateup_output.is_contiguous(), ( + f"{gateup_output.shape=} {gateup_output.stride()=}" + ) + assert gateup_output.ndim == 2, f"{gateup_output.shape=}" + + M = gateup_output.shape[0] + hidden_size = gateup_output.shape[1] + assert hidden_size % 2 == 0, f"{hidden_size=}" + N = hidden_size // 2 + + down_input = torch.empty( + (M, N), device=gateup_output.device, dtype=gateup_output.dtype + ) + if M == 0: + return down_input + + BLOCK_SIZE_N = max(8, min(256, triton.next_power_of_2(N))) + if M <= 1: + BLOCK_SIZE_M = 4 + elif M <= 256: + BLOCK_SIZE_M = 2 + elif M < 4096: + BLOCK_SIZE_M = 4 + else: + BLOCK_SIZE_M = 16 + BLOCK_SIZE_N = max(8, min(128, triton.next_power_of_2(N))) + max_grid_size = triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N) + num_sms = torch.cuda.get_device_properties( + gateup_output.device + ).multi_processor_count + grid_size = min(num_sms * 4, max_grid_size) + + _silu_and_mul_triton_kernel[(grid_size,)]( + gateup_out_ptr=gateup_output, + down_inp_ptr=down_input, + M=M, + N=N, + GRID_SIZE=grid_size, + NUM_STAGES=1, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + EVEN_N=N % BLOCK_SIZE_N == 0, + INT64_INDEX=gateup_output.nbytes >= 2**31, + num_warps=8, + ) + + return down_input + + +@triton.jit(do_not_specialize=["T"]) +def _sink_epilogue_kernel( + raw_ptr, # [T, S * 2F] gemm1 output, interleaved g/u pairs per expert block + alpha_ptr, # [S] fp32 per-expert pre-SiLU dequant scale + gamma_ptr, # [T, S] fp32 per-token sink weights (may be strided) + ratio_ptr, # [S] fp32 per-expert post-SiLU scale (gemm2 alpha ratio) + out_ptr, # [T, S * F] output + T, + stride_raw_0, + stride_gamma_0, + F: tl.constexpr, + S: tl.constexpr, + BLOCK_F: tl.constexpr, +): + pid_t = tl.program_id(0).to(tl.int64) + pid_sf = tl.program_id(1) + if pid_t >= T: + return + s = pid_sf // (F // BLOCK_F) + offs_f = (pid_sf % (F // BLOCK_F)) * BLOCK_F + tl.arange(0, BLOCK_F) + + base = pid_t * stride_raw_0 + s * 2 * F + gate = tl.load(raw_ptr + base + 2 * offs_f).to(tl.float32) + up = tl.load(raw_ptr + base + 2 * offs_f + 1).to(tl.float32) + alpha = tl.load(alpha_ptr + s) + weight = tl.load(gamma_ptr + pid_t * stride_gamma_0 + s) * tl.load(ratio_ptr + s) + + gate *= alpha + up *= alpha + h = gate * tl.sigmoid(gate) * up * weight + tl.store(out_ptr + pid_t * (S * F) + s * F + offs_f, h) + + +def sink_silu_mul_epilogue( + raw: torch.Tensor, # [T, S * 2F] gemm1 output (interleaved gate/up rows) + alphas: torch.Tensor, # [S] fp32 + gammas: torch.Tensor, # [T, S] fp32 + ratios: torch.Tensor, # [S] fp32 + n_experts: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Fused sink-expert epilogue: silu(g * a_e) * (u * a_e) * (gamma * r_e). + + One kernel replaces the per-expert dequant column scale, the SwiGLU, and + the per-token gamma multiply between the two sink GEMMs. + """ + tokens = raw.shape[0] + f = raw.shape[1] // (2 * n_experts) + out = torch.empty((tokens, n_experts * f), device=raw.device, dtype=out_dtype) + if tokens == 0: + return out + # raw may be a column-slice of a padded GEMM output (rows strided). + assert raw.stride(1) == 1 and gammas.stride(1) == 1 + # Largest power-of-two divisor of f (f = 768 -> 256), capped at 512. + block_f = min(512, f & (-f)) + _sink_epilogue_kernel[(tokens, n_experts * (f // block_f))]( + raw, + alphas, + gammas, + ratios, + out, + tokens, + raw.stride(0), + gammas.stride(0), + F=f, + S=n_experts, + BLOCK_F=block_f, + ) + return out diff --git a/vllm/models/inkling/nvidia/sconv_swa_attn.py b/vllm/models/inkling/nvidia/sconv_swa_attn.py new file mode 100644 index 000000000000..8a2b288d5ee9 --- /dev/null +++ b/vllm/models/inkling/nvidia/sconv_swa_attn.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short-conv state managed as a sliding-window KV cache. + +Each decoder layer owns one ``InklingConvState`` (an ``AttentionLayerBase``) that +emits a single ``SlidingWindowSpec`` for the layer's 4 sconv streams (K, V, +attn-output, mlp-output), packed head-major into one block: + + H = num_kv_heads (per-rank), N = block_size = sconv_kernel_size, + D = head_dim(K) + head_dim(V) + hidden/H(attn) + hidden/H(mlp) + +``D`` is TP-invariant; per rank we store ``H/TP`` heads of width ``D``. The conv +reads/writes this cache out-of-band via a custom backend; the (smaller) conv page +is padded up to the uniform attention page by ``unify_kv_cache_spec_page_size``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar + +import torch +from torch import nn + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.v1.attention.backend import ( + AttentionBackend, + AttentionCGSupport, + AttentionMetadata, + AttentionMetadataBuilder, + CommonAttentionMetadata, +) +from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, SlidingWindowSpec + +from .ops.sconv import sconv_seq_metadata + +# Stream order within the per-head packed D (== contiguous sub-ranges). +_K, _V, _ATTN, _MLP = 0, 1, 2, 3 + + +@dataclass +class InklingSconvMetadata(AttentionMetadata): + block_table: torch.Tensor # [num_reqs, max_blocks] physical blocks per req + slot_mapping: torch.Tensor # [T] int64 flat slot of each token (-1 => skip) + seq_idx: torch.Tensor # [T] int32 token -> batch request + query_start: torch.Tensor # [T] int32 first x-row of each token's request + + +class InklingSconvMetadataBuilder(AttentionMetadataBuilder[InklingSconvMetadata]): + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device) + assert isinstance(kv_cache_spec, SlidingWindowSpec) + # Persistent per-token buffers for CUDA graph capture. + max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.seq_idx_buffer = torch.empty( + max_num_tokens, dtype=torch.int32, device=device + ) + self.query_start_buffer = torch.empty( + max_num_tokens, dtype=torch.int32, device=device + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> InklingSconvMetadata: + num_reqs = common_attn_metadata.num_reqs + num_actual_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) + num_padded_tokens = common_attn_metadata.slot_mapping.shape[0] + assert num_padded_tokens >= num_actual_tokens + + # Per-token seq_idx (owning request) and query_start (first x-row of + # that request; the fused kernel uses it to tell same-forward taps, + # read from x, from pre-forward taps, read from cache) in one launch. + sconv_seq_metadata( + common_attn_metadata.query_start_loc, + num_reqs, + num_actual_tokens, + self.seq_idx_buffer, + self.query_start_buffer, + num_padded_tokens, + ) + + return InklingSconvMetadata( + block_table=common_attn_metadata.block_table_tensor, + slot_mapping=common_attn_metadata.slot_mapping[:num_padded_tokens], + seq_idx=self.seq_idx_buffer[:num_padded_tokens], + query_start=self.query_start_buffer[:num_padded_tokens], + ) + + +class InklingSconvBackend(AttentionBackend): + """Custom dummy backend for the sconv sliding-window cache management.""" + + @staticmethod + def get_name() -> str: + return "INKLING_SCONV_SWA" + + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the + # padded conv page is read through a strided view. + return True + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + # HND, num-blocks-first, head-major: [num_blocks, H, N, D]. + return (num_blocks, num_kv_heads, block_size, head_size) + + @staticmethod + def get_kv_cache_stride_order( + include_num_layers_dimension: bool = False, + ) -> tuple[int, ...]: + # Identity: physical layout == logical [num_blocks, H, N, D]. + if include_num_layers_dimension: + return (0, 1, 2, 3, 4) + return (0, 1, 2, 3) + + @staticmethod + def get_impl_cls(): + raise NotImplementedError( + "InklingSconvBackend has no attention impl; the conv runs out-of-band." + ) + + @staticmethod + def get_builder_cls() -> type[InklingSconvMetadataBuilder]: + return InklingSconvMetadataBuilder + + +class InklingConvState(nn.Module, AttentionLayerBase): + """Per-decoder-layer owner emitting one sliding-window conv-state spec.""" + + def __init__( + self, + *, + num_kv_heads: int, + head_dim: int, + hidden_size: int, + kernel_size: int, + prefix: str, + ) -> None: + super().__init__() + self.prefix = prefix + # Bound to the manager-allocated paged cache by bind_kv_cache; a + # placeholder until then. Read out-of-band by InklingShortConv. + self.kv_cache = torch.tensor([]) + tp_size = get_tensor_model_parallel_world_size() + # Guardrails for the conv-state layout below; only these are exercised. + # tp_size <= num_kv_heads keeps >=1 whole KV head per rank (no + # replication/clamping), so the per-head width stays TP-invariant. + assert tp_size <= num_kv_heads, ( + f"sconv SWA cache supports tp_size <= num_kv_heads ({num_kv_heads}), " + f"got {tp_size}" + ) + # Per-rank head count; D is TP-invariant (K/V heads and the hidden + # chunk both scale 1/TP together). The attn-/mlp-output sconv streams + # are hidden-sharded: each rank owns its H/tp chunk (the sublayer + # outputs are reduce-scattered / all-gathered around the conv). + self.num_kv_heads = num_kv_heads // tp_size + hidden_per_head = hidden_size // num_kv_heads + # Packed per-head width: K + V + attn-output chunk + mlp-output chunk, + # padded to a power of two so every layer's conv page is the same size + # and an exact multiple of the attention page (the page unifier then + # scales attention block sizes instead of padding). + raw_head_size = 2 * head_dim + 2 * hidden_per_head + self.head_size = 1 << (raw_head_size - 1).bit_length() + self.sliding_window = kernel_size + self.block_size = kernel_size + # Per-head D-sub-range (offset, width) for each stream. Streams share + # the cache; each writes/reads its own width across all H heads. + self.stream_ranges: tuple[tuple[int, int], ...] = ( + (0, head_dim), # _K + (head_dim, head_dim), # _V + (2 * head_dim, hidden_per_head), # _ATTN + (2 * head_dim + hidden_per_head, hidden_per_head), # _MLP + ) + vllm_config = get_current_vllm_config() + self._dtype = vllm_config.model_config.dtype + assert self._dtype == torch.bfloat16, ( + f"sconv SWA cache supports bfloat16 only, got {self._dtype}" + ) + # Register in the forward context so the runner enumerates this owner as + # an attention-like layer (get_kv_cache_spec / get_attn_backend). + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return InklingSconvBackend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + return SlidingWindowSpec( + block_size=self.block_size, + num_kv_heads=self.num_kv_heads, + head_size=self.head_size, + head_size_v=0, # all 4 streams packed into head_size + dtype=self._dtype, + sliding_window=self.sliding_window, + ) diff --git a/vllm/models/inkling/nvidia/short_conv.py b/vllm/models/inkling/nvidia/short_conv.py new file mode 100644 index 000000000000..363423bf2aad --- /dev/null +++ b/vllm/models/inkling/nvidia/short_conv.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling short convolution: depthwise causal conv1d (+ residual) over a paged +sliding-window conv-state cache. + +Each decoder layer owns one ``InklingConvState`` (``sconv_swa_attn.py``) holding the +manager-allocated paged cache for the layer's 4 sconv streams (K, V, attn-output, +mlp-output), packed head-major into one block. Each ``InklingShortConv`` is a +stateless weight + kernel launcher that, per forward (positions-addressed, the +same path for prefill / decode / mixed), inserts the current tokens' inputs +into their paged slot and convolves each token against the ``W`` taps ending +at its absolute position, reading pre-forward window positions out of the +paged cache via the block table. + +Per-forward metadata (``block_table`` / ``slot_mapping`` / ``seq_idx`` / +``query_start``) is built once by ``InklingSconvMetadataBuilder`` and published under +the owner's prefix in the forward context; the absolute ``positions`` are +threaded in from the model. The insert + conv run in a single ``fused_sconv`` +launch (same path for prefill / decode / mixed / spec). All inputs are +fixed-address persistent buffers and the grid is fixed, so the conv replays +correctly under eager, PIECEWISE, and FULL cudagraphs. +""" + +from __future__ import annotations + +import torch +from torch import nn +from torch.nn.parameter import Parameter + +from vllm.distributed import get_tensor_model_parallel_rank +from vllm.forward_context import get_forward_context +from vllm.model_executor.utils import set_weight_attrs + +from .ops import fused_sconv +from .sconv_swa_attn import InklingConvState, InklingSconvMetadata + + +class InklingShortConv(nn.Module): + def __init__( + self, dim: int, kernel_size: int, owner: InklingConvState, stream_idx: int + ) -> None: + super().__init__() + self.dim = dim + self.kernel_size = kernel_size + self.owner = owner + self.stream_idx = stream_idx + self.tp_rank = get_tensor_model_parallel_rank() + + # Depthwise conv weight; checkpoint stores (dim, 1, W). + self.weight = Parameter(torch.empty(dim, 1, kernel_size), requires_grad=False) + set_weight_attrs(self.weight, {"weight_loader": self.weight_loader}) + + def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor) -> None: + if loaded_weight.shape[0] != param.shape[0]: + shard = param.shape[0] + loaded_weight = loaded_weight.narrow(0, self.tp_rank * shard, shard) + param.data.copy_(loaded_weight) + + def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + # x: (num_tokens, dim); positions: (num_tokens,) absolute positions. + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + # Memory-profiling / no metadata: identity (residual). + return x + m = attn_metadata.get(self.owner.prefix) + if m is None: + return x + assert isinstance(m, InklingSconvMetadata) + cache = self.owner.kv_cache + if cache.numel() == 0: + # Cache not yet bound (profiling before KV alloc): identity. + return x + + off_s, ws = self.owner.stream_ranges[self.stream_idx] + block_size = self.owner.block_size + x = x.contiguous() + weight = self.weight.squeeze(1) # (dim, W) + + return fused_sconv( + x, + weight, + cache, + positions, + m.block_table, + m.seq_idx, + m.slot_mapping, + m.query_start, + off_s, + ws, + block_size, + activation=None, + use_residual=True, + ) diff --git a/vllm/models/minimax_m3/nvidia/indexer_msa.py b/vllm/models/minimax_m3/nvidia/indexer_msa.py index 16a9277d11af..8a015d4e6eae 100644 --- a/vllm/models/minimax_m3/nvidia/indexer_msa.py +++ b/vllm/models/minimax_m3/nvidia/indexer_msa.py @@ -13,10 +13,9 @@ than Triton for the wide prefill score, benchmarked ~3-5x), writing its ``max_score`` straight into the buffer's prefill region (stride-aware, no copy). -Decode scores with the Triton split-K ``minimax_m3_index_decode_score`` (a -purpose-built vector x matrix score, no wasted tensor-core tiles, 256-way -split-K, cudagraph-safe by shape-constant grids), writing into the decode -region. Its tuning heuristics are kept; only the top-k is shared with prefill. +Decode scores with CuteDSL when the flattened query tile is supported and fall +back to Triton otherwise, writing into the decode region. Only the top-k is +shared with prefill. ``fmha_sm100`` imports are function-local so this module is import-safe on AMD / non-SM100. @@ -39,6 +38,7 @@ from vllm.models.minimax_m3.common.ops.index_topk import ( minimax_m3_index_decode_score, ) +from vllm.models.minimax_m3.nvidia.ops import minimax_m3_index_decode_score_cutedsl from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -257,7 +257,7 @@ def build( class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl): - """Decode: Triton fused score+top-k. Prefill: fmha_sm100 OnlyScore + top-k.""" + """Decode: CuteDSL/Triton score. Prefill: fmha_sm100 OnlyScore + top-k.""" indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend @@ -296,7 +296,14 @@ def forward( # writes by strides). Top-k is deferred to the single unified call below. if md.decode is not None: d = md.decode - minimax_m3_index_decode_score( + # max_decode_query_len avoids recompiles across runtime decode sizes. + # Fall back when the flattened Q tile gets too wide for this kernel. + decode_score = ( + minimax_m3_index_decode_score_cutedsl + if self.num_index_heads * d.max_decode_query_len <= 32 + else minimax_m3_index_decode_score + ) + decode_score( index_q[:nd], kv, d.block_table, diff --git a/vllm/models/minimax_m3/nvidia/ops/__init__.py b/vllm/models/minimax_m3/nvidia/ops/__init__.py new file mode 100644 index 000000000000..f2557575536a --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/ops/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .index_decode_score import minimax_m3_index_decode_score_cutedsl + +__all__ = ["minimax_m3_index_decode_score_cutedsl"] diff --git a/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py b/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py new file mode 100644 index 000000000000..3d2919e39227 --- /dev/null +++ b/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py @@ -0,0 +1,483 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CuteDSL MiniMax M3 index decode score kernel. + +The kernel computes decode-time index block scores with TMA + ``mma.sync``. +We use ``mma.sync`` instead of tcgen05 because this score GEMM has a very small +N dimension and benefits more from higher CTA occupancy than from a deeper +single-CTA tcgen05 pipeline. + +The implementation should be portable to SM90/SM120 in principle, but it is +currently validated only for SM100. +""" + +from functools import cache + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass.cute.nvgpu import cpasync, warp +from quack.compile_utils import make_fake_tensor + +from vllm.cute_utils import ( + _TORCH_TO_CUTE_DTYPE, + EVICT_FIRST, + cvt, + mma_sync, + simple_tma_copy, +) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # This packed conversion is faster and emits fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = cvt.fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix gives four consecutive values along K. Split each group + # into the lower two and upper two values for two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + BLOCK_K = 128 + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + ): + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + dtype = self.dtype + num_heads = self.num_heads + head_dim = self.head_dim + BLOCK_K = self.BLOCK_K + num_stages = self.num_stages + MAX_DQL = self.max_decode_query_len + BLOCK_Q = num_heads * MAX_DQL + assert BLOCK_Q <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + swizzle_128B = cute.make_swizzle(3, 4, 3) + elems = 128 * 8 // dtype.width + + sQ_layout = cute.make_layout( + (MAX_DQL, num_heads, (elems, head_dim // elems)), + stride=(elems, MAX_DQL * elems, (1, BLOCK_Q * elems)), + ) + sQ_layout = cute.make_composed_layout(swizzle_128B, 0, sQ_layout) + Q_tma = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + sQ_layout, + cta_tiler=(MAX_DQL, num_heads, head_dim), + ) + + sK_layout = cute.make_layout( + (1, BLOCK_K, (elems, head_dim // elems), num_stages), + stride=(0, elems, (1, BLOCK_K * elems), BLOCK_K * head_dim), + ) + sK_layout = cute.make_composed_layout(swizzle_128B, 0, sK_layout) + K_tma = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + sK_layout, + cta_tiler=(1, BLOCK_K, head_dim), + ) + + self.kernel( + Q_tma, + K_tma, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=True) + + @cute.kernel + def kernel( + self, + Q_tma: cpasync.TmaInfo, + K_tma: cpasync.TmaInfo, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + tid, _, _ = cute.arch.thread_idx() + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + smem = cutlass.utils.SmemAllocator() + sK = smem.allocate_tensor( + dtype, + K_tma.smem_layout.outer, + byte_alignment=128, + swizzle=K_tma.smem_layout.inner, + )[0, None, None, None] + # alias sQ with the 1st stage of sK + sQ_tma = cute.make_tensor( + sK[None, None, 0].iterator, layout=Q_tma.smem_layout.outer + ) + # TMA sees Q as (query, head, dim), while ldmatrix consumes a + # flattened Q column mode. The target profile keeps the rank-2 view even + # for degenerate shapes like DQL1. + q_tma_elems = 128 * 8 // dtype.width + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(Q_tma.atom) + cpasync.prefetch_descriptor(K_tma.atom) + cute.arch.sync_threads() + + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + + gQ_tile = cute.local_tile( + cute.domain_offset( + (batch_id * decode_query_len, 0, 0), + Q_tma.tma_tensor, + ), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when runtime decode_query_len is smaller + # than MAX_DQL; padded Q columns are masked before global stores. + simple_tma_copy(Q_tma.atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(split_id, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = K_tma.tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + K_tma.atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps + # each warp handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile( + sK, (32, head_dim, num_stages), (warp_id, 0, 0) + ) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix address. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide( + sK_warp, (16, cute.make_layout((elems, 2)), 1) + ) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + # sm100 doesn't have native mma.sync.f8. ptxas lowers mma.sync.f8 + # to F2FP.F16.E4M3 + HMMA; doing the conversion explicitly gives + # better codegen while keeping the two FP16 k-fragments visible. + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor( + (4, head_dim // MMA_K, Q_TILES, 2), Float16 + ) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments( + rK[None, m, k] + ) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + # causal mask + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] + if q_pos >= k_pos + else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # thread-reduction along BLOCK_K dim + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # warp-reduction among lanes 0,4,8,12,... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + # store to smem for 4-warp reduction + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax( + final_score, epi_buffer[lane_id, i] + ) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + + @cache + @staticmethod + def compile( + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + ): + bs = cute.sym_int() + total_tokens = cute.sym_int() + BLOCK_K = IndexDecodeScoreKernel.BLOCK_K + + q = make_fake_tensor( + dtype, (total_tokens, num_heads, head_dim), divisibility=16 + ) + k_cache = make_fake_tensor( + dtype, (cute.sym_int(), BLOCK_K, head_dim), divisibility=16 + ) + block_table = make_fake_tensor(Int32, (bs, cute.sym_int()), divisibility=1) + score = make_fake_tensor( + Float32, (num_heads, total_tokens, cute.sym_int()), divisibility=4 + ) + seq_lens = make_fake_tensor(Int32, (bs,), divisibility=1) + kernel = IndexDecodeScoreKernel( + dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + ) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel, + q, + k_cache, + block_table, + score, + seq_lens, + stream, + options="--enable-tvm-ffi", + ) + + +def minimax_m3_index_decode_score_cutedsl( + idx_q: torch.Tensor, + index_kv_cache: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + init_blocks: int, + local_blocks: int, + num_kv_heads: int, + decode_query_len: int, + max_decode_query_len: int, + score_out: torch.Tensor, +) -> torch.Tensor: + if idx_q.dtype not in (torch.bfloat16, torch.float8_e4m3fn): + raise TypeError("CuteDSL index decode score supports BF16 and FP8 E4M3 only") + total_tokens, num_heads, head_dim = idx_q.shape + batch = block_table.shape[0] + assert index_kv_cache.shape[1] == IndexDecodeScoreKernel.BLOCK_K + assert total_tokens == batch * decode_query_len + assert 1 <= decode_query_len <= max_decode_query_len + assert num_heads * max_decode_query_len <= 32 + dtype = _TORCH_TO_CUTE_DTYPE[idx_q.dtype] + del max_seq_len, init_blocks, local_blocks, num_kv_heads + score = score_out + split_k = 256 + kernel = IndexDecodeScoreKernel.compile( + dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + ) + kernel(idx_q, index_kv_cache, block_table, score, seq_lens) + return score diff --git a/vllm/multimodal/utils.py b/vllm/multimodal/utils.py index 56c65dc8ea72..f34769d838bd 100644 --- a/vllm/multimodal/utils.py +++ b/vllm/multimodal/utils.py @@ -346,3 +346,41 @@ def fetch_video( allowed_local_media_path="/", ) return media_connector.fetch_video(video_url) + + +def set_mm_embedding_modality(embed: "torch.Tensor", modality: str) -> "torch.Tensor": + """Attach modality metadata to a gathered multimodal embedding tensor. + + Used by interleaved Omni merge paths that need to group embeddings by + modality without threading a parallel modalities list through + ``embed_input_ids``. + """ + embed.modality = modality # type: ignore[attr-defined] + return embed + + +def copy_mm_embedding_modality( + src: "torch.Tensor", dst: "torch.Tensor" +) -> "torch.Tensor": + """Copy ``modality`` from ``src`` onto ``dst`` if present.""" + modality = getattr(src, "modality", None) + if modality is not None: + dst.modality = modality # type: ignore[attr-defined] + return dst + + +def get_mm_embedding_modalities( + multimodal_embeddings: Sequence["torch.Tensor"], +) -> list[str]: + """Collect per-embedding modalities previously set on the tensors.""" + modalities: list[str] = [] + for i, emb in enumerate(multimodal_embeddings): + modality = getattr(emb, "modality", None) + if modality is None: + raise ValueError( + f"Missing modality on multimodal embedding at index {i}. " + "Encoder gather must set embed.modality before interleaved " + "audio-in-video merge." + ) + modalities.append(modality) + return modalities diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index fc2d653779ed..895f7516a53e 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -110,6 +110,8 @@ def __init__( self._has_reasoning = ( "THINK_END" in parser_engine_config.token_id_terminals + or "THINK_START" in parser_engine_config.terminals + or "THINK_END" in parser_engine_config.terminals or parser_engine_config.initial_state == ParserState.REASONING ) self._reasoning_ended: bool = not self._has_reasoning diff --git a/vllm/parser/engine/parser_engine_config.py b/vllm/parser/engine/parser_engine_config.py index f18f3f02e1e2..ad83e331490a 100644 --- a/vllm/parser/engine/parser_engine_config.py +++ b/vllm/parser/engine/parser_engine_config.py @@ -28,6 +28,7 @@ class ParserState(Enum): CONTENT = auto() REASONING = auto() + MESSAGE_HEADER = auto() TOOL_PREAMBLE = auto() TOOL_NAME = auto() TOOL_ARGS = auto() diff --git a/vllm/parser/engine/registered_adapters.py b/vllm/parser/engine/registered_adapters.py index 889c71504fd9..d259ade545c2 100644 --- a/vllm/parser/engine/registered_adapters.py +++ b/vllm/parser/engine/registered_adapters.py @@ -12,6 +12,7 @@ from vllm.parser.engine.adapters import make_adapters from vllm.parser.gemma4 import Gemma4Parser from vllm.parser.glm47_moe import Glm47MoeParser +from vllm.parser.inkling import InklingParser from vllm.parser.kimi_k2 import KimiK2Parser from vllm.parser.minimax_m2 import MinimaxM2Parser from vllm.parser.nemotron_v3 import NemotronV3Parser @@ -62,3 +63,8 @@ KimiK2ParserReasoningAdapter, KimiK2ParserToolAdapter, ) = make_adapters(KimiK2Parser) + +( + InklingParserReasoningAdapter, + InklingParserToolAdapter, +) = make_adapters(InklingParser) diff --git a/vllm/parser/engine/streaming_parser_engine.py b/vllm/parser/engine/streaming_parser_engine.py index 319791961c2e..8cafdf8e625b 100644 --- a/vllm/parser/engine/streaming_parser_engine.py +++ b/vllm/parser/engine/streaming_parser_engine.py @@ -269,6 +269,8 @@ def finish(self) -> list[SemanticEvent]: SemanticEvent(EventType.REASONING_END, tool_index=self.tool_index) ) self.state = ParserState.CONTENT + elif self.state == ParserState.MESSAGE_HEADER: + self.state = ParserState.CONTENT return events @@ -314,6 +316,15 @@ def _on_terminal(self, terminal: str, value: str) -> list[SemanticEvent]: return self._emit_for_state(value) if self.skip_tool_parsing and terminal in self._tool_terminals: + if self.state == ParserState.MESSAGE_HEADER: + self.state = ParserState.CONTENT + return [ + SemanticEvent( + EventType.TEXT_CHUNK, + value=value, + tool_index=self.tool_index, + ) + ] if EventType.REASONING_END in transition.events: self.state = ParserState.CONTENT return [ diff --git a/vllm/parser/inkling.py b/vllm/parser/inkling.py new file mode 100644 index 000000000000..84502be018e3 --- /dev/null +++ b/vllm/parser/inkling.py @@ -0,0 +1,411 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling parser: typed content blocks parsed by a single state machine. + +Inkling output format (every marker is a dedicated special token):: + + <|message_model|><|content_thinking|>...reasoning...<|end_message|> + <|message_model|><|content_text|>...visible text...<|end_message|> + <|message_model|><|content_invoke_tool_json|> + {"name":"get_weather","args":{"city":"SF"}}<|end_message|> + +Blocks are self-describing and may repeat in any order; sampling may +also end a block with the standalone ``<|content_model_end_sampling|>`` +token. The tool-call payload is a single JSON object whose ``name`` is +extracted by the engine's name-from-args path and whose ``args`` object +is carved out of the wrapper by :func:`_inkling_arg_converter`. + +Note the terminal *labels*: ``THINK_START``/``THINK_END`` are what the +engine keys its reasoning plumbing on (``is_reasoning_end``, +``count_reasoning_tokens``, initial-state seeding), so ``<|end_message|>`` +is labelled ``THINK_END`` here even though it ends every block kind — +the transition table, not the label, carries the semantics. +""" + +from __future__ import annotations + +import functools +import json +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from vllm.entrypoints.openai.engine.protocol import ( + ExtractedToolCallInformation, +) +from vllm.parser.engine.events import EventType +from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.engine.parser_engine_config import ( + ParserEngineConfig, + ParserState, + Transition, +) + +if TYPE_CHECKING: + from vllm.tokenizers import TokenizerLike + from vllm.tool_parsers.abstract_tool_parser import Tool + +MESSAGE_MODEL = "<|message_model|>" +CONTENT_TEXT = "<|content_text|>" +CONTENT_THINKING = "<|content_thinking|>" +CONTENT_INVOKE_TOOL_JSON = "<|content_invoke_tool_json|>" +CONTENT_INVOKE_TOOL_TEXT = "<|content_invoke_tool_text|>" +CONTENT_TOOL_ERROR = "<|content_tool_error|>" +CONTENT_MODEL_END_SAMPLING = "<|content_model_end_sampling|>" +END_MESSAGE = "<|end_message|>" + +INKLING_SPECIAL_TOKENS = ( + MESSAGE_MODEL, + CONTENT_TEXT, + CONTENT_THINKING, + CONTENT_INVOKE_TOOL_JSON, + CONTENT_INVOKE_TOOL_TEXT, + CONTENT_TOOL_ERROR, + CONTENT_MODEL_END_SAMPLING, + END_MESSAGE, +) + +_WS = " \t\r\n" + + +def _scan_json_value(raw: str, start: int) -> int | None: + """Return the end index (exclusive) of the JSON object starting at + ``raw[start]``, or ``None`` when the object is still unterminated.""" + depth = 0 + in_string = False + escape = False + for i in range(start, len(raw)): + ch = raw[i] + if escape: + escape = False + continue + if in_string: + if ch == "\\": + escape = True + elif ch == '"': + in_string = False + continue + if ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + 1 + return None + + +def _args_value_span(raw: str) -> str | None: + """Extract the raw text span of the top-level ``"args"`` value from a + (possibly incomplete) ``{"name":...,"args":{...}}`` wrapper. + + Returns the verbatim substring (prefix-stable across growing input, + which the engine's argument-delta diffing relies on), possibly an + unterminated object prefix; ``None`` when the value has not started. + Raises ``ValueError`` when the value is not a JSON object. + """ + depth = 0 + in_string = False + escape = False + string_start = -1 + last_string: str | None = None + for i, ch in enumerate(raw): + if escape: + escape = False + continue + if in_string: + if ch == "\\": + escape = True + elif ch == '"': + in_string = False + if depth == 1: + last_string = raw[string_start + 1 : i] + continue + if ch == '"': + in_string = True + string_start = i + elif ch == ":" and depth == 1 and last_string == "args": + value_start = i + 1 + while value_start < len(raw) and raw[value_start] in _WS: + value_start += 1 + if value_start >= len(raw): + return None + if raw[value_start] != "{": + raise ValueError("Inkling tool call args must be a JSON object") + value_end = _scan_json_value(raw, value_start) + if value_end is None: + return raw[value_start:] + return raw[value_start:value_end] + elif ch in "{[": + depth += 1 + elif ch in "}]": + depth -= 1 + return None + + +def _inkling_arg_converter(raw_args: str, partial: bool) -> str: + """Carve the ``args`` object out of the tool-call JSON wrapper. + + Why a converter at all: the engine's ``tool_args_json`` machinery + treats the *entire* TOOL_ARGS text as the tool arguments, but Inkling's + payload is the ``{"name":...,"args":{...}}`` wrapper — without a + converter, ``_compute_arg_delta`` streams the wrapper verbatim into + the OpenAI ``arguments`` field (``converter is None -> raw delta``, + unconditionally; ``stream_arg_deltas=False`` does not stop it). + + Why a hand-rolled scanner instead of (partial) ``json.loads`` + + ``json.dumps``: the engine diffs successive converter outputs and + requires each to extend the previous one (``startswith``); a + violation silently drops argument deltas. Re-serialization changes + whitespace and closes unterminated structures differently across + ticks, so the only prefix-stable output is a verbatim substring of + the input. The scanner is also string/escape-aware so an ``"args"`` + literal inside the name or a string value cannot mislead it, and it + still recovers a partial span when EOS truncates the wrapper (where + ``json.loads`` would fail). + """ + span = _args_value_span(raw_args) + if span is None: + # No args value yet (streaming) or none at all (treat as empty). + return "" if partial else "{}" + return span + + +@functools.cache +def inkling_config() -> ParserEngineConfig: + terminals = { + "MSG_MODEL": MESSAGE_MODEL, + "TEXT_START": CONTENT_TEXT, + "THINK_START": CONTENT_THINKING, + "THINK_END": END_MESSAGE, + "END_SAMPLING": CONTENT_MODEL_END_SAMPLING, + "TOOL_START": CONTENT_INVOKE_TOOL_JSON, + "TOOL_TEXT": CONTENT_INVOKE_TOOL_TEXT, + "TOOL_ERROR": CONTENT_TOOL_ERROR, + } + transitions: dict[tuple[ParserState, str], Transition] = { + # ── Between blocks / inside a text block ────────────────────── + (ParserState.CONTENT, "MSG_MODEL"): Transition( + ParserState.MESSAGE_HEADER, + (), + ), + (ParserState.CONTENT, "TEXT_START"): Transition( + ParserState.CONTENT, + (), + ), + (ParserState.CONTENT, "THINK_START"): Transition( + ParserState.REASONING, + (EventType.REASONING_START,), + ), + (ParserState.CONTENT, "TOOL_START"): Transition( + ParserState.TOOL_ARGS, + (EventType.TOOL_CALL_START,), + ), + # Raw / error tool blocks render as visible text. + (ParserState.CONTENT, "TOOL_TEXT"): Transition( + ParserState.CONTENT, + (), + ), + (ParserState.CONTENT, "TOOL_ERROR"): Transition( + ParserState.CONTENT, + (), + ), + # The optional function name between the model-role and content-kind + # markers is metadata, not visible assistant content. + (ParserState.MESSAGE_HEADER, "MSG_MODEL"): Transition( + ParserState.MESSAGE_HEADER, + (), + ), + (ParserState.MESSAGE_HEADER, "TEXT_START"): Transition( + ParserState.CONTENT, + (), + ), + (ParserState.MESSAGE_HEADER, "THINK_START"): Transition( + ParserState.REASONING, + (EventType.REASONING_START,), + ), + (ParserState.MESSAGE_HEADER, "TOOL_START"): Transition( + ParserState.TOOL_ARGS, + (EventType.TOOL_CALL_START,), + ), + (ParserState.MESSAGE_HEADER, "TOOL_TEXT"): Transition( + ParserState.CONTENT, + (), + ), + (ParserState.MESSAGE_HEADER, "TOOL_ERROR"): Transition( + ParserState.CONTENT, + (), + ), + # ── Inside a thinking block ─────────────────────────────────── + (ParserState.REASONING, "THINK_START"): Transition( + ParserState.REASONING, + (), + ), + # Defensive: tool call opening while a thinking block is still + # unclosed. + (ParserState.REASONING, "TOOL_START"): Transition( + ParserState.TOOL_ARGS, + (EventType.REASONING_END, EventType.TOOL_CALL_START), + ), + } + # Block-end terminals behave identically regardless of label. A + # closed tool block returns to CONTENT (Inkling has no section wrapper; + # blocks of any kind may follow), which also keeps the block-kind + # and role tokens out of the engine's tool-terminal set so the + # skip_tool_parsing reasoning pass still classifies reasoning. + for end in ("THINK_END", "END_SAMPLING"): + transitions[(ParserState.CONTENT, end)] = Transition( + ParserState.CONTENT, + (), + ) + transitions[(ParserState.REASONING, end)] = Transition( + ParserState.CONTENT, + (EventType.REASONING_END,), + ) + transitions[(ParserState.TOOL_ARGS, end)] = Transition( + ParserState.CONTENT, + (EventType.TOOL_CALL_END,), + ) + transitions[(ParserState.MESSAGE_HEADER, end)] = Transition( + ParserState.CONTENT, + (), + ) + + return ParserEngineConfig( + name="inkling", + # Normal generation continues after a prompt-prefilled + # `<|message_model|>`. Non-streaming parsing receives only the generated + # suffix, so begin in the corresponding message-header state as well. + initial_state=ParserState.MESSAGE_HEADER, + terminals=terminals, + # Inkling content-kind markers are the grammar. When the engine is + # used through DelegatingParser, the reasoning pass can hand the tool + # pass reconstructed text whose token-id slice no longer contains the + # content-kind marker that starts the tool block, so keep Inkling on + # the text grammar instead of token-id-only terminal matching. + token_id_terminals={}, + transitions=transitions, + arg_converter=_inkling_arg_converter, + stream_arg_deltas=True, + tool_args_json=True, + strip_trailing_reasoning_whitespace=True, + drop_whitespace_only_content_before_tools=True, + strip_content_whitespace_with_tools=False, + validate_tool_names=False, + ) + + +class InklingParser(ParserEngine): + CONFIG_NAME = "inkling" + + def __init__( + self, + tokenizer: TokenizerLike, + tools: list[Tool] | None = None, + **kwargs, + ) -> None: + kwargs.setdefault("parser_engine_config", inkling_config()) + super().__init__(tokenizer, tools, **kwargs) + + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None: + """Seed the initial parsing state from the prompt tail. + + Mirrors the Rust parser's ``initialize()``: scanning the prompt + backwards, the last relevant special token decides whether generation + continues inside a thinking block, a text block, a model message + header, or between blocks. + """ + vocab = self.vocab + thinking_id = vocab.get(CONTENT_THINKING) + text_id = vocab.get(CONTENT_TEXT) + model_id = vocab.get(MESSAGE_MODEL) + special_ids = {vocab[text] for text in INKLING_SPECIAL_TOKENS if text in vocab} + for token_id in reversed(prompt_token_ids): + if token_id == thinking_id: + self._engine.reset(initial_state=ParserState.REASONING) + self._streaming_initialized = True + return + if token_id == text_id: + self._engine.reset(initial_state=ParserState.CONTENT) + self._streaming_initialized = True + return + if token_id == model_id: + self._engine.reset(initial_state=ParserState.MESSAGE_HEADER) + self._streaming_initialized = True + return + if token_id in special_ids: + break + + def is_reasoning_end(self, input_ids: list[int]) -> bool: + vocab = self.vocab + thinking_id = vocab.get(CONTENT_THINKING) + text_id = vocab.get(CONTENT_TEXT) + model_id = vocab.get(MESSAGE_MODEL) + end_sampling_id = vocab.get(CONTENT_MODEL_END_SAMPLING) + for token_id in reversed(input_ids): + if token_id in (thinking_id, model_id): + return False + if token_id in (text_id, end_sampling_id): + return True + return False + + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + vocab = self.vocab + thinking_id = vocab.get(CONTENT_THINKING) + end_ids = { + token_id + for token_id in ( + vocab.get(END_MESSAGE), + vocab.get(CONTENT_MODEL_END_SAMPLING), + ) + if token_id is not None + } + in_reasoning = False + count = 0 + for token_id in token_ids: + if token_id == thinking_id: + in_reasoning = True + continue + if token_id in end_ids: + in_reasoning = False + continue + if in_reasoning: + count += 1 + return count + + def _single_pass_parse( + self, + text: str, + token_ids: Sequence[int], + initial_state: ParserState | None = None, + ) -> tuple[str | None, str | None, ExtractedToolCallInformation]: + reasoning, content, tool_call_info = super()._single_pass_parse( + text, token_ids, initial_state=initial_state + ) + # The engine defers content that follows tool-call events within a + # single pass; Inkling allows text blocks after tool-call blocks, so + # flush the trailing text (matching the Rust unified parser). + if self._deferred_content: + trailing = self._deferred_content + self._deferred_content = "" + content = self._strip_content_whitespace( + (content or "") + trailing, + tool_call_info.tools_called, + ) + tool_call_info = ExtractedToolCallInformation( + tools_called=tool_call_info.tools_called, + tool_calls=tool_call_info.tool_calls, + content=content, + ) + return reasoning, content, tool_call_info + + @staticmethod + def _extract_args_value(parsed: dict) -> str | None: + # Inkling wraps arguments under "args" rather than "arguments". + for key in ("args", "arguments", "parameters"): + if key in parsed: + val = parsed[key] + if isinstance(val, str): + return val + return json.dumps(val, ensure_ascii=False) + return None diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index cbfa579313b3..80f99acc5acb 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -124,14 +124,6 @@ def get_attn_backend_cls( attn_selector_config: "AttentionSelectorConfig", num_heads: int | None = None, ) -> str: - from vllm.v1.attention.backends.utils import set_kv_cache_layout - - set_kv_cache_layout("NHD") - logger.info_once( - "Setting VLLM_KV_CACHE_LAYOUT to 'NHD' for XPU; " - "only NHD layout is supported by XPU attention kernels." - ) - # TurboQuant KV cache: route directly to TQ backend kv_cache_dtype = attn_selector_config.kv_cache_dtype if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"): @@ -150,8 +142,18 @@ def get_attn_backend_cls( return AttentionBackendEnum.TRITON_ATTN.get_path() elif attn_selector_config.use_mm_prefix: # Flash Attention on XPU has no FA4 kernel, so it cannot apply the - # multimodal prefix-LM bidirectional mask. Fall back to Triton - # Attention, which supports mm_prefix. + # multimodal prefix-LM bidirectional mask. Honor an explicit Flash + # Attention request (for text-only workloads); otherwise fall back + # to Triton Attention, which supports mm_prefix. + if selected_backend == AttentionBackendEnum.FLASH_ATTN: + logger.warning_once( + "Using Flash Attention on XPU for a multimodal prefix-LM " + "model because it was explicitly requested. The prefix-LM " + "bidirectional mask cannot be applied, so image/video " + "inputs will produce incorrect results; only use this for " + "text-only workloads." + ) + return AttentionBackendEnum.FLASH_ATTN.get_path() logger.warning_once( "Flash Attention on XPU does not support multimodal prefix-LM " "attention. Falling back to Triton Attention backend." @@ -265,10 +267,6 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: if compilation_config.compile_sizes is None: compilation_config.compile_sizes = [] - attention_config = vllm_config.attention_config - if attention_config.backend is None: - attention_config.backend = AttentionBackendEnum.FLASH_ATTN - # lazy import to avoid circular import from vllm.utils.torch_utils import supports_xpu_graph diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index ba5ed4718521..84682cdc8bcb 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -128,6 +128,10 @@ "step3p5_reasoning_parser", "Step3p5ReasoningParser", ), + "inkling": ( + "inkling_reasoning_parser", + "InklingParserReasoningAdapter", + ), } diff --git a/vllm/reasoning/inkling_reasoning_parser.py b/vllm/reasoning/inkling_reasoning_parser.py new file mode 100644 index 000000000000..b9246e6c5863 --- /dev/null +++ b/vllm/reasoning/inkling_reasoning_parser.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.engine.registered_adapters import InklingParserReasoningAdapter + +__all__ = ["InklingParserReasoningAdapter"] diff --git a/vllm/renderers/deepseek_v32.py b/vllm/renderers/deepseek_v32.py index 45a46b23283a..ba2e0d60c1a1 100644 --- a/vllm/renderers/deepseek_v32.py +++ b/vllm/renderers/deepseek_v32.py @@ -8,7 +8,6 @@ parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from vllm.tokenizers.deepseek_v32 import DeepseekV32Tokenizer from vllm.utils.async_utils import make_async @@ -17,8 +16,6 @@ from .inputs.preprocess import parse_dec_only_prompt from .params import ChatParams -logger = init_logger(__name__) - class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]): def __init__( diff --git a/vllm/renderers/deepseek_v4.py b/vllm/renderers/deepseek_v4.py index 3dc82b9622e5..e93069209a20 100644 --- a/vllm/renderers/deepseek_v4.py +++ b/vllm/renderers/deepseek_v4.py @@ -8,7 +8,6 @@ parse_chat_messages, parse_chat_messages_async, ) -from vllm.logger import init_logger from vllm.tokenizers.deepseek_v4 import DeepseekV4Tokenizer from vllm.utils.async_utils import make_async @@ -17,8 +16,6 @@ from .inputs.preprocess import parse_dec_only_prompt from .params import ChatParams -logger = init_logger(__name__) - class DeepseekV4Renderer(BaseRenderer[DeepseekV4Tokenizer]): def __init__( diff --git a/vllm/renderers/inkling.py b/vllm/renderers/inkling.py new file mode 100644 index 000000000000..87eedfd69de4 --- /dev/null +++ b/vllm/renderers/inkling.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Native Inkling chat renderer for the Python frontend. + +Mirrors the Rust frontend's native Inkling renderer +(``rust/src/chat/src/renderer/inkling/mod.rs``): chat messages are rendered +directly to token ids — Inkling has no Jinja chat template and no faithful +text form. + +The encoding logic lives in ``inkling_encoding.py`` behind a narrow +"OpenAI messages + tools -> token ids" call; see the swap-point comment +in :meth:`InklingRenderer._render` for adopting a standalone Inkling +input-processing library (mistral-common style) later. +""" + +from vllm.config import VllmConfig +from vllm.entrypoints.chat_utils import ( + ChatCompletionMessageParam, + ConversationMessage, + parse_chat_messages, + parse_chat_messages_async, +) +from vllm.logger import init_logger +from vllm.tokenizers.hf import HfTokenizer +from vllm.utils.async_utils import make_async + +from .base import BaseRenderer +from .inkling_encoding import SPECIAL_TOKEN_SPELLINGS, render_inkling_messages +from .inputs import DictPrompt +from .inputs.preprocess import parse_dec_only_prompt +from .params import ChatParams + +logger = init_logger(__name__) + +_NAMED_REASONING_EFFORT = { + "none": 0.0, + "minimal": 0.1, + "low": 0.2, + "medium": 0.7, + "high": 0.9, + "xhigh": 0.99, + "max": 0.99, +} + +_DEFAULT_REASONING_EFFORT = 0.9 + + +def _resolve_reasoning_effort(value: object) -> float | int | None: + if value is None: + return _DEFAULT_REASONING_EFFORT + if isinstance(value, str): + return _NAMED_REASONING_EFFORT.get(value) + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return value + + +class _HfBackedTmlTokenizer: + """Adapts an HF tokenizer to the encoding core's tokenizer protocol. + + Special-token ids are resolved from the tokenizer vocab at + construction (never hardcoded), trying each known spelling — the Inkling + HF vocab exposes some semantic slots as ``<|unused_NNNNNN|>`` tokens. + """ + + def __init__(self, tokenizer: HfTokenizer) -> None: + self._tokenizer = tokenizer + + vocab = tokenizer.get_vocab() + special_ids: dict[str, int] = {} + missing: list[str] = [] + for token, spellings in SPECIAL_TOKEN_SPELLINGS.items(): + for spelling in spellings: + token_id = vocab.get(spelling) + if token_id is not None: + special_ids[token] = token_id + break + else: + missing.append(token) + if missing: + raise ValueError(f"Inkling tokenizer is missing special tokens: {missing}") + self._special_ids = special_ids + + def encode_text(self, text: str) -> list[int]: + return self._tokenizer.encode(text, add_special_tokens=False) + + def encode_special(self, token: str) -> int: + return self._special_ids[token] + + +class InklingRenderer(BaseRenderer[HfTokenizer]): + def __init__( + self, + config: VllmConfig, + tokenizer: HfTokenizer | None, + ) -> None: + super().__init__(config, tokenizer) + + self._inkling_tokenizer = _HfBackedTmlTokenizer(self.get_tokenizer()) + self._render_async = make_async(self._render, executor=self._executor) + + def _render( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> list[int]: + kwargs = params.chat_template_kwargs or {} + + if kwargs.get("continue_final_message"): + raise ValueError("Inkling renderer does not support continue_final_message") + + reasoning_effort = _resolve_reasoning_effort(kwargs.get("reasoning_effort")) + + try: + # Swap point: to adopt a standalone Inkling input-processing + # library, replace this call (and the _HfBackedTmlTokenizer + # adapter above) with the library's renderer. + return render_inkling_messages( + messages, + self._inkling_tokenizer, + add_generation_prompt=kwargs.get("add_generation_prompt", True), + tools=kwargs.get("tools"), + reasoning_effort=reasoning_effort, + ) + except ValueError: + raise + except (TypeError, KeyError) as e: + # Malformed request content; surface as a request error. + raise ValueError(str(e)) from e + except Exception as e: + logger.exception("Error while rendering Inkling chat messages") + raise ValueError(str(e)) from e + + def render_messages( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = parse_chat_messages( + messages, + self.model_config, + content_format="string", + media_io_kwargs=params.media_io_kwargs, + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + token_ids = self._render(messages, params) + + prompt = parse_dec_only_prompt(token_ids) + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return conversation, prompt + + async def render_messages_async( + self, + messages: list[ChatCompletionMessageParam], + params: ChatParams, + ) -> tuple[list[ConversationMessage], DictPrompt]: + conversation, mm_data, mm_uuids = await parse_chat_messages_async( + messages, + self.model_config, + content_format="string", + media_io_kwargs=params.media_io_kwargs, + mm_processor_kwargs=params.mm_processor_kwargs, + ) + + token_ids = await self._render_async(messages, params) + + prompt = parse_dec_only_prompt(token_ids) + if mm_data is not None: + prompt["multi_modal_data"] = mm_data + if mm_uuids is not None: + prompt["multi_modal_uuids"] = mm_uuids + + return conversation, prompt diff --git a/vllm/renderers/inkling_encoding.py b/vllm/renderers/inkling_encoding.py new file mode 100644 index 000000000000..38039e423289 --- /dev/null +++ b/vllm/renderers/inkling_encoding.py @@ -0,0 +1,383 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inkling chat-encoding core. + +Pure implementation of Inkling chat rendering, kept deliberately free of +vLLM imports: it depends only on the :class:`InklingTextTokenizer` protocol +and speaks OpenAI-style message dicts. If a standalone Inkling +input-processing library becomes available, this module is the unit to +swap out — the swap point is marked in ``vllm/renderers/inkling.py``. + +Multimodal parts follow the contract of vLLM's Inkling multimodal processor +(``InklingMultiModalProcessor`` anchors prompt updates on the bare +content-kind marker and inserts the per-patch placeholder run itself): + +* image parts emit only ``<|content_image|>`` — no seed placeholder id; +* audio parts emit ``<|content_audio_input|><|audio_end|>`` — no seed + placeholder id between them. + +``reasoning_effort`` (a float in [0, 0.99], sourced from +``chat_template_kwargs`` only) renders the ``Thinking effort level:`` +system control block after the tool declarations and initial system messages. +""" + +from __future__ import annotations + +import json +from collections.abc import Iterator, Mapping, Sequence +from typing import Any, Protocol + +END_OF_TEXT = "<|endoftext|>" +MESSAGE_USER = "<|message_user|>" +MESSAGE_MODEL = "<|message_model|>" +MESSAGE_SYSTEM = "<|message_system|>" +MESSAGE_TOOL = "<|message_tool|>" +CONTENT_TEXT = "<|content_text|>" +CONTENT_IMAGE = "<|content_image|>" +CONTENT_MODEL_END_SAMPLING = "<|content_model_end_sampling|>" +CONTENT_THINKING = "<|content_thinking|>" +CONTENT_AUDIO_INPUT = "<|content_audio_input|>" +CONTENT_TOOL_ERROR = "<|content_tool_error|>" +CONTENT_XML = "<|content_xml|>" +CONTENT_INVOKE_TOOL_JSON = "<|content_invoke_tool_json|>" +CONTENT_INVOKE_TOOL_TEXT = "<|content_invoke_tool_text|>" +END_MESSAGE = "<|end_message|>" +AUDIO_END = "<|audio_end|>" + +ROLE_MESSAGE_TOKENS: dict[str, str] = { + "user": MESSAGE_USER, + "assistant": MESSAGE_MODEL, + "system": MESSAGE_SYSTEM, + # The developer role folds into the system role on the Inkling wire. + "developer": MESSAGE_SYSTEM, + "tool": MESSAGE_TOOL, +} + +# Alternate vocab spellings per semantic token. The Inkling HF tokenizer +# exposes some semantic slots as ``<|unused_NNNNNN|>`` tokens (notably +# CONTENT_XML = <|unused_200024|>), so ID resolution must try these +# spellings in order. +SPECIAL_TOKEN_SPELLINGS: dict[str, tuple[str, ...]] = { + MESSAGE_USER: (MESSAGE_USER,), + MESSAGE_MODEL: (MESSAGE_MODEL,), + MESSAGE_SYSTEM: (MESSAGE_SYSTEM,), + MESSAGE_TOOL: (MESSAGE_TOOL,), + CONTENT_TEXT: (CONTENT_TEXT,), + CONTENT_IMAGE: (CONTENT_IMAGE,), + CONTENT_MODEL_END_SAMPLING: (CONTENT_MODEL_END_SAMPLING,), + CONTENT_THINKING: (CONTENT_THINKING,), + CONTENT_AUDIO_INPUT: (CONTENT_AUDIO_INPUT,), + CONTENT_XML: (CONTENT_XML, "<|unused_200024|>"), + CONTENT_INVOKE_TOOL_JSON: (CONTENT_INVOKE_TOOL_JSON,), + END_MESSAGE: (END_MESSAGE,), + AUDIO_END: (AUDIO_END,), +} + + +class InklingTextTokenizer(Protocol): + """Structural tokenizer contract required by the renderer.""" + + def encode_text(self, text: str) -> list[int]: ... + + def encode_special(self, token: str) -> int: ... + + +# OpenAI content-part type spellings that mean image / audio (rendering only +# needs the kind, not the bytes — the bytes are handled by the MM processor). +_IMAGE_PART_TYPES = frozenset({"image", "input_image", "image_url"}) +_AUDIO_PART_TYPES = frozenset({"audio", "input_audio", "audio_url"}) + +_MAX_REASONING_EFFORT = 0.99 + + +def render_inkling_messages( + messages: Sequence[Mapping[str, Any]], + tokenizer: InklingTextTokenizer, + *, + add_generation_prompt: bool = True, + tools: Sequence[Mapping[str, Any]] | None = None, + reasoning_effort: float | None = None, +) -> list[int]: + """Render chat messages to Inkling input ids. + + PURE renderer: emits Inkling framing plus bare media markers; media + encoding and placeholder expansion happen later in the MM processor. + ``add_generation_prompt`` appends the assistant turn opener so the + model continues into the response. + """ + input_ids: list[int] = [] + tool_call_id_to_name: dict[str, str] = {} + + # Request-level tools plus per-developer-message tools (Rust renderer + # semantics) are declared in a single leading system block. + all_tools = list(tools or []) + for message in messages: + if message.get("role") == "developer": + all_tools.extend(message.get("tools") or []) + + if all_tools: + _append_message( + input_ids, + tokenizer, + "system", + "xml", + _tool_declare_json(all_tools), + author_name="tool_declare", + ) + + for message in messages: + role = _expect_role(message) + if role not in {"system", "developer"} and reasoning_effort is not None: + _append_reasoning_effort(input_ids, tokenizer, reasoning_effort) + reasoning_effort = None + + if role == "tool": + tool_name = message.get("name") or tool_call_id_to_name.get( + str(message.get("tool_call_id") or ""), "" + ) + _append_message( + input_ids, + tokenizer, + "tool", + "text", + _flatten_text_content(message.get("content")), + author_name=str(tool_name), + ) + continue + + if role == "assistant": + reasoning_content = message.get("reasoning") + if reasoning_content is None: + reasoning_content = message.get("reasoning_content") + if reasoning_content: + if not isinstance(reasoning_content, str): + raise TypeError( + "assistant reasoning_content must be a string for " + "Inkling rendering" + ) + _append_message( + input_ids, + tokenizer, + "assistant", + "thinking", + reasoning_content, + ) + + for kind, text in _iter_render_parts(message.get("content")): + _append_message(input_ids, tokenizer, role, kind, text) + + if role == "assistant": + for tool_call in message.get("tool_calls") or []: + name, args = _tool_call_name_and_args(tool_call) + tool_call_id = _as_mapping(tool_call).get("id") + if tool_call_id: + tool_call_id_to_name[str(tool_call_id)] = name + _append_message( + input_ids, + tokenizer, + "assistant", + "invoke_tool_json", + _tool_call_json(name, args), + author_name=name, + ) + + input_ids.append(tokenizer.encode_special(CONTENT_MODEL_END_SAMPLING)) + + if reasoning_effort is not None: + _append_reasoning_effort(input_ids, tokenizer, reasoning_effort) + + if add_generation_prompt: + input_ids.append(tokenizer.encode_special(MESSAGE_MODEL)) + return input_ids + + +def _append_reasoning_effort( + input_ids: list[int], + tokenizer: InklingTextTokenizer, + reasoning_effort: float, +) -> None: + _append_message( + input_ids, + tokenizer, + "system", + "text", + _thinking_effort_text(reasoning_effort), + ) + + +def _thinking_effort_text(reasoning_effort: Any) -> str: + if isinstance(reasoning_effort, bool) or not isinstance( + reasoning_effort, (int, float) + ): + raise TypeError( + "Inkling reasoning_effort must be a number in [0.0, 0.99], " + f"got {type(reasoning_effort).__name__}" + ) + value = float(reasoning_effort) + if not 0.0 <= value <= _MAX_REASONING_EFFORT: + raise ValueError( + f"Inkling reasoning_effort must be in [0.0, 0.99], got {value}" + ) + + effort_text = f"{value:.2f}".rstrip("0").rstrip(".") + if effort_text in {"0", "-0"}: + effort_text = "0.0" + return f"Thinking effort level: {effort_text}" + + +def _append_message( + input_ids: list[int], + tokenizer: InklingTextTokenizer, + role: str, + kind: str, + text: str, + *, + author_name: str | None = None, +) -> None: + input_ids.append(tokenizer.encode_special(ROLE_MESSAGE_TOKENS[role])) + if author_name: + input_ids.extend(tokenizer.encode_text(author_name)) + + if kind == "text": + input_ids.append(tokenizer.encode_special(CONTENT_TEXT)) + input_ids.extend(tokenizer.encode_text(text)) + elif kind == "image": + # Bare marker only: the MM processor replaces the marker with + # `marker + placeholder * num_patches` (see mm_preprocess.py). + input_ids.append(tokenizer.encode_special(CONTENT_IMAGE)) + elif kind == "audio": + input_ids.append(tokenizer.encode_special(CONTENT_AUDIO_INPUT)) + input_ids.append(tokenizer.encode_special(AUDIO_END)) + elif kind == "thinking": + input_ids.append(tokenizer.encode_special(CONTENT_THINKING)) + input_ids.extend(tokenizer.encode_text(text)) + elif kind == "xml": + input_ids.append(tokenizer.encode_special(CONTENT_XML)) + input_ids.extend(tokenizer.encode_text(text)) + elif kind == "invoke_tool_json": + input_ids.append(tokenizer.encode_special(CONTENT_INVOKE_TOOL_JSON)) + input_ids.extend(tokenizer.encode_text(text)) + else: + raise ValueError(f"unsupported Inkling render part kind: {kind!r}") + + input_ids.append(tokenizer.encode_special(END_MESSAGE)) + + +def _iter_render_parts(content: Any) -> Iterator[tuple[str, str]]: + """Yield (kind, text) per content part: kind in {text, image, audio}.""" + if content is None: + return + if isinstance(content, str): + if content: + yield ("text", content) + return + if not isinstance(content, Sequence) or isinstance(content, (bytes, bytearray)): + raise TypeError("message content must be a string or a sequence of parts") + for part in content: + if isinstance(part, str): + yield ("text", part) + continue + if not isinstance(part, Mapping): + raise TypeError(f"content part must be mapping, got {type(part).__name__}") + ptype = part.get("type") + if ptype in (None, "text", "input_text"): + text = part.get("text", "") + yield ("text", text if isinstance(text, str) else "") + elif ptype in _IMAGE_PART_TYPES: + yield ("image", "") + elif ptype in _AUDIO_PART_TYPES: + yield ("audio", "") + else: + raise ValueError(f"unsupported content part type: {ptype!r}") + + +def _flatten_text_content(content: Any) -> str: + """Flatten a tool-response content (string or text parts) to text.""" + if content is None: + return "" + if isinstance(content, str): + return content + parts: list[str] = [] + for kind, text in _iter_render_parts(content): + if kind != "text": + raise ValueError( + "Inkling tool response content must be text, " + f"got a part of kind {kind!r}" + ) + parts.append(text) + return "".join(parts) + + +def _expect_role(message: Mapping[str, Any]) -> str: + role = message.get("role") + if role not in ROLE_MESSAGE_TOKENS: + raise ValueError( + f"unsupported Inkling message role {role!r}; " + f"expected one of {sorted(ROLE_MESSAGE_TOKENS)}" + ) + return str(role) + + +def _as_mapping(value: Any) -> Mapping[str, Any]: + if isinstance(value, Mapping): + return value + if hasattr(value, "model_dump"): + dumped = value.model_dump() + if isinstance(dumped, Mapping): + return dumped + raise TypeError(f"expected mapping, got {type(value).__name__}") + + +def _canonical_json(value: Any) -> str: + return json.dumps( + _sort_json(value), + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + ) + + +def _sort_json(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _sort_json(value[key]) for key in sorted(value)} + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return [_sort_json(item) for item in value] + return value + + +def _tool_declare_json(tools: Sequence[Mapping[str, Any]]) -> str: + tool_specs = [] + for tool_value in tools: + tool = _as_mapping(tool_value) + function = _as_mapping(tool.get("function", {})) + tool_specs.append( + { + "description": function.get("description") or "", + "name": function["name"], + "parameters": function.get("parameters") or {}, + "type": tool.get("type", "function"), + } + ) + return _canonical_json(tool_specs) + + +def _tool_call_name_and_args(tool_call_value: Any) -> tuple[str, Mapping[str, Any]]: + tool_call = _as_mapping(tool_call_value) + function = _as_mapping(tool_call.get("function", {})) + name = function.get("name") + if not isinstance(name, str): + raise TypeError("tool call function name must be a string") + + raw_args = function.get("arguments") or {} + if isinstance(raw_args, str): + args = json.loads(raw_args) if raw_args.strip() else {} + else: + args = raw_args + if not isinstance(args, Mapping): + raise TypeError("tool call function arguments must decode to an object") + return name, args + + +def _tool_call_json(name: str, args: Mapping[str, Any]) -> str: + name_json = json.dumps(name, ensure_ascii=False, allow_nan=False) + return f'{{"name":{name_json},"args":{_canonical_json(args)}}}' diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 20c54eb07ffd..3ce4f74d9d82 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -7,6 +7,7 @@ from vllm.entrypoints.generate.base.serving import resolve_token_id_placeholder from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionLogProbs, + ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionResponseChoice, ChatMessage, @@ -136,6 +137,13 @@ async def derender_chat( else [] ) + is_named_tool_choice = ( + type(chat_request.tool_choice) is ChatCompletionNamedToolChoiceParam + ) + is_required_tool_choice = chat_request.tool_choice == "required" + if is_named_tool_choice or is_required_tool_choice: + content = content or "" + message = ChatMessage( role="assistant", reasoning=reasoning, diff --git a/vllm/renderers/registry.py b/vllm/renderers/registry.py index 098a58e8edcd..395132c56e15 100644 --- a/vllm/renderers/registry.py +++ b/vllm/renderers/registry.py @@ -26,6 +26,7 @@ "kimi_audio": ("hf", "HfRenderer"), "mistral": ("mistral", "MistralRenderer"), "terratorch": ("terratorch", "TerratorchRenderer"), + "inkling": ("inkling", "InklingRenderer"), } diff --git a/vllm/third_party/flash_linear_attention/LICENSE b/vllm/third_party/flash_linear_attention/LICENSE new file mode 100644 index 000000000000..559c4d95cdff --- /dev/null +++ b/vllm/third_party/flash_linear_attention/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/vllm/model_executor/layers/fla/__init__.py b/vllm/third_party/flash_linear_attention/__init__.py similarity index 100% rename from vllm/model_executor/layers/fla/__init__.py rename to vllm/third_party/flash_linear_attention/__init__.py diff --git a/vllm/model_executor/layers/fla/ops/__init__.py b/vllm/third_party/flash_linear_attention/ops/__init__.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/__init__.py rename to vllm/third_party/flash_linear_attention/ops/__init__.py diff --git a/vllm/model_executor/layers/fla/ops/chunk.py b/vllm/third_party/flash_linear_attention/ops/chunk.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/chunk.py rename to vllm/third_party/flash_linear_attention/ops/chunk.py diff --git a/vllm/model_executor/layers/fla/ops/chunk_delta_h.py b/vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/chunk_delta_h.py rename to vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py diff --git a/vllm/model_executor/layers/fla/ops/chunk_o.py b/vllm/third_party/flash_linear_attention/ops/chunk_o.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/chunk_o.py rename to vllm/third_party/flash_linear_attention/ops/chunk_o.py diff --git a/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py b/vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py rename to vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py diff --git a/vllm/model_executor/layers/fla/ops/cumsum.py b/vllm/third_party/flash_linear_attention/ops/cumsum.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/cumsum.py rename to vllm/third_party/flash_linear_attention/ops/cumsum.py diff --git a/vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py b/vllm/third_party/flash_linear_attention/ops/fused_gdn_prefill_post_conv.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py rename to vllm/third_party/flash_linear_attention/ops/fused_gdn_prefill_post_conv.py diff --git a/vllm/model_executor/layers/fla/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/fused_recurrent.py rename to vllm/third_party/flash_linear_attention/ops/fused_recurrent.py diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py rename to vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py diff --git a/vllm/model_executor/layers/fla/ops/index.py b/vllm/third_party/flash_linear_attention/ops/index.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/index.py rename to vllm/third_party/flash_linear_attention/ops/index.py diff --git a/vllm/model_executor/layers/fla/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/kda.py rename to vllm/third_party/flash_linear_attention/ops/kda.py diff --git a/vllm/model_executor/layers/fla/ops/l2norm.py b/vllm/third_party/flash_linear_attention/ops/l2norm.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/l2norm.py rename to vllm/third_party/flash_linear_attention/ops/l2norm.py diff --git a/vllm/model_executor/layers/fla/ops/layernorm_guard.py b/vllm/third_party/flash_linear_attention/ops/layernorm_guard.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/layernorm_guard.py rename to vllm/third_party/flash_linear_attention/ops/layernorm_guard.py diff --git a/vllm/model_executor/layers/fla/ops/op.py b/vllm/third_party/flash_linear_attention/ops/op.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/op.py rename to vllm/third_party/flash_linear_attention/ops/op.py diff --git a/vllm/model_executor/layers/fla/ops/solve_tril.py b/vllm/third_party/flash_linear_attention/ops/solve_tril.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/solve_tril.py rename to vllm/third_party/flash_linear_attention/ops/solve_tril.py diff --git a/vllm/model_executor/layers/fla/ops/utils.py b/vllm/third_party/flash_linear_attention/ops/utils.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/utils.py rename to vllm/third_party/flash_linear_attention/ops/utils.py diff --git a/vllm/model_executor/layers/fla/ops/wy_fast.py b/vllm/third_party/flash_linear_attention/ops/wy_fast.py similarity index 100% rename from vllm/model_executor/layers/fla/ops/wy_fast.py rename to vllm/third_party/flash_linear_attention/ops/wy_fast.py diff --git a/vllm/tokenizers/deepseek_v32_encoding.py b/vllm/tokenizers/deepseek_v32_encoding.py index 249b5326275e..b02449020e77 100644 --- a/vllm/tokenizers/deepseek_v32_encoding.py +++ b/vllm/tokenizers/deepseek_v32_encoding.py @@ -7,8 +7,6 @@ import json from typing import Any -import regex as re - # flake8: noqa: E501 TOOLS_SYSTEM_TEMPLATE = """## Tools You have access to a set of tools you can use to answer the user's question. @@ -79,19 +77,6 @@ def tool_calls_from_openai_format(tool_calls): ] -def tool_calls_to_openai_format(tool_calls): - return [ - { - "type": "function", - "function": { - "name": tool_call["name"], - "arguments": tool_call["arguments"], - }, - } - for tool_call in tool_calls - ] - - def encode_arguments_to_dsml(tool_call: dict[str, str]) -> str: p_dsml_template = """<{dsml_token}parameter name="{key}" string="{is_str}">{value}""" P_dsml_strs = [] @@ -113,24 +98,6 @@ def encode_arguments_to_dsml(tool_call: dict[str, str]) -> str: return "\n".join(P_dsml_strs) -def decode_dsml_to_arguments( - tool_name: str, tool_args: dict[str, tuple[str, str]] -) -> dict[str, str]: - def _decode_value(key: str, value: str, string: str): - if string == "true": - value = to_json(value) - return f"{to_json(key)}: {value}" - - tool_args_json = ( - "{" - + ", ".join( - [_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()] - ) - + "}" - ) - return dict(name=tool_name, arguments=tool_args_json) - - def render_tools(tools: list[dict[str, str | dict[str, Any]]]) -> str: tools_json = [to_json(t) for t in tools] @@ -333,139 +300,3 @@ def encode_messages( ) return prompt - - -def _read_until_stop( - index: int, text: str, stop: list[str] -) -> tuple[int, str, None | str]: - min_pos = len(text) - matched_stop = None - - for s in stop: - pos = text.find(s, index) - if pos != -1 and pos < min_pos: - min_pos = pos - matched_stop = s - - if matched_stop: - content = text[index:min_pos] - return min_pos + len(matched_stop), content, matched_stop - else: - content = text[index:] - return len(text), content, None - - -def parse_tool_calls(index: int, text: str): - tool_calls: list[dict[str, Any]] = [] - stop_token = None - tool_calls_end_token = f"" - - while index < len(text): - index, _, stop_token = _read_until_stop( - index, text, [f"<{dsml_token}invoke", tool_calls_end_token] - ) - if _ != ">\n": - raise RuntimeError("Tool call format error") - - if stop_token == tool_calls_end_token: - break - - if stop_token is None: - raise RuntimeError("Missing special token") - - index, tool_name_content, stop_token = _read_until_stop( - index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL - ) - if len(p_tool_name) != 1: - raise RuntimeError("Tool name format error") - tool_name = p_tool_name[0] - - tool_args: dict[str, tuple[str, str]] = {} - while stop_token == f"<{dsml_token}parameter": - index, param_content, stop_token = _read_until_stop( - index, text, [f"/{dsml_token}parameter"] - ) - - param_kv = re.findall( - r'^ name="(.*?)" string="(true|false)">(.*?)<$', - param_content, - flags=re.DOTALL, - ) - if len(param_kv) != 1: - raise RuntimeError("Parameter format error") - param_name, string, param_value = param_kv[0] - - if param_name in tool_args: - raise RuntimeError("Duplicate parameter name") - tool_args[param_name] = (param_value, string) - - index, content, stop_token = _read_until_stop( - index, text, [f"<{dsml_token}parameter", f"\n": - raise RuntimeError("Parameter format error") - - tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args) - tool_calls.append(tool_call) - - return index, stop_token, tool_calls - - -# NOTE: This function is designed to parse only correctly -# formatted string and will not attempt to correct malformed output -# that may be generated by the model. -def parse_message_from_completion_text(text: str, thinking_mode: str): - summary_content, reasoning, tool_calls = "", "", [] - index, stop_token = 0, None - tool_calls_start_token = f"\n\n<{dsml_token}function_calls" - - is_thinking, is_tool_calling = thinking_mode == "thinking", False - - if is_thinking: - index, content_delta, stop_token = _read_until_stop( - index, text, [thinking_end_token, tool_calls_start_token] - ) - reasoning = content_delta - if stop_token != thinking_end_token: - raise RuntimeError("Invalid thinking format") - - index, content_delta, stop_token = _read_until_stop( - index, text, [eos_token, tool_calls_start_token] - ) - summary_content = content_delta - if stop_token == tool_calls_start_token: - is_tool_calling = True - else: - if stop_token != eos_token: - raise RuntimeError("Invalid summary format") - - if is_tool_calling: - index, stop_token, tool_calls = parse_tool_calls(index, text) - - index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token]) - if tool_ends_text: - raise RuntimeError("Unexpected content after tool calls") - - if not (len(text) == index and stop_token in [eos_token, None]): - raise RuntimeError("Unexpected content at end") - - for sp_token in [ - bos_token, - eos_token, - thinking_start_token, - thinking_end_token, - dsml_token, - ]: - if sp_token in summary_content or sp_token in reasoning: - raise RuntimeError("Unexpected special token in content") - - return { - "role": "assistant", - "content": summary_content, - "reasoning": reasoning, - "tool_calls": tool_calls_to_openai_format(tool_calls), - } diff --git a/vllm/tokenizers/deepseek_v4_encoding.py b/vllm/tokenizers/deepseek_v4_encoding.py index 6895771e2f59..16bfa1a99a1b 100644 --- a/vllm/tokenizers/deepseek_v4_encoding.py +++ b/vllm/tokenizers/deepseek_v4_encoding.py @@ -6,16 +6,14 @@ """ DeepSeek-V4 Encoding -A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages -with tool calling, thinking mode, and quick instruction task support. +A self-contained implementation for encoding DeepSeek-V4 chat messages with tool +calling, thinking mode, and quick instruction task support. """ -from typing import Any, Dict, List, Union, Optional, Tuple +from typing import Any, Dict, List, Union, Optional import copy import json -import regex as re - # ============================================================ # Special Tokens # ============================================================ @@ -128,20 +126,6 @@ def tool_calls_from_openai_format(tool_calls): ] -def tool_calls_to_openai_format(tool_calls): - """Convert internal tool calls to OpenAI format.""" - return [ - { - "type": "function", - "function": { - "name": tool_call["name"], - "arguments": tool_call["arguments"], - } - } - for tool_call in tool_calls - ] - - def encode_arguments_to_dsml(tool_call: Dict[str, Any]) -> str: """ Encode tool call arguments into DSML parameter format. @@ -172,26 +156,6 @@ def encode_arguments_to_dsml(tool_call: Dict[str, Any]) -> str: return "\n".join(P_dsml_strs) -def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]: - """ - Decode DSML parameters back to a tool call dict. - - Args: - tool_name: Name of the tool. - tool_args: Dict mapping param_name -> (value, is_string_flag). - - Returns: - Dict with "name" and "arguments" (JSON string) keys. - """ - def _decode_value(key: str, value: str, string: str): - if string == "true": - value = to_json(value) - return f"{to_json(key)}: {value}" - - tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}" - return dict(name=tool_name, arguments=tool_args_json) - - def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str: """ Render tool schemas into the system prompt format. @@ -605,153 +569,4 @@ def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, An return result -# ============================================================ -# Parsing (Decoding model output) -# ============================================================ - -def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]: - """ - Read text from index until one of the stop strings is found. - - Returns: - Tuple of (new_index, content_before_stop, matched_stop_string_or_None). - """ - min_pos = len(text) - matched_stop = None - - for s in stop: - pos = text.find(s, index) - if pos != -1 and pos < min_pos: - min_pos = pos - matched_stop = s - - if matched_stop: - content = text[index:min_pos] - return min_pos + len(matched_stop), content, matched_stop - else: - content = text[index:] - return len(text), content, None - - -def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]: - """ - Parse DSML tool calls from text starting at the given index. - - Args: - index: Starting position in text. - text: The full text to parse. - - Returns: - Tuple of (new_index, last_stop_token, list_of_tool_call_dicts). - Each tool call dict has "name" and "arguments" keys. - """ - tool_calls: List[Dict[str, Any]] = [] - stop_token = None - tool_calls_end_token = f"" - - while index < len(text): - index, content_before, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token]) - if content_before != ">\n": - raise ValueError(f"Tool call format error: expected '>\\n' but got '{content_before}'") - - if stop_token == tool_calls_end_token: - break - - if stop_token is None: - raise ValueError("Missing special token in tool calls") - - index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL) - if len(p_tool_name) != 1: - raise ValueError(f"Tool name format error: '{tool_name_content}'") - tool_name = p_tool_name[0] - - tool_args: Dict[str, Tuple[str, str]] = {} - while stop_token == f"<{dsml_token}parameter": - index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"]) - - param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL) - if len(param_kv) != 1: - raise ValueError(f"Parameter format error: '{param_content}'") - param_name, string, param_value = param_kv[0] - - if param_name in tool_args: - raise ValueError(f"Duplicate parameter name: '{param_name}'") - tool_args[param_name] = (param_value, string) - - index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n": - raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'") - - tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args) - tool_calls.append(tool_call) - - return index, stop_token, tool_calls - - -def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]: - """ - Parse a model completion text into a structured assistant message. - - This function takes the raw text output from the model (a single assistant turn) - and extracts: - - reasoning (thinking block) - - content (summary/response) - - tool_calls (if any) - - NOTE: This function is designed to parse only correctly formatted strings and - will raise ValueError for malformed output. - - Args: - text: The raw completion text (including EOS token). - thinking_mode: Either "chat" or "thinking". - - Returns: - Dict with keys: "role", "content", "reasoning", "tool_calls". - tool_calls are in OpenAI format. - """ - summary_content, reasoning = "", "" - tool_calls: List[Dict[str, str]] = [] - index, stop_token = 0, None - tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}" - - is_thinking = thinking_mode == "thinking" - is_tool_calling = False - - if is_thinking: - index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token]) - reasoning = content_delta - if stop_token != thinking_end_token: - raise ValueError("Invalid thinking format: missing ") - - index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token]) - summary_content = content_delta - if stop_token == tool_calls_start_token: - is_tool_calling = True - else: - if stop_token != eos_token: - raise ValueError("Invalid format: missing EOS token") - - if is_tool_calling: - index, stop_token, tool_calls = parse_tool_calls(index, text) - - index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token]) - if tool_ends_text: - raise ValueError("Unexpected content after tool calls") - - if len(text) != index or stop_token not in [eos_token, None]: - raise ValueError("Unexpected content at end") - - for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]: - if sp_token in summary_content or sp_token in reasoning: - raise ValueError(f"Unexpected special token '{sp_token}' in content") - - return { - "role": "assistant", - "content": summary_content, - "reasoning": reasoning, - "tool_calls": tool_calls_to_openai_format(tool_calls) - } - # fmt: on diff --git a/vllm/tokenizers/registry.py b/vllm/tokenizers/registry.py index cef2f7645fe8..e6c12ccc3bc7 100644 --- a/vllm/tokenizers/registry.py +++ b/vllm/tokenizers/registry.py @@ -44,6 +44,10 @@ "hf": ("hf", "CachedHfTokenizer"), "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"), "mistral": ("mistral", "MistralTokenizer"), + # Inkling uses the plain HF tokenizer for token operations; the "inkling" + # mode exists to select the InklingRenderer, which renders chat to + # token ids natively (Inkling has no Jinja chat template). + "inkling": ("hf", "CachedHfTokenizer"), } diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 26362ebf0ed6..64b1342fd253 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -174,6 +174,10 @@ "step3p5_tool_parser", "Step3p5ToolParser", ), + "inkling": ( + "inkling_tool_parser", + "InklingEngineToolParser", + ), "xlam": ( "xlam_tool_parser", "xLAMToolParser", diff --git a/vllm/tool_parsers/inkling_tool_parser.py b/vllm/tool_parsers/inkling_tool_parser.py new file mode 100644 index 000000000000..ab760015af73 --- /dev/null +++ b/vllm/tool_parsers/inkling_tool_parser.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.engine.registered_adapters import InklingParserToolAdapter + + +class InklingEngineToolParser(InklingParserToolAdapter): # type: ignore[valid-type, misc] + # No Inkling structural-tag grammar is wired up yet; fall back to auto + # parsing for named/required tool choice. + structural_tag_model = None + supports_required_and_named = False diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index bb72087a95d4..b457a6e2f029 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -127,6 +127,8 @@ def __getitem__(self, key): laguna="LagunaConfig", lfm2_moe="Lfm2MoeConfig", **{"unlimited-ocr": "UnlimitedOCRConfig"}, + inkling_mm_model="InklingMMConfig", + inkling_model="InklingModelConfig", ) _SPECULATIVE_DECODING_CONFIGS: set[str] = {"eagle", "speculators", "medusa"} diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index b64310507b7a..4bb7674ddbb4 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -94,6 +94,10 @@ "Qwen3_5TextConfig": "vllm.transformers_utils.configs.qwen3_5", "Qwen3_5MoeConfig": "vllm.transformers_utils.configs.qwen3_5_moe", "Qwen3_5MoeTextConfig": "vllm.transformers_utils.configs.qwen3_5_moe", + "InklingModelConfig": "vllm.models.inkling.configs", + "InklingAudioConfig": "vllm.models.inkling.configs", + "InklingVisionConfig": "vllm.models.inkling.configs", + "InklingMMConfig": "vllm.models.inkling.configs", # Special case: DeepseekV3Config is from HuggingFace Transformers "DeepseekV3Config": "transformers", } @@ -174,6 +178,10 @@ "Qwen3_5TextConfig", "Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig", + "InklingModelConfig", + "InklingAudioConfig", + "InklingVisionConfig", + "InklingMMConfig", ] diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index e35218c4be3a..60b56b75a00c 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -44,6 +44,9 @@ "Ovis2_5Processor", "Qwen3ASRProcessor", "Step3VLProcessor", + "InklingProcessor", + "InklingImageProcessor", + "InklingAudioFeatureExtractor", ] _CLASS_TO_MODULE: dict[str, str] = { @@ -80,6 +83,9 @@ "Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5", "Qwen3ASRProcessor": "vllm.transformers_utils.processors.qwen3_asr", "Step3VLProcessor": "vllm.transformers_utils.processors.step3_vl", + "InklingProcessor": "vllm.transformers_utils.processors.inkling", + "InklingImageProcessor": "vllm.transformers_utils.processors.inkling", + "InklingAudioFeatureExtractor": "vllm.transformers_utils.processors.inkling", } diff --git a/vllm/transformers_utils/processors/inkling.py b/vllm/transformers_utils/processors/inkling.py new file mode 100644 index 000000000000..09d34f7806cb --- /dev/null +++ b/vllm/transformers_utils/processors/inkling.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Vendored HuggingFace-convention processors for the Inkling Titan model. + +Implements the image processor (numba patchifier), the audio feature extractor +(STFT/dMel path), the composite processor, and the MM token-id constants. + +Besides raw bytes / file paths, the extractors also accept the dummy inputs +vLLM generates during profiling (PIL images / numpy audio arrays). +""" + +from __future__ import annotations + +import io +import math +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn.functional as F +from numba import njit +from transformers.feature_extraction_utils import ( + BatchFeature, + FeatureExtractionMixin, +) +from transformers.image_processing_utils import BaseImageProcessor +from transformers.image_utils import ImageInput + +# --------------------------------------------------------------------------- +# MM token-id constants +# --------------------------------------------------------------------------- + +# Block-start marker token ids. These are real tokens the model was trained on +# (``<|content_image|>`` / ``<|content_audio_input|>``) that mark the start of an +# image/audio embedding block; they are kept verbatim in ``input_ids``. +IMAGE_MARKER_ID = 200005 # <|content_image|> +AUDIO_MARKER_ID = 200020 # <|content_audio_input|> + +# Per-patch / per-frame placeholder ids marking where tower embeddings are +# scattered in. These are unused slots in the padded vocabulary and the +# corresponding positions are always overwritten by tower embeddings. +IMAGE_TOKEN_ID = 200054 # <|unused_200054|> +AUDIO_TOKEN_ID = 200053 # <|unused_200053|> + +# --------------------------------------------------------------------------- +# Image processing +# --------------------------------------------------------------------------- + +IMAGE_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32) +IMAGE_STD = np.array([0.26862954, 0.2613026, 0.2757771], dtype=np.float32) +PAD_RAW_VALUE = np.float32(-1.0 / 255.0) +PAD_NORM = (np.full((3,), PAD_RAW_VALUE, dtype=np.float32) - IMAGE_MEAN) / IMAGE_STD + + +def _validate_image_rescale( + rescale_image_frac: float | None, + rescale_image_max_upscaled_long_edge: int | None, +) -> None: + if rescale_image_frac is not None and ( + not math.isfinite(rescale_image_frac) or rescale_image_frac <= 0 + ): + raise ValueError( + "rescale_image_frac must be positive and finite or None, " + f"got {rescale_image_frac}" + ) + if rescale_image_max_upscaled_long_edge is None: + return + if rescale_image_max_upscaled_long_edge <= 0: + raise ValueError( + "rescale_image_max_upscaled_long_edge must be positive or None, " + f"got {rescale_image_max_upscaled_long_edge}" + ) + if rescale_image_frac is None or rescale_image_frac <= 1.0: + raise ValueError( + "rescale_image_max_upscaled_long_edge requires rescale_image_frac > 1, " + f"got {rescale_image_frac}" + ) + + +def _scaled_image_dimensions( + width: int, + height: int, + rescale_image_frac: float | None, + rescale_image_max_upscaled_long_edge: int | None, +) -> tuple[int, int]: + """Return the long-edge-scaled ``(width, height)``.""" + if rescale_image_frac is None: + return width, height + + long_edge = max(width, height) + if long_edge == 0: + return width, height + + target_long_edge = float(long_edge) * rescale_image_frac + if rescale_image_max_upscaled_long_edge is not None: + effective_cap = max(rescale_image_max_upscaled_long_edge, long_edge) + target_long_edge = min(target_long_edge, float(effective_cap)) + + ratio = target_long_edge / float(long_edge) + if ratio == 1.0: + return width, height + + def scale(value: int) -> int: + return max(1, math.floor(float(value) * ratio + 0.5)) + + return scale(width), scale(height) + + +def _load_image_bytes(image) -> bytes: + """Encode a PIL image as raw PNG bytes for preprocessing. + + The HF processor is always handed ``PIL.Image`` instances by vLLM, so no + other input types need to be supported here. + """ + if image.mode != "RGB": + image = image.convert("RGB") + buf = io.BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + +@njit(cache=True) +def _fill_patches_numba( + arr: np.ndarray, + patch_size: int, + patches: np.ndarray, + mean: np.ndarray, + std: np.ndarray, + pad_norm: np.ndarray, +) -> None: + h = arr.shape[0] + w = arr.shape[1] + nph = (h + patch_size - 1) // patch_size + npw = w // patch_size + 1 + inv255 = np.float32(1.0 / 255.0) + + for k in range(nph * npw): + i = k // npw + j = k - i * npw + y_base = i * patch_size + x_base = j * patch_size + + for y in range(patch_size): + iy = y_base + y + for x in range(patch_size): + ix = x_base + x + if iy < h and ix < w: + for c in range(3): + raw = np.float32(arr[iy, ix, c]) * inv255 + patches[k, y, x, c] = (raw - mean[c]) / std[c] + else: + for c in range(3): + patches[k, y, x, c] = pad_norm[c] + + +def _encode_image_bytes( + image_bytes: bytes, + *, + patch_size: int, + rescale_image_frac: float | None, + rescale_image_max_upscaled_long_edge: int | None, +) -> torch.Tensor: + if patch_size <= 0: + raise ValueError("patch_size must be greater than zero") + _validate_image_rescale( + rescale_image_frac, + rescale_image_max_upscaled_long_edge, + ) + + from PIL import Image + + image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + scaled_size = _scaled_image_dimensions( + image.width, + image.height, + rescale_image_frac=rescale_image_frac, + rescale_image_max_upscaled_long_edge=rescale_image_max_upscaled_long_edge, + ) + if scaled_size != image.size: + image = image.resize(scaled_size, resample=Image.Resampling.LANCZOS) + arr = np.array(image, dtype=np.uint8, copy=True) + height, width, _ = arr.shape + + nph = (height + patch_size - 1) // patch_size + npw = width // patch_size + 1 + num_patches = nph * npw + + patches = np.empty((num_patches, patch_size, patch_size, 3), dtype=np.float32) + _fill_patches_numba(arr, patch_size, patches, IMAGE_MEAN, IMAGE_STD, PAD_NORM) + + return ( + torch.from_numpy(patches) + .to(torch.bfloat16) + .view(num_patches, 1, patch_size, patch_size, 3) + .expand(num_patches, 2, patch_size, patch_size, 3) + ) + + +class InklingImageProcessor(BaseImageProcessor): + r"""Turn raw images into ``vision_patches_bthwc`` for Inkling hMLP. + + ``rescale_image_frac`` scales the long edge while preserving aspect ratio. + ``rescale_image_max_upscaled_long_edge`` optionally caps only upscaling and + therefore requires a scale factor greater than one. The defaults, ``2.0`` and + ``2048``, grow images toward a 2048-pixel long edge by at most 2x, while leaving + images already at or above 2048 unchanged. + """ + + model_input_names = ["vision_patches_bthwc"] + + def __init__( + self, + patch_size: int = 40, + rescale_image_frac: float | None = 2.0, + rescale_image_max_upscaled_long_edge: int | None = 2048, + **kwargs, + ): + if patch_size <= 0: + raise ValueError("patch_size must be greater than zero") + _validate_image_rescale( + rescale_image_frac, + rescale_image_max_upscaled_long_edge, + ) + super().__init__(**kwargs) + self.patch_size = patch_size + self.rescale_image_frac = rescale_image_frac + self.rescale_image_max_upscaled_long_edge = rescale_image_max_upscaled_long_edge + + def _encode_one(self, image) -> torch.Tensor: + return _encode_image_bytes( + _load_image_bytes(image), + patch_size=self.patch_size, + rescale_image_frac=self.rescale_image_frac, + rescale_image_max_upscaled_long_edge=self.rescale_image_max_upscaled_long_edge, + ) + + def preprocess( + self, + images: ImageInput | list, + return_tensors: str | None = "pt", + **kwargs, + ) -> BatchFeature: + del return_tensors, kwargs + if not isinstance(images, (list, tuple)): + images = [images] + + per_image_patches: list[torch.Tensor] = [] + num_patches: list[int] = [] + num_tokens: list[int] = [] + for img in images: + vp = self._encode_one(img) + n_patches = int(vp.shape[0]) + per_image_patches.append(vp) + num_patches.append(n_patches) + num_tokens.append(n_patches) + + if len(per_image_patches) == 1: + vision_patches_bthwc = per_image_patches[0] + elif per_image_patches: + vision_patches_bthwc = torch.cat(per_image_patches, dim=0) + else: + vision_patches_bthwc = torch.empty(0) + + data = { + "vision_patches_bthwc": vision_patches_bthwc, + "num_patches": num_patches, + "num_tokens": num_tokens, + } + return BatchFeature(data=data, tensor_type=None) + + +# --------------------------------------------------------------------------- +# Audio feature extraction +# --------------------------------------------------------------------------- + + +@dataclass +class InklingAudioEncoderParams: + """Audio preprocessing parameters used to convert raw audio into dMel bins.""" + + sample_rate: int = 16_000 + window_size_multiplier: float = 2.0 + n_fft: int | None = None + n_mels: int = 80 + num_dmel_bins: int = 16 + dmel_min_value: float = -7.0 + dmel_max_value: float = 2.0 + audio_token_duration_s: float = 0.05 + + +def _to_exact_int(value: float, name: str, tolerance: float = 1e-6) -> int: + rounded = round(value) + if abs(value - rounded) > tolerance: + raise ValueError(f"{name} must resolve to an integer sample count, got {value}") + return int(rounded) + + +def _hz_to_mel(frequencies: np.ndarray) -> np.ndarray: + """Slaney mel scale, matching the librosa/torchaudio convention.""" + frequencies = np.asarray(frequencies, dtype=np.float64) + f_sp = 200.0 / 3.0 + min_log_hz = 1000.0 + min_log_mel = min_log_hz / f_sp + logstep = np.log(6.4) / 27.0 + linear = frequencies / f_sp + log = ( + min_log_mel + np.log(np.maximum(frequencies, min_log_hz) / min_log_hz) / logstep + ) + return np.where(frequencies >= min_log_hz, log, linear) + + +def _mel_to_hz(mels: np.ndarray) -> np.ndarray: + mels = np.asarray(mels, dtype=np.float64) + f_sp = 200.0 / 3.0 + min_log_hz = 1000.0 + min_log_mel = min_log_hz / f_sp + logstep = np.log(6.4) / 27.0 + linear = mels * f_sp + log = min_log_hz * np.exp(logstep * (mels - min_log_mel)) + return np.where(mels >= min_log_mel, log, linear) + + +_MEL_BASIS_CACHE: dict[tuple[int, int, int], torch.Tensor] = {} + + +def _mel_basis(sample_rate: int, n_fft: int, n_mels: int) -> torch.Tensor: + key = (sample_rate, n_fft, n_mels) + cached = _MEL_BASIS_CACHE.get(key) + if cached is not None: + return cached + + fft_bins = n_fft // 2 + 1 + fft_freqs = np.arange(fft_bins, dtype=np.float64) * sample_rate / n_fft + mel_edges = _mel_to_hz( + np.linspace( + _hz_to_mel(np.array([0.0]))[0], + _hz_to_mel(np.array([sample_rate / 2.0]))[0], + n_mels + 2, + dtype=np.float64, + ) + ) + mel_widths = np.diff(mel_edges) + lower = (fft_freqs[None, :] - mel_edges[:-2, None]) / mel_widths[:-1, None] + upper = (mel_edges[2:, None] - fft_freqs[None, :]) / mel_widths[1:, None] + weights = np.maximum(0.0, np.minimum(lower, upper)) + + # Slaney area normalization. + weights *= (2.0 / (mel_edges[2:] - mel_edges[:-2]))[:, None] + basis = torch.from_numpy(weights.astype(np.float32, copy=False)).contiguous() + _MEL_BASIS_CACHE[key] = basis + return basis + + +def _dmel_bins(audio: torch.Tensor, params: InklingAudioEncoderParams) -> torch.Tensor: + hop_length = _to_exact_int( + params.audio_token_duration_s * params.sample_rate, + "audio_token_duration_s * sample_rate", + ) + window_size = _to_exact_int( + params.audio_token_duration_s + * params.window_size_multiplier + * params.sample_rate, + "audio_token_duration_s * window_size_multiplier * sample_rate", + ) + n_fft = params.n_fft or window_size + if hop_length <= 0 or window_size <= 0 or n_fft <= 0: + raise ValueError("audio hop length, window size, and n_fft must be positive") + if audio.numel() == 0: + return torch.empty((0, params.n_mels), dtype=torch.int32) + + right_pad = math.ceil(audio.numel() / hop_length) * hop_length - audio.numel() + left_pad = max(n_fft - hop_length, 0) + audio = F.pad(audio, (left_pad, right_pad)) + + window = torch.hann_window(window_size, periodic=True, dtype=torch.float32) + spec = torch.stft( + audio.unsqueeze(0), + n_fft=n_fft, + hop_length=hop_length, + win_length=window_size, + window=window, + center=False, + normalized=False, + onesided=True, + return_complex=True, + ) + spec_ri = torch.view_as_real(spec) + magnitude = ( + (spec_ri[..., 0].square() + spec_ri[..., 1].square()) + .clamp_min(1e-10) + .sqrt() + .squeeze(0) + ) + + mel = ( + _mel_basis(params.sample_rate, n_fft, params.n_mels) + .matmul(magnitude) + .clamp_min(1e-10) + .log10() + ) + mel = mel.to(torch.float64).clamp( + min=params.dmel_min_value, max=params.dmel_max_value + ) + bin_centers = torch.linspace( + params.dmel_min_value, + params.dmel_max_value, + params.num_dmel_bins, + dtype=torch.float64, + ) + dmel_bins = (mel.unsqueeze(-1) - bin_centers).abs().argmin(dim=-1) + return dmel_bins.to(torch.int32).T.contiguous() + + +class InklingAudioFeatureExtractor(FeatureExtractionMixin): + """Convert raw audio into Inkling dMel bins in the HF feature-extractor API.""" + + model_input_names = ["dmel_bins"] + + def __init__(self, params: dict | None = None, **kwargs): + super().__init__(**kwargs) + merged = InklingAudioEncoderParams() + if params: + for k, v in params.items(): + if hasattr(merged, k): + setattr(merged, k, v) + # also accept flat kwargs (HF config style) + for k in list(kwargs.keys()): + if hasattr(merged, k): + setattr(merged, k, kwargs[k]) + self.params = merged + + def _decode_one(self, audio) -> torch.Tensor: + # vLLM hands the feature extractor numpy arrays (the dummy-input builder + # during profiling, and MultiModalDataParser after resampling to the + # target sample rate), so no other input types need to be supported. + return torch.from_numpy( + np.ascontiguousarray(audio.astype(np.float32, copy=False)) + ).flatten() + + def _encode_one(self, audio) -> torch.Tensor: + return _dmel_bins(self._decode_one(audio), self.params) + + def __call__( + self, + audios: Sequence | None, + return_tensors: str | None = None, + **kwargs, + ) -> BatchFeature: + del return_tensors, kwargs + if audios is None: + audios = [] + if not isinstance(audios, (list, tuple)): + audios = [audios] + + dmel_bins = [self._encode_one(a) for a in audios] + data = { + # per-clip feature: dmel bins as float32 [T, n_mels] + "dmel_bins": [bins.to(torch.float32) for bins in dmel_bins], + "num_audio_tokens": [int(bins.shape[0]) for bins in dmel_bins], + } + # return_tensors intentionally ignored: per-clip features have ragged T. + return BatchFeature(data=data, tensor_type=None) + + +# --------------------------------------------------------------------------- +# Composite processor +# --------------------------------------------------------------------------- + + +class InklingProcessor: + """Bundle Inkling image + audio preprocessing with the MM token ids.""" + + def __init__( + self, + image_processor: InklingImageProcessor | None = None, + audio_feature_extractor: InklingAudioFeatureExtractor | None = None, + tokenizer=None, + ): + self.image_processor = image_processor or InklingImageProcessor() + self.audio_feature_extractor = ( + audio_feature_extractor or InklingAudioFeatureExtractor() + ) + self.tokenizer = tokenizer + + def process_images(self, images: list): + """Raw images -> BatchFeature(vision_patches_bthwc, num_patches, num_tokens).""" + return self.image_processor.preprocess(images, return_tensors="pt") + + def process_audios(self, audios: list): + """Raw audios -> BatchFeature(dmel_bins, num_audio_tokens).""" + return self.audio_feature_extractor(audios) + + +__all__ = [ + "InklingImageProcessor", + "InklingAudioFeatureExtractor", + "InklingAudioEncoderParams", + "InklingProcessor", + "IMAGE_MARKER_ID", + "AUDIO_MARKER_ID", + "IMAGE_TOKEN_ID", + "AUDIO_TOKEN_ID", +] diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 1334d110b491..8b999c84783d 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -618,14 +618,16 @@ def flashinfer_mm_fp4_fake( ) def flashinfer_mxfp4_quantize( a: torch.Tensor, + backend: str, ) -> tuple[torch.Tensor, torch.Tensor]: from flashinfer import mxfp4_quantize as _mxfp4_quantize - return _mxfp4_quantize(a) + return _mxfp4_quantize(a, backend=backend) @torch.library.register_fake("vllm::flashinfer_mxfp4_quantize") def flashinfer_mxfp4_quantize_fake( a: torch.Tensor, + backend: str, ) -> tuple[torch.Tensor, torch.Tensor]: m, k = a.shape sf_vec_size = 32 diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index 8228e24c1c3b..4c3327283a3c 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -267,6 +267,26 @@ def _log_cutedsl_jit_compile(fn_name: str) -> None: ) +class _MonitoredCuteCompile: + """Logs JIT compilations; a plain function would break ``cute.compile[opts]``.""" + + def __init__(self, inner): + self._inner = inner + + def __getitem__(self, options) -> "_MonitoredCuteCompile": + return _MonitoredCuteCompile(self._inner[options]) + + def __call__(self, *args, **kwargs): + kernel = args[0] if args else kwargs.get("function") + kernel_name = getattr(kernel, "__name__", None) + if kernel_name is None: + kernel_name = ( + kernel.__class__.__name__ if kernel is not None else "" + ) + _log_cutedsl_jit_compile(kernel_name) + return self._inner(*args, **kwargs) + + def _setup_cutedsl_jit_hook() -> None: """Wrap ``cutlass.cute.compile`` to warn on compilation.""" global _cutedsl_hook_installed @@ -279,20 +299,7 @@ def _setup_cutedsl_jit_hook() -> None: logger.debug("CuTeDSL is not available; skipping CuTeDSL JIT monitor.") return - original_compile = cute.compile - - @functools.wraps(original_compile) - def _compile_with_monitor(*args, **kwargs): - kernel = args[0] if args else kwargs.get("function") - kernel_name = getattr(kernel, "__name__", None) - if kernel_name is None: - kernel_name = ( - kernel.__class__.__name__ if kernel is not None else "" - ) - _log_cutedsl_jit_compile(kernel_name) - return original_compile(*args, **kwargs) - - cute.compile = _compile_with_monitor + cute.compile = _MonitoredCuteCompile(cute.compile) _cutedsl_hook_installed = True diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 2708e0ab1f0e..a82d49dadb25 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -757,7 +757,9 @@ def use_cascade_attention( class AttentionLayer(Protocol): _q_scale: torch.Tensor _k_scale: torch.Tensor + _k_scale_cpu: torch.Tensor _v_scale: torch.Tensor + _v_scale_cpu: torch.Tensor _q_scale_float: float _k_scale_float: float _v_scale_float: float @@ -923,6 +925,14 @@ def fused_output_quant_supported(self, quant_key: "QuantKey") -> bool: """ return False + def fused_qk_norm_rope_kvcache_supported(self): + """ + Does this attention implementation support fused QKNorm+RoPE+KVCache fusion. + This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops + with the RoPE ops and the KV cache update for implementations that support it. + """ + return False + def fused_rope_kvcache_supported(self): """ Does this attention implementation support RoPE+KVCache fusion. @@ -931,6 +941,29 @@ def fused_rope_kvcache_supported(self): """ return False + def do_qk_norm_rope_kvcache_update( + self, + layer: AttentionLayer, + qkv: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache: torch.Tensor, + layer_slot_mapping: torch.Tensor, + ): + """ + If `fused_qk_norm_rope_kvcache_supported` returns True, this method + will be called by the fused custom op. Applies QK-norm + RoPE and + writes K/V to the KV cache. Results are written to the pre-allocated + q_out and k_out tensors; V is split from QKV at the graph level. + """ + raise NotImplementedError + def do_rope_and_kv_cache_update( self, layer: AttentionLayer, diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 9eb1cdbef321..5a029747247d 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -1573,6 +1573,8 @@ def __init__( ) self.sinks: torch.Tensor | None = None + # Keep the source so RL weight updates can refresh the runtime tensor. + self._sinks_source = sinks if sinks is not None: if sinks.shape[0] != num_heads: raise ValueError( @@ -1633,8 +1635,15 @@ def fused_output_quant_supported(self, quant_key: QuantKey): # FlashInfer requires attention sinks to be float32 def process_weights_after_loading(self, act_dtype: torch.dtype): - if self.sinks is not None and self.sinks.dtype != torch.float32: - self.sinks = self.sinks.to(torch.float32) + source_sinks = self._sinks_source + if source_sinks is None: + return + if source_sinks.dtype == torch.float32: + self.sinks = source_sinks + elif self.sinks is None or self.sinks.dtype != torch.float32: + self.sinks = source_sinks.to(torch.float32) + else: + self.sinks.copy_(source_sinks) def get_xqa_bmm1_scale(self, layer: torch.nn.Module, q_data_type: torch.dtype): bmm1_scale = self.scale diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 340a304030ea..70794be0751a 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -331,7 +331,9 @@ def build( # type: ignore[override] prefill_state_indices: torch.Tensor | None = None prefill_has_initial_state: torch.Tensor | None = None if num_prefills > 0: - from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE + from vllm.third_party.flash_linear_attention.ops.utils import ( + FLA_CHUNK_SIZE, + ) # In a mixed non-spec batch, decodes are peeled off to the recurrent # kernel (decode-first front slice), so build chunk metadata from the @@ -371,7 +373,7 @@ def build( # type: ignore[override] # Only prefill batches use FLA chunk ops. # Pre-compute on CPU and async-copy to GPU to avoid # GPU→CPU sync (.tolist()) in prepare_chunk_indices. - from vllm.model_executor.layers.fla.ops.index import ( + from vllm.third_party.flash_linear_attention.ops.index import ( prepare_chunk_indices, prepare_chunk_offsets, ) diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 16e292e21d2f..6fce7d0dc7e9 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -341,23 +341,36 @@ def _compute_prefix_caching_block_indices( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() # Block index of the last computed token - block_idx_last_computed_token = cdiv(num_computed_tokens, mamba_block_size) - 1 + block_idx_last_computed_token = ( + torch.div( + num_computed_tokens + mamba_block_size - 1, + mamba_block_size, + rounding_mode="floor", + ) + - 1 + ) # which is <= block index for the first scheduled token block_idx_first_scheduled_token = ( - cdiv(num_computed_tokens + 1, mamba_block_size) - 1 + torch.div( + num_computed_tokens + mamba_block_size, + mamba_block_size, + rounding_mode="floor", + ) + - 1 ) # which is <= block index of the last scheduled token block_idx_last_scheduled_token = ( - cdiv(common_attn_metadata.seq_lens, mamba_block_size) - 1 + torch.div( + common_attn_metadata.seq_lens + mamba_block_size - 1, + mamba_block_size, + rounding_mode="floor", + ) + - 1 ) # -1 in case it's non-computed and causes later issues with indexing - block_idx_last_computed_token = torch.clamp( - block_idx_last_computed_token, min=0 - ) + block_idx_last_computed_token.clamp_(min=0) # -1 in the case we have a padded request (0 seq-len) - block_idx_last_scheduled_token = torch.clamp( - block_idx_last_scheduled_token, min=0 - ) + block_idx_last_scheduled_token.clamp_(min=0) return ( block_idx_last_computed_token, @@ -447,9 +460,8 @@ def _compute_common_metadata( common_attn_metadata, mamba_block_size ) if self.use_spec_decode and prev_last_scheduled_idx is not None: - fallback = torch.clamp( - (num_computed_tokens - 1) // mamba_block_size, min=0 - ) + fallback = (num_computed_tokens - 1) // mamba_block_size + fallback.clamp_(min=0) block_idx_last_scheduled_token_prev_step = torch.where( prev_last_scheduled_idx >= 0, prev_last_scheduled_idx, diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 4dea36a3decc..e8be7ad5bb39 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -32,7 +32,6 @@ from vllm.v1.attention.backends.utils import ( reshape_attn_output_for_spec_decode, reshape_query_for_spec_decode, - split_decodes_and_prefills, split_prefill_chunks, ) from vllm.v1.attention.ops.flashmla import ( @@ -57,8 +56,8 @@ # the FP8 decode kernel for decode. # Currently we use #1 when the number of heads per rank is low (i.e. TP) since the BF16 # prefill kernel requires padding the number of heads to 128 while the decode does not -# so when the per ranke head count is below MIN_HEADS_FOR_BF16_PREFILL we use the mixed -# batch mode (#2). +# so when the per-rank head count is below MIN_HEADS_FOR_BF16_PREFILL we use the mixed +# batch mode (#1). MIN_HEADS_FOR_BF16_PREFILL = 32 """ @@ -182,10 +181,6 @@ class Decode: @dataclass class Prefill: - # Sequence lengths (context + query) for prefill requests - # Shape: [num_prefill_reqs] - seq_lens: torch.Tensor - # Request ID for each token: -1 for decode tokens, request index # (0, 1, 2, ...) for prefill tokens. # Shape: [num_actual_tokens] @@ -204,7 +199,6 @@ class Chunk: Prefill requests may be chunked to fit within the fixed workspace size. """ - seq_lens: torch.Tensor tokens_slice: slice block_table: torch.Tensor req_start_idx: int @@ -239,6 +233,7 @@ class FlashMLASparseMetadataBuilder( SparseMLACommonMetadataBuilder[FlashMLASparseMetadata] ): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + require_uniform_decodes: ClassVar[bool] = True metadata_cls = FlashMLASparseMetadata def __init__( @@ -322,10 +317,11 @@ def _build_fp8_mixed_decode_prefill( self, common_attn_metadata: CommonAttentionMetadata, ) -> "FlashMLASparseMetadata.FP8KernelMetadata": - """Build FP8 metadata treating all tokens as one mixed batch. + """Build FP8 metadata treating MQA tokens as one batch. - This matches main branch's approach and avoids the BF16 prefill kernel - which has head padding overhead when num_heads is small (high TP case). + The scheduler initializes lazily from the runtime query shape, which may + be the full batch or only decodes when prefills use dense MHA. This avoids + the BF16 prefill kernel's head-padding overhead at high TP. """ num_tokens = common_attn_metadata.num_actual_tokens @@ -353,20 +349,29 @@ def _build_fp8_mixed_decode_prefill( def _build_fp8_separate_prefill_decode( self, common_attn_metadata: CommonAttentionMetadata, + metadata: FlashMLASparseMetadata, ) -> "FlashMLASparseMetadata.FP8SeparatePrefillDecode": num_tokens = common_attn_metadata.num_actual_tokens (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = ( - split_decodes_and_prefills( - common_attn_metadata, - decode_threshold=self.reorder_batch_threshold or 1, - require_uniform=True, - ) + metadata.num_decodes, + metadata.num_prefills, + metadata.num_decode_tokens, + num_tokens - metadata.num_decode_tokens, ) + decode_query_len = 0 + active_num_decodes = num_decodes + if num_decodes > 0: + query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu + decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item() + assert decode_query_len > 0 + active_num_decodes = num_decode_tokens // decode_query_len + assert active_num_decodes * decode_query_len == num_decode_tokens + FP8Meta = FlashMLASparseMetadata.FP8SeparatePrefillDecode fp8_metadata = FP8Meta( - num_decodes=num_decodes, + num_decodes=active_num_decodes, num_prefills=num_prefills, num_decode_tokens=num_decode_tokens, num_prefill_tokens=num_prefill_tokens, @@ -374,7 +379,6 @@ def _build_fp8_separate_prefill_decode( # Extract prefill sequence lengths (context + query, not just query) # Decode requests come first in the batch, prefill requests follow - prefill_seq_lens = None prefill_request_id = None prefill_workspace_starts = None prefill_chunks = None @@ -386,11 +390,9 @@ def _build_fp8_separate_prefill_decode( # slice below), so no D2H sync is needed. seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound assert seq_lens_cpu is not None - seq_lens = common_attn_metadata.seq_lens query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:] - prefill_seq_lens = seq_lens[num_decodes:] # Build prefill_request_id: -1 for decode, request index for # prefill. This enables a single @@ -438,7 +440,6 @@ def _build_fp8_separate_prefill_decode( offset = prefill_workspace_starts_cpu[chunk_start].item() prefill_workspace_starts_cpu[chunk_start:chunk_end] -= offset - chunk_seq_lens = prefill_seq_lens[chunk_start:chunk_end] chunk_tot_seqlen = prefill_seq_lens_cpu[chunk_start:chunk_end].sum() token_start = query_start_loc_cpu[num_decodes + chunk_start].item() token_end = query_start_loc_cpu[num_decodes + chunk_end].item() @@ -452,7 +453,6 @@ def _build_fp8_separate_prefill_decode( prefill_chunks.append( FP8Meta.Prefill.Chunk( - seq_lens=chunk_seq_lens, tokens_slice=tokens_slice, block_table=chunk_block_table, req_start_idx=chunk_start, @@ -466,27 +466,22 @@ def _build_fp8_separate_prefill_decode( ) fp8_metadata.prefill = FP8Meta.Prefill( - seq_lens=prefill_seq_lens, request_ids=prefill_request_id, workspace_starts=prefill_workspace_starts, chunks=prefill_chunks, ) if num_decodes > 0: - # Compute decode_query_len for spec decode (uniform due to require_uniform) - query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu - decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item() - # Use padded head count since that's what the kernel will see scheduler_metadata, _ = get_mla_metadata() kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata( scheduler_metadata=scheduler_metadata, - dummy_block_table=self.dummy_block_table[:num_decodes], - cache_lens=self.max_model_len_tensor[:num_decodes], + dummy_block_table=self.dummy_block_table[:active_num_decodes], + cache_lens=self.max_model_len_tensor[:active_num_decodes], ) fp8_metadata.decode = FP8Meta.Decode( - seq_lens=common_attn_metadata.seq_lens[:num_decodes], + seq_lens=common_attn_metadata.seq_lens[:active_num_decodes], kernel_metadata=kernel_meta, decode_query_len=decode_query_len, ) @@ -510,7 +505,7 @@ def build( ) else: metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode( - common_attn_metadata + common_attn_metadata, metadata ) return metadata @@ -597,9 +592,10 @@ def _forward_bf16_kv( attn_metadata: FlashMLASparseMetadata, ) -> torch.Tensor: # Convert per-request indices to global slots (decode) or workspace - # offsets (prefill). + # offsets (prefill). req_id_per_token covers the whole batch; slice it + # to the MQA tokens (q may exclude prefill tokens routed to dense MHA). topk_indices, topk_length = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token, + attn_metadata.req_id_per_token[: topk_indices.shape[0]], attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, @@ -624,11 +620,18 @@ def _forward_fp8_kv_separate_prefill_decode( fp8_metadata = attn_metadata.fp8_extra_metadata assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode) num_decodes = fp8_metadata.num_decodes + num_mqa_tokens = q.shape[0] + num_decode_tokens = fp8_metadata.num_decode_tokens + num_prefill_tokens = num_mqa_tokens - num_decode_tokens + assert num_prefill_tokens in (0, fp8_metadata.num_prefill_tokens), ( + "FP8 sparse MLA expects either the decode subset or the full batch" + ) prefill_request_ids = None prefill_workspace_starts = None has_prefill_workspace = False - if fp8_metadata.prefill is not None: + if num_prefill_tokens > 0: + assert fp8_metadata.prefill is not None prefill_request_ids = fp8_metadata.prefill.request_ids prefill_workspace_starts = fp8_metadata.prefill.workspace_starts has_prefill_workspace = True @@ -640,7 +643,7 @@ def _forward_fp8_kv_separate_prefill_decode( # prefill_workspace_starts has been adjusted in-place per chunk so # prefill indices automatically come out chunk-local topk_indices, topk_length = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token, + attn_metadata.req_id_per_token[: topk_indices.shape[0]], attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, @@ -676,9 +679,6 @@ def _fp8_decode( # -> (num_decode_tokens, num_heads, head_dim_v) return reshape_attn_output_for_spec_decode(attn_out) - num_decode_tokens = fp8_metadata.num_decode_tokens - num_prefill_tokens = fp8_metadata.num_prefill_tokens - # Pure decode: direct call without allocation if num_decode_tokens > 0 and num_prefill_tokens == 0: assert fp8_metadata.decode is not None @@ -686,7 +686,7 @@ def _fp8_decode( else: # Mixed or pure prefill: allocate output tensor attn_out = q.new_empty( - (attn_metadata.num_actual_tokens, self.num_heads, self.kv_lora_rank), + (num_mqa_tokens, self.num_heads, self.kv_lora_rank), dtype=q.dtype, device=q.device, ) @@ -704,7 +704,6 @@ def _fp8_decode( kv_c_and_k_pe_cache, chunk_workspace, chunk.block_table, - chunk.seq_lens, chunk.workspace_starts, len(chunk.block_table), ) @@ -738,7 +737,7 @@ def _forward_fp8_kv_mixed_batch( # Convert per-request indices to global slots (decode) or workspace # offsets (prefill). topk_indices = triton_convert_req_index_to_global_index( - attn_metadata.req_id_per_token, + attn_metadata.req_id_per_token[: topk_indices.shape[0]], attn_metadata.block_table, topk_indices, BLOCK_SIZE=attn_metadata.block_size, diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index afc4efa1ae9c..30c0c6d8a681 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -63,11 +63,16 @@ def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool: qk_rope_head_dim=64, v_head_dim=256, ) + dims_mistral_s4 = MLADimensions( + qk_nope_head_dim=64, + qk_rope_head_dim=64, + v_head_dim=128, + ) fa_version = get_flash_attn_version() if fa_version == 4: - return mla_dimensions == dims_deepseek + return mla_dimensions in [dims_deepseek, dims_mistral_s4] else: - return mla_dimensions in [dims_deepseek, dims_glm] + return mla_dimensions in [dims_deepseek, dims_glm, dims_mistral_s4] def __init__( self, diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index 522b52b0dcd2..681ab7cfd9fa 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -154,6 +154,11 @@ def triton_convert_req_index_to_global_index( assert req_id.dtype == torch.int32 assert block_table.dtype == torch.int32 assert token_indices.dtype == torch.int32 + assert req_id.shape[0] == token_indices.shape[0], ( + f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) " + "must cover the same tokens; the grid is sized by req_id but the output " + "is allocated like token_indices, so a longer req_id writes out of bounds" + ) assert token_indices.shape[1] == NUM_TOPK_TOKENS assert NUM_TOPK_TOKENS % BLOCK_N == 0, ( f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})" diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index acc9c9cb5010..5c10c16a78f4 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/vllm/v1/attention/backends/mla/triton_mla.py @@ -180,6 +180,32 @@ def __init__( "TritonMLAImpl" ) + if current_platform.is_cuda(): + cap = current_platform.get_device_capability() + cap_str = cap.as_version_str() if cap is not None else "unknown" + dev = current_platform.get_device_name() + if self.kv_cache_dtype.startswith("fp8") and not ( + current_platform.has_device_capability(89) + ): + suggested = ( + "float16" if (cap is None or cap.to_int() < 80) else "bfloat16" + ) + raise ValueError( + f"FP8 KV cache is not supported by the Triton MLA backend " + f"on {dev} (compute capability {cap_str}); native FP8 " + f"(fp8e4nv) requires SM89+. Re-run with " + f"--kv-cache-dtype {suggested}." + ) + if self.kv_cache_dtype == "bfloat16" and not ( + current_platform.has_device_capability(80) + ): + raise ValueError( + f"bfloat16 KV cache is not supported by the Triton MLA " + f"backend on {dev} (compute capability {cap_str}); " + f"bfloat16 requires SM80+. Re-run with " + f"--kv-cache-dtype float16." + ) + # For FP8 KV cache, we dequantize to BF16 on load inside the # Triton kernel. Tell the common layer not to quantize queries # to FP8 — we handle FP8 KV cache with BF16 queries (Mode 1). diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index b2edc2d05158..1f46dddabfc5 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -593,9 +593,8 @@ def build( chunk_ends = torch.min( computed_kv_lens.unsqueeze(0), chunk_starts + max_context_chunk ) - chunk_seq_lens = (chunk_ends - chunk_starts).clamp( - min=0 - ) # [num_chunks, num_extends] + chunk_seq_lens = chunk_ends - chunk_starts + chunk_seq_lens.clamp_(min=0) # [num_chunks, num_extends] cu_seq_lens_cpu = torch.zeros( [num_chunks, num_extends + 1], dtype=torch.int32, pin_memory=True ) @@ -1449,6 +1448,47 @@ def fused_rope_kvcache_supported(self): and not rocm_aiter_ops.is_shuffle_kv_cache_enabled() ) + def fused_qk_norm_rope_kvcache_supported(self): + return rocm_aiter_ops.is_enabled() + + def do_qk_norm_rope_kvcache_update( + self, + layer: AttentionLayer, + qkv: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache: torch.Tensor, + layer_slot_mapping: torch.Tensor, + ): + key_cache, value_cache = kv_cache.unbind(1) + rocm_aiter_ops.do_qk_norm_rope_kvcache_update( + qkv=qkv, + q_weight=q_weight, + k_weight=k_weight, + cos_sin_cache=cos_sin_cache, + positions=positions, + num_heads_q=self.num_heads, + num_heads_k=self.num_kv_heads, + head_dim=self.head_size, + is_neox=is_neox, + rms_norm_eps=rms_norm_eps, + q_out=q_out, + k_out=k_out, + key_cache=key_cache, + value_cache=value_cache, + slot_mapping=layer_slot_mapping, + k_scale=layer._k_scale_cpu, + v_scale=layer._v_scale_cpu, + kv_cache_dtype=self.kv_cache_dtype, + use_shuffle_layout=rocm_aiter_ops.is_shuffle_kv_cache_enabled(), + ) + def do_rope_and_kv_cache_update( self, layer: AttentionLayer, diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index cc45fade5ab4..08cd86968679 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -311,6 +311,50 @@ def do_kv_cache_update( def fused_rope_kvcache_supported(self): return rocm_aiter_ops.is_enabled() + def fused_qk_norm_rope_kvcache_supported(self): + return rocm_aiter_ops.is_enabled() + + def do_qk_norm_rope_kvcache_update( + self, + layer: AttentionLayer, + qkv: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + is_neox: bool, + kv_cache: torch.Tensor, + layer_slot_mapping: torch.Tensor, + ): + # _split_kv_cache picks the unbind dim per layout (incl. the K/V-first + # encoder-decoder path). unified reads NHD, so never write the shuffle + # layout here. + key_cache, value_cache = self._split_kv_cache(kv_cache) + rocm_aiter_ops.do_qk_norm_rope_kvcache_update( + qkv=qkv, + q_weight=q_weight, + k_weight=k_weight, + cos_sin_cache=cos_sin_cache, + positions=positions, + num_heads_q=self.num_heads, + num_heads_k=self.num_kv_heads, + head_dim=self.head_size, + is_neox=is_neox, + rms_norm_eps=rms_norm_eps, + q_out=q_out, + k_out=k_out, + key_cache=key_cache, + value_cache=value_cache, + slot_mapping=layer_slot_mapping, + k_scale=layer._k_scale_cpu, + v_scale=layer._v_scale_cpu, + kv_cache_dtype=self.kv_cache_dtype, + use_shuffle_layout=False, + ) + def do_rope_and_kv_cache_update( self, layer: AttentionLayer, diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 1e12f43caacb..c8c9a7334a29 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -950,10 +950,8 @@ def mamba_get_block_table_tensor( assert isinstance(kv_cache_spec, MambaSpec) # NOTE: For 0-length requests in CUDA graph, use a start_index of 0 # to handle the invalid block table. - start_indices = torch.clamp( - (seq_lens - 1) // kv_cache_spec.block_size, - min=0, - ) + start_indices = (seq_lens - 1) // kv_cache_spec.block_size + start_indices.clamp_(min=0) # Use int32 for arithmetic to avoid dtype promotion overhead, # then convert to int64 for gather (which requires Long indices) offsets = torch.arange( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 2153a460f696..c3faee402461 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -1187,10 +1187,12 @@ def _sparse_attn_prefill_ragged_kernel( kv_len = kv_end - kv_start k_offsets = tl.arange(0, BLOCK_K) + slot = tl.load( + kv_indices_ptr + kv_start + k_offsets, mask=k_offsets < kv_len, other=-1 + ) for k_start in tl.range(0, kv_len, BLOCK_K): k_pos = k_start + k_offsets in_range = k_pos < kv_len - slot = tl.load(kv_indices_ptr + kv_start + k_pos, mask=in_range, other=-1) valid = in_range & (slot >= 0) & (slot < num_kv) safe_slot = tl.where(valid, slot, 0) @@ -1201,7 +1203,11 @@ def _sparse_attn_prefill_ragged_kernel( mask=valid[:, None] & dim_mask[None, :], other=0.0, ) - kv = tl.where(valid[:, None] & dim_mask[None, :], kv, 0.0) + + next_k_pos = k_start + BLOCK_K + k_offsets + slot = tl.load( + kv_indices_ptr + kv_start + next_k_pos, mask=next_k_pos < kv_len, other=-1 + ) scores = tl.dot(q, tl.trans(kv)) * scale scores = tl.where(head_mask[:, None] & valid[None, :], scores, neg_large) @@ -1865,6 +1871,7 @@ def _rocm_sparse_attn_prefill_ragged_triton( block_h = 16 block_d = triton.next_power_of_2(head_dim) block_k = 16 if head_dim >= 256 else 32 + num_warps = 4 out = torch.empty_like(q, dtype=torch.bfloat16) _sparse_attn_prefill_ragged_kernel[(num_queries, triton.cdiv(num_heads, block_h))]( q, @@ -1889,7 +1896,7 @@ def _rocm_sparse_attn_prefill_ragged_triton( BLOCK_H=block_h, BLOCK_D=block_d, BLOCK_K=block_k, - num_warps=8, + num_warps=num_warps, ) return out @@ -2097,7 +2104,7 @@ def _rocm_sparse_attn_decode_ragged_triton( comb_dim = nope_head_dim + rope_head_dim is_fnuz = current_platform.is_fp8_fnuz() - if not _ON_GFX950: # Fallback path for un-tuned architectures. + if not (_ON_GFX942 or _ON_GFX950): # Fallback path for un-tuned architectures. block_k = 16 if head_dim >= 256 else 32 _sparse_attn_decode_ragged_kernel[(num_queries, heads_blocks)]( q, diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index e3388a9f85a2..3d3d8c1573a9 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -681,6 +681,34 @@ def take_new_block_ids(self) -> list[int]: ids.extend(mgr.take_new_block_ids()) return ids + def get_zeroing_block_ids_in_range( + self, request_id: str, start_token: int, end_token: int + ) -> list[int]: + """The request's block ids covering [start_token, end_token), from + the groups whose new blocks are zeroed by the worker.""" + ids: list[int] = [] + for mgr in self.coordinator.single_type_managers: + if mgr.records_new_block_ids: + start_idx = start_token // mgr.block_size + end_idx = cdiv(end_token, mgr.block_size) + blocks = mgr.req_to_blocks[request_id] + ids.extend(blk.block_id for blk in blocks[start_idx:end_idx]) + return ids + + def record_blocks_for_zeroing(self, request_id: str, start_token: int) -> None: + """Re-record the request's blocks from start_token onwards for + zeroing, e.g. blocks a failed async KV load left unwritten. + + start_token must be block-aligned: zeroing a partially-valid block + would wipe its valid prefix. + """ + for mgr in self.coordinator.single_type_managers: + if mgr.records_new_block_ids: + assert start_token % mgr.block_size == 0 + start_idx = start_token // mgr.block_size + blocks = mgr.req_to_blocks[request_id] + mgr.new_block_ids.extend(blk.block_id for blk in blocks[start_idx:]) + def take_kv_cache_block_copies( self, ) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]: diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 4401fb050b3f..5667482e948b 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -179,6 +179,14 @@ def make_empty(cls) -> "CachedRequestData": ) +@dataclass +class ScheduledEncoderInputStats: + """Stats for encoder inputs scheduled in one iteration.""" + + num_inputs: int = 0 + output_tokens: int = 0 + + @dataclass class SchedulerOutput: # list of the requests that are scheduled for the first time. @@ -216,6 +224,8 @@ class SchedulerOutput: # freed from the encoder cache. free_encoder_mm_hashes: list[str] + scheduled_encoder_input_stats: ScheduledEncoderInputStats | None = None + # Request IDs that are preempted in this step. # Only used for v2 model runner. preempted_req_ids: set[str] | None = None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 371139fa7226..d6f3c3ad0e38 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -44,6 +44,7 @@ CachedRequestData, GrammarOutput, NewRequestData, + ScheduledEncoderInputStats, SchedulerOutput, ) from vllm.v1.core.sched.request_queue import ( @@ -296,6 +297,9 @@ def __init__( self.has_mamba_layers = kv_cache_config.has_mamba_layers self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing + # Blocks that async KV loads will overwrite this step, skipped from + # zeroing since the zeroing could race the out-of-band write. + self._skip_zero_block_ids: set[int] = set() self.need_mamba_block_aligned_split = ( self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" ) @@ -985,6 +989,16 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # only the successfully loaded tokens. request.num_computed_tokens = num_computed_tokens self._inflight_prefills.add(request) + if self.needs_kv_cache_zeroing: + # Skip zeroing of the blocks the async load will + # overwrite; the zeroing could race the write. + self._skip_zero_block_ids.update( + self.kv_cache_manager.get_zeroing_block_ids_in_range( + request.request_id, + num_new_local_computed_tokens, + num_computed_tokens, + ) + ) continue self.running.append(request) @@ -1097,12 +1111,6 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) - # Drain new attention block ids every step so the manager-side list - # does not grow unbounded; only kv-cache zeroing consumes them. - new_attn_block_ids = self.kv_cache_manager.take_new_block_ids() - new_block_ids_to_zero = ( - (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None - ) kv_cache_block_copies, cow_retained_blocks = ( self.kv_cache_manager.take_kv_cache_block_copies() ) @@ -1121,6 +1129,15 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: len(num_scheduled_tokens) ] + scheduled_encoder_input_stats = None + if ( + self.log_stats + and self.observability_config.enable_logging_iteration_details + ): + scheduled_encoder_input_stats = self._make_scheduled_encoder_input_stats( + scheduled_encoder_inputs + ) + scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, scheduled_cached_reqs=cached_reqs_data, @@ -1128,6 +1145,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: total_num_scheduled_tokens=total_num_scheduled_tokens, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, scheduled_encoder_inputs=scheduled_encoder_inputs, + scheduled_encoder_input_stats=scheduled_encoder_input_stats, num_common_prefix_blocks=num_common_prefix_blocks, preempted_req_ids=self.reset_preempted_req_ids, # finished_req_ids is an existing state in the scheduler, @@ -1136,7 +1154,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # the previous and the current steps. finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), - new_block_ids_to_zero=new_block_ids_to_zero, + new_block_ids_to_zero=self._get_new_block_ids_to_zero(), kv_cache_block_copies=pending_kv_cache_block_copies, num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, ) @@ -1170,6 +1188,20 @@ def _build_kv_connector_meta( ) -> KVConnectorMetadata: return connector.build_connector_meta(scheduler_output) + def _get_new_block_ids_to_zero(self) -> list[int] | None: + # Drain new attention block ids every step so the manager-side list + # does not grow unbounded; only kv-cache zeroing consumes them. + new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids() + if not self.needs_kv_cache_zeroing: + return None + + if self._skip_zero_block_ids: + skip = self._skip_zero_block_ids + new_block_ids_to_zero = [b for b in new_block_ids_to_zero if b not in skip] + skip.clear() + + return new_block_ids_to_zero or None + def _preempt_request(self, request: Request, timestamp: float) -> None: """Preempt a request and put it back to the waiting queue. @@ -1506,6 +1538,23 @@ def _try_schedule_encoder_inputs( external_load_encoder_input, ) + def _make_scheduled_encoder_input_stats( + self, scheduled_encoder_inputs: dict[str, list[int]] + ) -> ScheduledEncoderInputStats | None: + stats = ScheduledEncoderInputStats() + + for req_id, input_ids in scheduled_encoder_inputs.items(): + request = self.requests.get(req_id) + if request is None: + continue + + for input_id in input_ids: + mm_feature = request.mm_features[input_id] + stats.num_inputs += 1 + stats.output_tokens += mm_feature.mm_position.get_num_embeds() + + return stats if stats.num_inputs else None + def get_grammar_bitmask( self, scheduler_output: SchedulerOutput ) -> GrammarOutput | None: @@ -1877,7 +1926,10 @@ def update_from_output( if ( stats := self.make_stats( - spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats + spec_decoding_stats, + kv_connector_stats, + cudagraph_stats, + perf_stats, ) ) is not None: # Return stats to only one of the front-ends. @@ -2487,9 +2539,18 @@ def _update_waiting_for_remote_kv(self, request: Request) -> None: if request.num_computed_tokens: # Cache any valid computed tokens. self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens) + if self.needs_kv_cache_zeroing: + # The failed load left the blocks beyond the valid + # prefix unwritten and their zeroing was skipped; zero + # them before they are recomputed locally. + self.kv_cache_manager.record_blocks_for_zeroing( + request.request_id, request.num_computed_tokens + ) else: # No valid computed tokens, release allocated blocks. # There may be a local cache hit on retry. + # (Freed blocks are re-recorded for zeroing when + # reallocated, so the skipped blocks need no handling.) self.kv_cache_manager.free(request) self.failed_recving_kv_req_ids.remove(request.request_id) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index de47ce0a2c78..244048390947 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -359,6 +359,11 @@ def allocate_new_blocks( self.new_block_ids.extend(b.block_id for b in new_blocks) return cow_blocks + new_blocks + @property + def records_new_block_ids(self) -> bool: + """Whether this manager's new blocks are zeroed by the worker.""" + return self._record_new_block_ids + def take_new_block_ids(self) -> list[int]: """Drain and return block IDs allocated since the last call.""" ids = self.new_block_ids diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 8bcd4ba89a4e..93e02abf7479 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -1067,17 +1067,8 @@ async def init_weight_transfer_engine( Args: request: Weight transfer initialization request with backend-specific info """ - from vllm.distributed.weight_transfer.base import ( - WeightTransferInitRequest, - ) - - if isinstance(request, WeightTransferInitRequest): - init_info_dict = request.init_info - else: - raise TypeError(f"Expected WeightTransferInitRequest, got {type(request)}") - await self.collective_rpc( - "init_weight_transfer_engine", kwargs={"init_info": init_info_dict} + "init_weight_transfer_engine", kwargs={"init_info": request.init_info} ) async def start_weight_update(self) -> None: @@ -1095,16 +1086,8 @@ async def update_weights(self, request: WeightTransferUpdateRequest) -> None: Args: request: Weight update request with backend-specific update info """ - - if isinstance(request, WeightTransferUpdateRequest): - update_info_dict = request.update_info - else: - raise TypeError( - f"Expected WeightTransferUpdateRequest, got {type(request)}" - ) - await self.collective_rpc( - "update_weights", kwargs={"update_info": update_info_dict} + "update_weights", kwargs={"update_info": request.update_info} ) async def finish_weight_update(self) -> None: diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 301c5892a58a..383853807db6 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -79,16 +79,17 @@ ) from vllm.v1.executor import Executor from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind -from vllm.v1.metrics.stats import SchedulerStats +from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder from vllm.v1.structured_output import StructuredOutputManager -from vllm.v1.utils import IterationDetails, compute_iteration_details +from vllm.v1.utils import compute_iteration_details from vllm.version import __version__ as VLLM_VERSION logger = init_logger(__name__) + HANDSHAKE_TIMEOUT_MINS = 5 _R = TypeVar("_R") # Return type for collective_rpc @@ -498,45 +499,74 @@ def log_error_detail(self, scheduler_output: SchedulerOutput): raise err @contextmanager - def log_iteration_details(self, scheduler_output: SchedulerOutput | None): - if not self.vllm_config.observability_config.enable_logging_iteration_details: - yield + def capture_iteration_details( + self, scheduler_output: SchedulerOutput | None + ) -> Generator[SchedulerIterationDetails | None, None, None]: + enable_details = ( + self.vllm_config.observability_config.enable_logging_iteration_details + ) + if not self.log_stats or not enable_details: + yield None return # 0-token step: let the dummy_batch wrapper log it (avoids double-log). - if scheduler_output and scheduler_output.total_num_scheduled_tokens == 0: - yield + if ( + scheduler_output is not None + and scheduler_output.total_num_scheduled_tokens == 0 + ): + yield None return - self._iteration_index = getattr(self, "_iteration_index", 0) + + iteration_index = getattr(self, "_iteration_index", 0) # scheduler_output=None marks a DP dummy iteration. if scheduler_output is None: - iteration_details = IterationDetails(0, 0, 0, 0) - is_dummy = True + iteration_details = SchedulerIterationDetails( + iteration_index=iteration_index, + num_ctx_requests=0, + num_ctx_tokens=0, + num_generation_requests=0, + num_generation_tokens=0, + elapsed_ms=0.0, + is_dummy=True, + ) else: - iteration_details = compute_iteration_details(scheduler_output) - is_dummy = False - before = time.monotonic() - yield - logger.info( - "".join( - [ - "Iteration(", - str(self._iteration_index), - "): ", - str(iteration_details.num_ctx_requests), - " context requests, ", - str(iteration_details.num_ctx_tokens), - " context tokens, ", - str(iteration_details.num_generation_requests), - " generation requests, ", - str(iteration_details.num_generation_tokens), - " generation tokens, iteration elapsed time: ", - format((time.monotonic() - before) * 1000, ".2f"), - " ms", - " (dummy)" if is_dummy else "", - ] + details = compute_iteration_details(scheduler_output) + iteration_details = SchedulerIterationDetails( + iteration_index=iteration_index, + num_ctx_requests=details.num_ctx_requests, + num_ctx_tokens=details.num_ctx_tokens, + num_generation_requests=details.num_generation_requests, + num_generation_tokens=details.num_generation_tokens, + elapsed_ms=0.0, + num_encoder_inputs=details.num_encoder_inputs, + num_encoder_output_tokens=details.num_encoder_output_tokens, ) - ) - self._iteration_index += 1 + + start_time = time.monotonic() + yield iteration_details + iteration_details.elapsed_ms = (time.monotonic() - start_time) * 1000 + self._iteration_index = iteration_index + 1 + + def _make_iteration_details_stats( + self, iteration_details: SchedulerIterationDetails + ) -> SchedulerStats: + stats = self.scheduler.make_stats() or SchedulerStats() + stats.iteration_details = iteration_details + return stats + + def _attach_iteration_details( + self, + outputs: dict[int, EngineCoreOutputs], + iteration_details: SchedulerIterationDetails | None, + ) -> None: + if iteration_details is None: + return + + if (eco := next(iter(outputs.values()), None)) is None: + outputs[0] = eco = EngineCoreOutputs() + if eco.scheduler_stats is None: + eco.scheduler_stats = self._make_iteration_details_stats(iteration_details) + else: + eco.scheduler_stats.iteration_details = iteration_details def _should_throttle_prefills(self) -> bool: """Whether to defer new prefills this step (DP prefill balancing). @@ -558,8 +588,8 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: future = self.model_executor.execute_model(scheduler_output, non_block=True) grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output) with ( + self.capture_iteration_details(scheduler_output) as iteration_details, self.log_error_detail(scheduler_output), - self.log_iteration_details(scheduler_output), ): model_output = future.result() if model_output is None: @@ -571,6 +601,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: engine_core_outputs = self.scheduler.update_from_output( scheduler_output, model_output ) + self._attach_iteration_details(engine_core_outputs, iteration_details) return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 @@ -656,8 +687,8 @@ def step_with_batch_queue( # Block until the next result is available. future, scheduler_output, exec_model_fut = batch_queue.pop() with ( + self.capture_iteration_details(scheduler_output) as iteration_details, self.log_error_detail(scheduler_output), - self.log_iteration_details(scheduler_output), ): model_output = future.result() if model_output is None: @@ -672,6 +703,7 @@ def step_with_batch_queue( engine_core_outputs = self.scheduler.update_from_output( scheduler_output, model_output ) + self._attach_iteration_details(engine_core_outputs, iteration_details) # NOTE(nick): We can either handle the deferred tasks here or save # in a field and do it immediately once step_with_batch_queue is @@ -2019,8 +2051,13 @@ def run_busy_loop(self): # Execute a dummy pass when no ready requests ran, unless the # engine is sleeping. elif not self.model_executor.is_sleeping: - with self.log_iteration_details(None): + with self.capture_iteration_details(None) as iteration_details: self.execute_dummy_batch() + if iteration_details is not None and not self.has_coordinator: + stats = self._make_iteration_details_stats(iteration_details) + self.output_queue.put_nowait( + (0, EngineCoreOutputs(scheduler_stats=stats)) + ) # 3) All-reduce operation to determine global unfinished reqs. self.engines_running = self._has_global_unfinished_reqs( diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 5a2e3c184d39..60c0aa4374b8 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -14,14 +14,13 @@ import torch from vllm.logger import init_logger -from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes if TYPE_CHECKING: - from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) - from vllm.v1.kv_cache_interface import KVCacheConfig + +from vllm.v1.kv_offload.config import OffloadingConfig # `OffloadKey` identifies an offloaded block. It combines a block hash with # its KV cache group index, encoded as raw bytes to avoid tuple GC overhead. @@ -98,12 +97,20 @@ class PrepareStoreOutput: evicted_keys: list[OffloadKey] +class Locality(Enum): + """Locality of a tier's storage relative to the publishing instance.""" + + LOCAL = "LOCAL" + REMOTE = "REMOTE" + + @dataclass class OffloadingEvent: keys: list[OffloadKey] medium: str # True if blocks are removed, False if stored removed: bool + locality: Locality | None = None """ @@ -482,22 +489,15 @@ def build_metric_definitions( """Return Prometheus metric definitions emitted by this spec.""" return {} - def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"): + def __init__(self, config: OffloadingConfig): logger.warning( "Initializing OffloadingSpec. This API is experimental and " "subject to change in the future as we iterate the design." ) - self.vllm_config = vllm_config - self.kv_cache_config = kv_cache_config - - kv_transfer_config = vllm_config.kv_transfer_config - assert kv_transfer_config is not None - self.extra_config = kv_transfer_config.kv_connector_extra_config - kv_events_config = vllm_config.kv_events_config + self.config = config + self.extra_config = config.extra_config self.kv_events_config = OffloadingKVEventsConfig( - enable_kv_cache_events=( - kv_events_config is not None and kv_events_config.enable_kv_cache_events - ), + enable_kv_cache_events=config.enable_kv_cache_events, self_describing_kv_events=bool( self.extra_config.get("self_describing_kv_events", False) ), @@ -511,48 +511,9 @@ def __init__(self, vllm_config: "VllmConfig", kv_cache_config: "KVCacheConfig"): self.extra_config.get("offload_prompt_only", True) ) - parallel_config = vllm_config.parallel_config - context_parallel_factor = ( - parallel_config.decode_context_parallel_size - * parallel_config.prefill_context_parallel_size - ) - - # gpu block size per group - self.gpu_block_size: tuple[int, ...] = tuple( - kv_cache_group.kv_cache_spec.block_size * context_parallel_factor - for kv_cache_group in kv_cache_config.kv_cache_groups - ) - - # hash_block_size must match what the scheduler uses for - # Request.block_hashes (resolved via resolve_kv_cache_block_sizes). - _, self.hash_block_size = resolve_kv_cache_block_sizes( - kv_cache_config, vllm_config - ) - - for block_size in self.gpu_block_size: - assert block_size % self.hash_block_size == 0, ( - f"gpu_block_size={block_size} not divisible by " - f"hash_block_size={self.hash_block_size}. " - f"Hybrid models (e.g. Mamba+Attention) need " - f"--enable-prefix-caching to align block sizes." - ) - - # offloaded_block_size / gpu_block_size - self.block_size_factor: int = 1 - - offloaded_block_size = self.extra_config.get("block_size") - if offloaded_block_size is not None: - offloaded_block_size_int = int(offloaded_block_size) - gpu_block_sizes = set(self.gpu_block_size) - assert len(gpu_block_sizes) == 1, ( - "If 'block_size' is specified in kv_connector_extra_config, " - "there must be at least one KV cache group, " - "and all groups must have the same block size." - ) - gpu_block_size = gpu_block_sizes.pop() - - assert offloaded_block_size_int % gpu_block_size == 0 - self.block_size_factor = offloaded_block_size_int // gpu_block_size + self.tokens_per_block = tuple(group.tokens_per_block for group in config.groups) + self.tokens_per_hash = config.cache.tokens_per_hash + self.blocks_per_chunk = config.cache.blocks_per_chunk @abstractmethod def get_manager(self) -> OffloadingManager: diff --git a/vllm/v1/kv_offload/config.py b/vllm/v1/kv_offload/config.py new file mode 100644 index 000000000000..cd7b3ee2075a --- /dev/null +++ b/vllm/v1/kv_offload/config.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Normalized configuration consumed by native offloading backends.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class OffloadingGroupConfig: + # Total token span covered by one block across all workers + # (accounts for context parallelism). + tokens_per_block: int + # Layer names belonging to this group. + layer_names: tuple[str, ...] + + +@dataclass(frozen=True) +class OffloadingModelConfig: + # Model identifier (e.g. HuggingFace model path). + name: str + # KV cache data type (e.g. "float16"). + dtype: str + + +@dataclass(frozen=True) +class OffloadingCacheConfig: + # Tokens per block hash. + tokens_per_hash: int + # Blocks coalesced into one offload chunk. + blocks_per_chunk: int + + +@dataclass(frozen=True) +class OffloadingParallelConfig: + # Worker index in [0, world_size). 0 on the scheduler side. + rank: int + # Total number of workers. + world_size: int + # Tensor parallel size. + tp_size: int + # Pipeline parallel size. + pp_size: int + # Prefill context parallel size. + pcp_size: int + # Decode context parallel size. + dcp_size: int + # Data parallel replica index of this engine. + data_parallel_index: int + # True when concatenating a block's data across all workers yields + # the same result regardless of the parallelism configuration. + is_parallelism_agnostic: bool + + +@dataclass(frozen=True) +class OffloadingConfig: + groups: tuple[OffloadingGroupConfig, ...] + # KV bytes stored by one worker per block. + worker_kv_bytes_per_block: int + # Whether the scheduler emits KV cache events. When true, + # the offloading backend should emit events as well. + enable_kv_cache_events: bool + # Offloading-specific configuration from kv_connector_extra_config. + extra_config: Mapping[str, Any] + # Unique identifier for this engine, distinct per DP rank. + engine_id: str + model: OffloadingModelConfig + cache: OffloadingCacheConfig + parallel: OffloadingParallelConfig diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index c8b9915a1e56..baf9a66719a0 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -72,7 +72,7 @@ class Transfer: def compute_sub_block_ptrs( block_ids: np.ndarray, - block_size_factor: int, + blocks_per_chunk: int, output: np.ndarray, tensor: torch.Tensor, skip_count: int = 0, @@ -80,38 +80,38 @@ def compute_sub_block_ptrs( """ Compute byte pointers for sub-blocks of the given block IDs. - Each block in block_ids contains block_size_factor sub-blocks. + Each block in block_ids contains blocks_per_chunk sub-blocks. The pointer for sub-block j of block b is: - base_ptr + b * row_stride + j * sub_block_size + base_ptr + b * row_stride + j * block_page_size - where sub_block_size = tensor.shape[1] // block_size_factor (gpu page size). + where block_page_size = tensor.shape[1] // blocks_per_chunk (gpu page size). - This handles tensors where row_stride != block_size_factor * sub_block_size + This handles tensors where row_stride != blocks_per_chunk * block_page_size (e.g. non-contiguous CPU tensors). Args: block_ids: array of block IDs at the tensor's native granularity. - block_size_factor: number of sub-blocks per block. + blocks_per_chunk: number of sub-blocks per block. output: pre-allocated pointer array to write pointers into. tensor: the source or destination tensor. skip_count: sub-blocks to skip in the first block. """ - assert skip_count < block_size_factor + assert skip_count < blocks_per_chunk num_sub_blocks = len(output) base_ptr = tensor.data_ptr() row_stride = tensor.stride(0) - if block_size_factor == 1: + if blocks_per_chunk == 1: # Fast path: 1:1 mapping, no sub-block expansion needed. output[:] = base_ptr + block_ids.astype(np.uint64)[:num_sub_blocks] * row_stride return - # Vectorized expansion for block_size_factor > 1. - assert tensor.shape[1] % block_size_factor == 0 - sub_block_size = tensor.shape[1] // block_size_factor - sub_offsets = np.arange(block_size_factor, dtype=np.uint64) * sub_block_size - # (num_blocks, 1) + (1, block_size_factor) -> (num_blocks, block_size_factor) + # Vectorized expansion for blocks_per_chunk > 1. + assert tensor.shape[1] % blocks_per_chunk == 0 + block_page_size = tensor.shape[1] // blocks_per_chunk + sub_offsets = np.arange(blocks_per_chunk, dtype=np.uint64) * block_page_size + # (num_blocks, 1) + (1, blocks_per_chunk) -> (num_blocks, blocks_per_chunk) all_ptrs = ( base_ptr + block_ids.astype(np.uint64)[:, np.newaxis] * row_stride ) + sub_offsets[np.newaxis, :] @@ -175,7 +175,7 @@ def __init__( self, gpu_tensors: list[torch.Tensor], cpu_tensors: list[torch.Tensor], - block_size_factor: int, + blocks_per_chunk: int, kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, mmap_region: SharedOffloadRegion | None = None, @@ -205,7 +205,7 @@ def __init__( assert cpu_tensor.device.type == "cpu" _, gpu_page_size = gpu_tensor.shape _, cpu_page_size = cpu_tensor.shape - assert cpu_page_size == gpu_page_size * block_size_factor + assert cpu_page_size == gpu_page_size * blocks_per_chunk self.src_tensors: list[torch.Tensor] = ( gpu_tensors if gpu_to_cpu else cpu_tensors @@ -220,9 +220,9 @@ def __init__( ) # GPU blocks may be smaller - # cpu_page_size = gpu_page_size * block_size_factor. - self.src_block_size_factor = 1 if self.gpu_to_cpu else block_size_factor - self.dst_block_size_factor = block_size_factor if self.gpu_to_cpu else 1 + # cpu_page_size = gpu_page_size * blocks_per_chunk. + self.src_blocks_per_chunk = 1 if self.gpu_to_cpu else blocks_per_chunk + self.dst_blocks_per_chunk = blocks_per_chunk if self.gpu_to_cpu else 1 # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it) self._mmap_region = mmap_region @@ -313,20 +313,16 @@ def transfer_async( if group_size == 0: continue - src_logical_blocks_to_skip = block_idx % self.src_block_size_factor - dst_logical_blocks_to_skip = block_idx % self.dst_block_size_factor + src_logical_blocks_to_skip = block_idx % self.src_blocks_per_chunk + dst_logical_blocks_to_skip = block_idx % self.dst_blocks_per_chunk src_logical_blocks_count = group_size + src_logical_blocks_to_skip dst_logical_blocks_count = group_size + dst_logical_blocks_to_skip - dst_blocks_count = cdiv( - dst_logical_blocks_count, self.dst_block_size_factor - ) + dst_blocks_count = cdiv(dst_logical_blocks_count, self.dst_blocks_per_chunk) dst_end_offset = dst_offset + dst_blocks_count assert dst_end_offset <= num_dst_blocks - src_blocks_count = cdiv( - src_logical_blocks_count, self.src_block_size_factor - ) + src_blocks_count = cdiv(src_logical_blocks_count, self.src_blocks_per_chunk) src_end_offset = src_offset + src_blocks_count assert src_end_offset <= num_src_blocks @@ -339,14 +335,14 @@ def transfer_async( compute_sub_block_ptrs( group_src, - self.src_block_size_factor, + self.src_blocks_per_chunk, all_src[op_idx:end_idx], self.src_tensors[t_idx], skip_count=src_logical_blocks_to_skip, ) compute_sub_block_ptrs( group_dst, - self.dst_block_size_factor, + self.dst_blocks_per_chunk, all_dst[op_idx:end_idx], self.dst_tensors[t_idx], skip_count=dst_logical_blocks_to_skip, @@ -476,7 +472,7 @@ class CPUOffloadingWorker(OffloadingWorker): def __init__( self, kv_caches: CanonicalKVCaches, - block_size_factor: int, + blocks_per_chunk: int, num_cpu_blocks: int, mmap_region: SharedOffloadRegion | None = None, ): @@ -492,7 +488,7 @@ def __init__( gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view( (-1, gpu_page_size_bytes) ) - cpu_page_size_bytes = gpu_page_size_bytes * block_size_factor + cpu_page_size_bytes = gpu_page_size_bytes * blocks_per_chunk if mmap_region is not None: cpu_tensor = mmap_region.create_next_view(cpu_page_size_bytes) @@ -518,7 +514,7 @@ def __init__( self._store_handler = SingleDirectionOffloadingHandler( gpu_tensors=gpu_tensors, cpu_tensors=cpu_tensors, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, kv_cache_groups_data_refs=kv_caches.group_data_refs, gpu_to_cpu=True, mmap_region=mmap_region, @@ -527,7 +523,7 @@ def __init__( self._load_handler = SingleDirectionOffloadingHandler( gpu_tensors=gpu_tensors, cpu_tensors=cpu_tensors, - block_size_factor=block_size_factor, + blocks_per_chunk=blocks_per_chunk, kv_cache_groups_data_refs=kv_caches.group_data_refs, gpu_to_cpu=False, ) diff --git a/vllm/v1/kv_offload/cpu/spec.py b/vllm/v1/kv_offload/cpu/spec.py index 4ad6974857cf..f6d1c29aff93 100644 --- a/vllm/v1/kv_offload/cpu/spec.py +++ b/vllm/v1/kv_offload/cpu/spec.py @@ -4,10 +4,8 @@ from typing_extensions import override -from vllm.config import VllmConfig from vllm.platforms import current_platform from vllm.utils.math_utils import round_up -from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.kv_offload.base import ( CanonicalKVCaches, OffloadingCounterMetadata, @@ -18,6 +16,7 @@ OffloadingSpec, OffloadingWorker, ) +from vllm.v1.kv_offload.config import OffloadingConfig from vllm.v1.kv_offload.cpu.common import CPUOffloadingMetrics from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager @@ -73,8 +72,8 @@ def build_metric_definitions( ) return definitions - def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): - super().__init__(vllm_config, kv_cache_config) + def __init__(self, config: OffloadingConfig): + super().__init__(config) cpu_bytes_to_use = self.extra_config.get("cpu_bytes_to_use") if not cpu_bytes_to_use: @@ -82,42 +81,28 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): "cpu_bytes_to_use must be specified in kv_connector_extra_config" ) - world_size = vllm_config.parallel_config.world_size + world_size = config.parallel.world_size self.num_blocks = 0 - self.kv_bytes_per_offloaded_block = 0 + self.kv_bytes_per_chunk = 0 self.cpu_page_size_per_worker = 0 - assert kv_cache_config is not None - if kv_cache_config.num_blocks > 0 and world_size > 0: - is_packed = any(t.block_stride for t in kv_cache_config.kv_cache_tensors) - assert not is_packed or all( - t.block_stride for t in kv_cache_config.kv_cache_tensors - ) - total_gpu_kv_bytes = ( - kv_cache_config.kv_cache_tensors[0].size - if is_packed - else sum(t.size for t in kv_cache_config.kv_cache_tensors) - ) - kv_bytes_per_block = ( - total_gpu_kv_bytes // kv_cache_config.num_blocks - ) * world_size - kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor + if config.worker_kv_bytes_per_block > 0 and world_size > 0: + kv_bytes_per_block = config.worker_kv_bytes_per_block * world_size + kv_bytes_per_chunk = kv_bytes_per_block * self.blocks_per_chunk # calculate cpu_page_size_per_worker - self.cpu_page_size_per_worker = kv_bytes_per_offloaded_block // world_size + self.cpu_page_size_per_worker = kv_bytes_per_chunk // world_size # calculate num_blocks - aligned_kv_bytes_per_offloaded_block = round_up( - kv_bytes_per_offloaded_block, self.BLOCK_SIZE_ALIGNMENT - ) - self.num_blocks = ( - int(cpu_bytes_to_use) // aligned_kv_bytes_per_offloaded_block + aligned_kv_bytes_per_chunk = round_up( + kv_bytes_per_chunk, self.BLOCK_SIZE_ALIGNMENT ) + self.num_blocks = int(cpu_bytes_to_use) // aligned_kv_bytes_per_chunk - # Expose aligned_kv_bytes_per_offloaded_block as - # kv_bytes_per_offloaded_block. Note that this might contain + # Expose aligned_kv_bytes_per_chunk as + # kv_bytes_per_chunk. Note that this might contain # some padding. i.e. each offloaded block is of the form, # |--- W0-B0---|---- W1-B0---| ... |---- Wn-B0---| *** maybe-pad *** | - self.kv_bytes_per_offloaded_block = aligned_kv_bytes_per_offloaded_block + self.kv_bytes_per_chunk = aligned_kv_bytes_per_chunk # scheduler-side self._manager: OffloadingManager | None = None @@ -150,7 +135,7 @@ def get_manager(self) -> OffloadingManager: def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: return CPUOffloadingWorker( kv_caches=kv_caches, - block_size_factor=self.block_size_factor, + blocks_per_chunk=self.blocks_per_chunk, num_cpu_blocks=self.num_blocks, ) diff --git a/vllm/v1/kv_offload/factory.py b/vllm/v1/kv_offload/factory.py index abbc9c0ede79..931fda8308fb 100644 --- a/vllm/v1/kv_offload/factory.py +++ b/vllm/v1/kv_offload/factory.py @@ -1,15 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import importlib -from collections.abc import Callable -from typing import TYPE_CHECKING +from collections.abc import Callable, Mapping +from typing import Any from vllm.logger import init_logger from vllm.v1.kv_offload.base import OffloadingSpec - -if TYPE_CHECKING: - from vllm.config import VllmConfig - from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_offload.config import OffloadingConfig logger = init_logger(__name__) @@ -30,10 +27,7 @@ def loader() -> type[OffloadingSpec]: cls._registry[name] = loader @classmethod - def get_spec_cls(cls, config: "VllmConfig") -> type[OffloadingSpec]: - kv_transfer_config = config.kv_transfer_config - assert kv_transfer_config is not None - extra_config = kv_transfer_config.kv_connector_extra_config + def get_spec_cls(cls, extra_config: Mapping[str, Any]) -> type[OffloadingSpec]: spec_name = extra_config.get("spec_name", "CPUOffloadingSpec") if spec_name in cls._registry: spec_cls = cls._registry[spec_name]() @@ -47,19 +41,11 @@ def get_spec_cls(cls, config: "VllmConfig") -> type[OffloadingSpec]: return spec_cls @classmethod - def create_spec( - cls, - config: "VllmConfig", - kv_cache_config: "KVCacheConfig", - ) -> OffloadingSpec: - kv_transfer_config = config.kv_transfer_config - assert kv_transfer_config is not None - spec_name = kv_transfer_config.kv_connector_extra_config.get( - "spec_name", "CPUOffloadingSpec" - ) - spec_cls = cls.get_spec_cls(config) + def create_spec(cls, config: OffloadingConfig) -> OffloadingSpec: + spec_name = config.extra_config.get("spec_name", "CPUOffloadingSpec") + spec_cls = cls.get_spec_cls(config.extra_config) logger.info("Creating offloading spec with name: %s", spec_name) - return spec_cls(config, kv_cache_config) + return spec_cls(config) # Register various specs here. diff --git a/vllm/v1/kv_offload/file_mapper.py b/vllm/v1/kv_offload/file_mapper.py index d8fadb09988e..b85d4d069790 100644 --- a/vllm/v1/kv_offload/file_mapper.py +++ b/vllm/v1/kv_offload/file_mapper.py @@ -4,7 +4,6 @@ import hashlib import json -from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec from vllm.v1.kv_offload.base import ( OffloadingSpec, OffloadKey, @@ -25,8 +24,8 @@ def __init__( self, root_dir: str, model_name: str, - hash_block_size: int, - gpu_blocks_per_file: int, + tokens_per_hash: int, + blocks_per_file: int, tp_size: int, pp_size: int, pcp_size: int, @@ -49,8 +48,8 @@ def __init__( self.rank: int = rank self.fields: dict = { "model_name": model_name, - "hash_block_size": hash_block_size, - "gpu_blocks_per_file": gpu_blocks_per_file, + "tokens_per_hash": tokens_per_hash, + "blocks_per_file": blocks_per_file, "tp_size": tp_size, "pp_size": pp_size, "pcp_size": pcp_size, @@ -66,47 +65,32 @@ def from_offloading_spec( cls, root_dir: str, offloading_spec: OffloadingSpec, - gpu_blocks_per_file: int = 1, + blocks_per_file: int = 1, parallel_agnostic: bool = False, ) -> "FileMapper": """Build a FileMapper from an OffloadingSpec.""" - vllm_config = offloading_spec.vllm_config - kv_cache_config = offloading_spec.kv_cache_config - - parallel_config = vllm_config.parallel_config - dtype = str(vllm_config.cache_config.cache_dtype).replace("torch.", "") + config = offloading_spec.config kv_cache_groups = [ { - "block_size": group.kv_cache_spec.block_size, + "tokens_per_block": group.tokens_per_block, "layer_names": list(group.layer_names), } - for group in kv_cache_config.kv_cache_groups + for group in config.groups ] - # Only a single full-attention group is parallelism-invariant. MLA is - # excluded: its latent KV is replicated per rank, never head-sharded. - # The V2 model runner is excluded: its KV layout is not known to be - # parallelism-invariant. - groups = kv_cache_config.kv_cache_groups - spec = groups[0].kv_cache_spec if len(groups) == 1 else None - parallel_agnostic = ( - parallel_agnostic - and not vllm_config.use_v2_model_runner - and isinstance(spec, FullAttentionSpec) - and not isinstance(spec, MLAAttentionSpec) - ) + parallel = config.parallel return cls( root_dir=root_dir, - model_name=vllm_config.model_config.model, - hash_block_size=vllm_config.cache_config.block_size, - gpu_blocks_per_file=gpu_blocks_per_file, - tp_size=parallel_config.tensor_parallel_size, - pp_size=parallel_config.pipeline_parallel_size, - pcp_size=parallel_config.prefill_context_parallel_size, - dcp_size=parallel_config.decode_context_parallel_size, - rank=parallel_config.rank, - dtype=dtype, + model_name=config.model.name, + tokens_per_hash=config.cache.tokens_per_hash, + blocks_per_file=blocks_per_file, + tp_size=parallel.tp_size, + pp_size=parallel.pp_size, + pcp_size=parallel.pcp_size, + dcp_size=parallel.dcp_size, + rank=parallel.rank, + dtype=config.model.dtype, kv_cache_groups=kv_cache_groups, - parallel_agnostic=parallel_agnostic, + parallel_agnostic=(parallel_agnostic and parallel.is_parallelism_agnostic), ) def get_file_name(self, key: OffloadKey) -> str: diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index f83113e137cb..f5614171963b 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -31,6 +31,13 @@ JobId = int +class TieringOffloadingMetrics: + """Metric names for TieringOffloadingManager.""" + + LOOKUP_SYNC_DELAY = "vllm:kv_offload_tiering_lookup_sync_delay_seconds" + LOOKUP_ASYNC_DELAY = "vllm:kv_offload_tiering_lookup_async_delay_seconds" + + @dataclass class JobMetadata: """Metadata for an in-flight async transfer job.""" diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index d8d17002856a..f12ef2c6d4d3 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -33,6 +33,7 @@ from vllm.distributed.kv_events import MEDIUM_FS from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( + Locality, LookupResult, OffloadingEvent, OffloadKey, @@ -110,11 +111,12 @@ def __init__( n_read_threads: int = 16, n_write_threads: int = 16, enable_kv_events: bool = False, + locality: str | None = None, ): """ Args: - offloading_spec: contains the vllm_config, kv_cache_config - and block_size_factor. + offloading_spec: Contains normalized offloading configuration and + blocks_per_chunk. primary_kv_view: Memoryview of the primary tier's CPU KV cache. tier_type: Tier type identifier, set by SecondaryTierFactory. root_dir: Root directory for block files. @@ -123,8 +125,11 @@ def __init__( enable_kv_events: Emit BlockStored KV events for blocks successfully stored to this tier. Effective only when KV cache events are enabled globally (kv_events_config). + locality: Whether this tier's storage is LOCAL or REMOTE relative + to the publishing vLLM instance. """ super().__init__(offloading_spec, primary_kv_view, tier_type) + self.locality = Locality(locality) if locality is not None else None self.events: list[OffloadingEvent] | None = None if enable_kv_events: @@ -150,7 +155,7 @@ def __init__( self.file_mapper = FileMapper.from_offloading_spec( root_dir=root_dir, offloading_spec=offloading_spec, - gpu_blocks_per_file=offloading_spec.block_size_factor, + blocks_per_file=offloading_spec.blocks_per_chunk, parallel_agnostic=True, ) @@ -223,7 +228,12 @@ def get_finished_jobs(self) -> Iterable[JobResult]: keys = self._store_job_keys.pop(job_id, None) if success and keys: self.events.append( - OffloadingEvent(keys=keys, medium=self.medium, removed=False) + OffloadingEvent( + keys=keys, + medium=self.medium, + removed=False, + locality=self.locality, + ) ) results.append(JobResult(job_id=job_id, success=success)) return results diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 728ac8dc86e4..2f5b52e22204 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -20,6 +20,7 @@ protecting blocks from eviction until complete_read() is called """ +import time from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass, field @@ -50,6 +51,7 @@ JobMetadata, ParentManager, SecondaryTierManager, + TieringOffloadingMetrics, ) logger = init_logger(__name__) @@ -70,6 +72,10 @@ class RequestState: pending_primary_stores: int = 0 is_finished: bool = False request_level_tiers: set[SecondaryTierManager] | None = None + sync_lookup_delay: float = 0.0 + # time.monotonic() of this request's first deferred secondary-tier lookup; + # None once consumed (observed) or while no secondary lookup is pending. + secondary_lookup_start_time: float | None = None class CPUPrimaryTierOffloadingManager(CPUOffloadingManager): @@ -212,6 +218,10 @@ def __init__( for tier in self.secondary_tiers } + # Buffers manager-level observations (e.g. lookup delay) between + # get_stats() calls; merged in and reset each time get_stats() runs. + self._stats = OffloadingConnectorStats() + def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" job_id = self._job_id_counter @@ -301,28 +311,68 @@ def lookup( # in time for a promotion this lookup may initiate. self._maybe_process_finished_jobs() + req_state = self._req_state.get(req_context.req_id) + primary_hit = self.primary_tier.lookup(key, req_context) if primary_hit is LookupResult.HIT: return LookupResult.HIT if primary_hit is LookupResult.HIT_PENDING: return LookupResult.HIT_PENDING + lookup_start = time.monotonic() any_retry = False for tier in self.secondary_tiers: if tier is exclude_tier: continue result = tier.lookup(key, req_context) if result is LookupResult.HIT: - if not self._initiate_promotion(tier, key, req_context): - return LookupResult.MISS - return LookupResult.RETRY + promoted = self._initiate_promotion(tier, key, req_context) + self._accumulate_lookup_sync_delay(req_state, lookup_start) + if ( + req_state is not None + and promoted + and req_state.secondary_lookup_start_time is None + ): + req_state.secondary_lookup_start_time = lookup_start + return LookupResult.MISS if not promoted else LookupResult.RETRY if result is LookupResult.RETRY: any_retry = True + self._accumulate_lookup_sync_delay(req_state, lookup_start) if any_retry: + if req_state is not None and req_state.secondary_lookup_start_time is None: + req_state.secondary_lookup_start_time = lookup_start return LookupResult.RETRY return LookupResult.MISS + def _accumulate_lookup_sync_delay( + self, req_state: RequestState | None, start_time: float + ) -> None: + """Accumulate secondary-tier lookup time until allocation or finish.""" + if req_state is not None: + req_state.sync_lookup_delay += time.monotonic() - start_time + + def _maybe_observe_lookup_sync_delay(self, req_state: RequestState) -> None: + delay = req_state.sync_lookup_delay + if delay == 0: + return + req_state.sync_lookup_delay = 0.0 + self._stats.observe_histogram( + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, + delay, + ) + + def _maybe_observe_lookup_async_delay(self, req_state: RequestState) -> None: + """Flush a pending deferred secondary-tier lookup timer, if any.""" + start_time = req_state.secondary_lookup_start_time + if start_time is None: + return + req_state.secondary_lookup_start_time = None + self._stats.observe_histogram( + TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY, + time.monotonic() - start_time, + ) + def _initiate_promotion( self, tier: SecondaryTierManager, @@ -666,6 +716,8 @@ def _maybe_finalize_request( if tier is exclude_tier: continue tier.on_request_finished(state.req_context) + self._maybe_observe_lookup_sync_delay(state) + self._maybe_observe_lookup_async_delay(state) del self._req_state[req_id] @override @@ -692,6 +744,13 @@ def on_schedule_end(self, context: ScheduleEndContext) -> None: for tier in self.secondary_tiers: tier.on_schedule_end(context) + for req_id in context.new_req_ids: + state = self._req_state.get(req_id) + if state is None: + continue + self._maybe_observe_lookup_sync_delay(state) + self._maybe_observe_lookup_async_delay(state) + @override def has_pending_work(self) -> bool: # In-flight primary<->secondary transfers (pending promotions are @@ -745,6 +804,8 @@ def reset_cache(self) -> None: continue for tier in self.secondary_tiers: tier.on_request_finished(state.req_context) + self._maybe_observe_lookup_sync_delay(state) + self._maybe_observe_lookup_async_delay(state) finished_req_ids.append(req_id) self.primary_tier.reset_cache() @@ -769,6 +830,13 @@ def get_stats(self) -> OffloadingConnectorStats | None: else: stats.aggregate(tier_stats) + if not self._stats.is_empty(): + if stats is None: + stats = self._stats + else: + stats.aggregate(self._stats) + self._stats = OffloadingConnectorStats() + return stats @override diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 954f86ed162f..f4af8c44c531 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -12,6 +12,7 @@ from vllm.distributed.nixl_utils import nixl_agent_config from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( + Locality, LookupResult, OffloadingEvent, OffloadKey, @@ -108,6 +109,7 @@ def __init__( prefix: str = "", io_threads: int = 4, enable_kv_events: bool = False, + locality: str | None = None, ): """ Args: @@ -120,8 +122,11 @@ def __init__( enable_kv_events: Emit BlockStored KV events for blocks successfully stored to this tier. Effective only when KV cache events are enabled globally (kv_events_config). + locality: Whether this tier's storage is LOCAL or REMOTE relative + to the publishing vLLM instance. """ super().__init__(offloading_spec, primary_kv_view, tier_type) + self.locality = Locality(locality) if locality is not None else None self.events: list[OffloadingEvent] | None = None if enable_kv_events: @@ -323,7 +328,12 @@ def get_finished_jobs(self) -> Iterable[JobResult]: keys = self._store_job_keys.pop(result.job_id, None) if result.success and keys: self.events.append( - OffloadingEvent(keys=keys, medium=self.medium, removed=False) + OffloadingEvent( + keys=keys, + medium=self.medium, + removed=False, + locality=self.locality, + ) ) return results diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 0cb37dff0e08..95605fd17358 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -127,8 +127,8 @@ def __init__( configuration reference. Args: - offloading_spec: Owning ``OffloadingSpec`` (provides - ``vllm_config`` and the offloaded block layout). + offloading_spec: Owning ``OffloadingSpec`` (provides normalized + model, parallel, and cache layout configuration). primary_kv_view: Memoryview over the CPU primary tier; the NIXL agent registers this region for RDMA transfers. tier_type: Tier identifier (defaults to ``"p2p"``). @@ -164,7 +164,7 @@ def __init__( # One control socket per DP replica: offset the base by the global # data-parallel index so replicas on a host don't collide (mirrors # NIXL). For DP=1 the index is 0, leaving the base port unchanged. - dp_index = offloading_spec.vllm_config.parallel_config.data_parallel_index + dp_index = offloading_spec.config.parallel.data_parallel_index port = int(port) + dp_index # Two decoupled identities: # _local_id (``host:port``): the ZMQ control identity that peers @@ -181,7 +181,7 @@ def __init__( config_fields = FileMapper.from_offloading_spec( root_dir="", offloading_spec=offloading_spec, - gpu_blocks_per_file=offloading_spec.block_size_factor, + blocks_per_file=offloading_spec.blocks_per_chunk, parallel_agnostic=True, ).get_run_config() self._data: DataTransport = NixlTransport( diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index 5f9e8cdc2379..964699968c9b 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -36,17 +36,18 @@ import torch from typing_extensions import override -from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.kv_offload.base import ( CanonicalKVCaches, + OffloadingHistogramMetadata, OffloadingManager, OffloadingMetricMetadata, ) +from vllm.v1.kv_offload.config import OffloadingConfig from vllm.v1.kv_offload.cpu.gpu_worker import CPUOffloadingWorker from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec +from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.manager import ( CPUPrimaryTierOffloadingManager, @@ -77,6 +78,50 @@ def build_metric_definitions( cls, extra_config: dict[str, Any] ) -> dict[str, OffloadingMetricMetadata]: metrics = super().build_metric_definitions(extra_config) + metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY] = ( + OffloadingHistogramMetadata( + documentation=( + "Histogram of total blocking time spent querying secondary " + "tiers for a request, accumulated from first lookup until " + "the request is allocated or finishes, in seconds." + ), + buckets=( + 0.00001, + 0.00005, + 0.0001, + 0.0005, + 0.001, + 0.005, + 0.01, + 0.05, + 0.1, + 0.5, + 1, + ), + ) + ) + metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY] = ( + OffloadingHistogramMetadata( + documentation=( + "Histogram of wall-clock time from a request's first deferred " + "secondary-tier lookup until the request is allocated or " + "finishes, in seconds." + ), + buckets=( + 0.0001, + 0.0005, + 0.001, + 0.005, + 0.01, + 0.05, + 0.1, + 0.5, + 1, + 5, + 10, + ), + ) + ) secondary_tier_configs = extra_config.get("secondary_tiers", []) if not isinstance(secondary_tier_configs, list): raise ValueError("secondary_tiers must be a list of tier configurations") @@ -87,8 +132,8 @@ def build_metric_definitions( metrics.update(tier_cls.build_metric_definitions(tier_config)) return metrics - def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): - super().__init__(vllm_config, kv_cache_config) + def __init__(self, config: OffloadingConfig): + super().__init__(config) # Redeclare for mypy: parent sets this but `--follow-imports skip` hides it self._manager: OffloadingManager | None = None if self.kv_events_config.self_describing_kv_events: @@ -110,10 +155,8 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig): # engine_id is unique per DP replica (suffixed with _dp{rank} in both # the Ray and multiprocessing paths), so it names a per-replica offload - # region. Non-None is guaranteed by OffloadingSpec.__init__. - assert vllm_config.kv_transfer_config is not None - assert vllm_config.kv_transfer_config.engine_id is not None - self._engine_id: str = vllm_config.kv_transfer_config.engine_id + # region. + self._engine_id = config.engine_id @override def get_manager(self) -> OffloadingManager: @@ -134,7 +177,7 @@ def get_manager(self) -> OffloadingManager: engine_id=self._engine_id, num_blocks=self.num_blocks, rank=None, - kv_bytes_per_block=self.kv_bytes_per_offloaded_block, + kv_bytes_per_block=self.kv_bytes_per_chunk, cpu_page_size=self.cpu_page_size_per_worker, ) self._scheduler_mmap = scheduler_mmap @@ -196,18 +239,18 @@ def get_manager(self) -> OffloadingManager: def create_worker(self, kv_caches: CanonicalKVCaches) -> CPUOffloadingWorker: # Fold the global physical device index into the replica-local # [0, world_size) slot range. - world_size = self.vllm_config.parallel_config.world_size + world_size = self.config.parallel.world_size rank = torch.accelerator.current_device_index() % world_size worker_mmap = SharedOffloadRegion( engine_id=self._engine_id, num_blocks=self.num_blocks, rank=rank, - kv_bytes_per_block=self.kv_bytes_per_offloaded_block, + kv_bytes_per_block=self.kv_bytes_per_chunk, cpu_page_size=self.cpu_page_size_per_worker, ) return CPUOffloadingWorker( kv_caches=kv_caches, - block_size_factor=self.block_size_factor, + blocks_per_chunk=self.blocks_per_chunk, num_cpu_blocks=self.num_blocks, mmap_region=worker_mmap, ) diff --git a/vllm/v1/metrics/loggers.py b/vllm/v1/metrics/loggers.py index 021019dc1cdc..692106cf3967 100644 --- a/vllm/v1/metrics/loggers.py +++ b/vllm/v1/metrics/loggers.py @@ -160,6 +160,42 @@ def _get_throughput(self, tracked_stats: int, now: float) -> float: def log_prefix(self): return "Engine {:03d}: ".format(self.engine_index) + def _log_prefix_for_engine(self, engine_idx: int) -> str: + if self.engine_index == engine_idx: + return self.log_prefix + return "Engine {:03d}: ".format(engine_idx) + + def _log_iteration_details( + self, scheduler_stats: SchedulerStats, engine_idx: int + ) -> None: + details = scheduler_stats.iteration_details + if details is None: + return + + encoder_msg = "" + if details.num_encoder_inputs: + encoder_msg = ( + f", encoder inputs: {details.num_encoder_inputs}, " + f"encoder output embeddings: {details.num_encoder_output_tokens}" + ) + + logger.info( + "%sIteration(%d): %d context requests, %d context tokens, " + "%d generation requests, %d generation tokens, " + "iteration elapsed time: %.2f ms%s, " + "GPU KV cache usage: %.1f%%%s", + self._log_prefix_for_engine(engine_idx), + details.iteration_index, + details.num_ctx_requests, + details.num_ctx_tokens, + details.num_generation_requests, + details.num_generation_tokens, + details.elapsed_ms, + " (dummy)" if details.is_dummy else "", + scheduler_stats.kv_cache_usage * 100, + encoder_msg, + ) + def record( self, scheduler_stats: SchedulerStats | None, @@ -172,6 +208,7 @@ def record( self._track_iteration_stats(iteration_stats) if scheduler_stats is not None: + self._log_iteration_details(scheduler_stats, engine_idx) self.prefix_caching_metrics.observe(scheduler_stats.prefix_cache_stats) if scheduler_stats.connector_prefix_cache_stats is not None: diff --git a/vllm/v1/metrics/stats.py b/vllm/v1/metrics/stats.py index a7a5fb7a2d2f..20bb3e1caa64 100644 --- a/vllm/v1/metrics/stats.py +++ b/vllm/v1/metrics/stats.py @@ -167,6 +167,21 @@ class KVCacheEvictionEvent: reuse_gaps_seconds: tuple[float, ...] +@dataclass +class SchedulerIterationDetails: + """Scheduler-side details for one engine iteration.""" + + iteration_index: int + num_ctx_requests: int + num_ctx_tokens: int + num_generation_requests: int + num_generation_tokens: int + elapsed_ms: float + num_encoder_inputs: int = 0 + num_encoder_output_tokens: int = 0 + is_dummy: bool = False + + @dataclass class SchedulerStats: """Stats associated with the scheduler.""" @@ -181,6 +196,7 @@ class SchedulerStats: current_wave: int = 0 kv_cache_usage: float = 0.0 + iteration_details: SchedulerIterationDetails | None = None prefix_cache_stats: PrefixCacheStats = field(default_factory=PrefixCacheStats) connector_prefix_cache_stats: PrefixCacheStats | None = None diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 756c5f3b3717..f8b52d079a89 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1300,6 +1300,15 @@ def _create_draft_vllm_config(self) -> VllmConfig: ), ) + if spec_cfg.kv_cache_dtype is not None: + base = replace( + base, + cache_config=replace( + base.cache_config, + cache_dtype=spec_cfg.kv_cache_dtype, + ), + ) + return base def _get_model(self) -> nn.Module: diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py index ed544bb27c1c..2de5c3be007f 100644 --- a/vllm/v1/spec_decode/ngram_proposer_gpu.py +++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py @@ -91,7 +91,8 @@ def _find_first_and_extract_all_n_parallel( suffix_indices = suffix_starts.unsqueeze(1) + torch.arange( ngram_len, device=device ) - suffix = torch.gather(token_ids, 1, suffix_indices.clamp(min=0)) + suffix_indices.clamp_(min=0) + suffix = torch.gather(token_ids, 1, suffix_indices) # Window matches for each sequence. matches = (search_windows == suffix.unsqueeze(1)).all(dim=-1) @@ -134,7 +135,7 @@ def _find_first_and_extract_all_n_parallel( draft_indices = draft_start.unsqueeze(1) + torch.arange( num_draft_tokens, device=device ) - draft_indices = draft_indices.clamp(min=0, max=max_seq_len - 1) + draft_indices.clamp_(min=0, max=max_seq_len - 1) # Extract draft tokens; gather always runs. draft_tokens = torch.gather(token_ids, 1, draft_indices) @@ -357,7 +358,8 @@ def propose( valid_write_mask & (valid_sampled_token_ids_gpu != -1) & in_bounds ) - write_positions_long = write_positions.clamp(max=max_seq_len - 1).long() + write_positions.clamp_(max=max_seq_len - 1) + write_positions_long = write_positions.long() existing_values = token_ids_gpu.gather(1, write_positions_long) tokens_cast = valid_sampled_token_ids_gpu.to(token_ids_gpu.dtype) @@ -427,7 +429,9 @@ def update_token_ids_ngram( ) # Backup last valid token before speculative tokens. - backup_indices = (num_tokens_no_spec[:num_reqs] - 1).clamp(min=0).long() + backup_indices = num_tokens_no_spec[:num_reqs] - 1 + backup_indices.clamp_(min=0) + backup_indices = backup_indices.long() backup_next_token_ids = torch.gather( token_ids_gpu[:num_reqs], dim=1, index=backup_indices.unsqueeze(1) ).squeeze(1) @@ -447,16 +451,17 @@ def update_token_ids_ngram( # Rightmost valid index per row. last_valid_indices = valid_sampled_tokens_count - 1 - last_valid_indices_safe = torch.clamp(last_valid_indices, min=0) + has_valid_sample = last_valid_indices >= 0 + last_valid_indices.clamp_(min=0) # Last valid token from each row; undefined if none. selected_tokens = torch.gather( - valid_sampled_token_ids_gpu, 1, last_valid_indices_safe.unsqueeze(1) + valid_sampled_token_ids_gpu, 1, last_valid_indices.unsqueeze(1) ).squeeze(1) # Use last token if valid; otherwise fallback to backup. next_token_ids = torch.where( - last_valid_indices != -1, + has_valid_sample, selected_tokens, backup_next_token_ids, ) diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py index 043f3f2be2bb..47821cc9b2d5 100644 --- a/vllm/v1/spec_decode/step3p5.py +++ b/vllm/v1/spec_decode/step3p5.py @@ -107,7 +107,9 @@ def _update_positions_dependent_metadata( if block_table is None: continue n_blocks = block_table.shape[1] - bn = (new_positions_1d // block_size).clamp(max=n_blocks - 1).to(torch.long) + bn = new_positions_1d // block_size + bn.clamp_(max=n_blocks - 1) + bn = bn.to(torch.long) block_ids = block_table[:batch_size].gather(1, bn.unsqueeze(1)).squeeze(1) sm = block_ids * block_size + (new_positions_1d % block_size) sm.masked_fill_(exceeds, PADDING_SLOT_ID) diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index b485d838f3e8..e17ffed93403 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -782,12 +782,16 @@ class IterationDetails: num_ctx_tokens: int num_generation_requests: int num_generation_tokens: int + num_encoder_inputs: int = 0 + num_encoder_output_tokens: int = 0 def __repr__(self) -> str: return f"IterationDetails(num_ctx_requests={self.num_ctx_requests},\ num_ctx_tokens={self.num_ctx_tokens}, \ num_generation_requests={self.num_generation_requests}, \ - num_generation_tokens={self.num_generation_tokens})" + num_generation_tokens={self.num_generation_tokens}, \ + num_encoder_inputs={self.num_encoder_inputs}, \ + num_encoder_output_tokens={self.num_encoder_output_tokens})" def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDetails: @@ -818,9 +822,18 @@ def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDet else: num_generation_requests += 1 num_generation_tokens += num_tokens + scheduled_encoder_input_stats = scheduler_output.scheduled_encoder_input_stats + num_encoder_inputs = 0 + num_encoder_output_tokens = 0 + if scheduled_encoder_input_stats is not None: + num_encoder_inputs = scheduled_encoder_input_stats.num_inputs + num_encoder_output_tokens = scheduled_encoder_input_stats.output_tokens + return IterationDetails( num_context_requests, num_context_tokens, num_generation_requests, num_generation_tokens, + num_encoder_inputs, + num_encoder_output_tokens, ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 9898fadc25e9..27fd5547257c 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -138,8 +138,6 @@ def __init__( self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {} self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {} - self._init_candidates() - # Breakable CUDA graph (PW CUDA graph without torch.compile) self.use_breakable_cg = ( is_breakable_cudagraph_enabled() @@ -147,6 +145,8 @@ def __init__( ) self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None + self._init_candidates() + def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]: """Precompute actual num_active_loras -> effective captured case. @@ -255,11 +255,11 @@ def _init_candidates(self) -> None: # for PIECEWISE graphs there is no limit on requests when replaying # i.e. no request padding is needed # so we leave it as None - num_reqs = ( - min(num_tokens, self.max_num_reqs) - if mixed_mode == CUDAGraphMode.FULL - else None - ) + num_reqs = None + if mixed_mode == CUDAGraphMode.FULL or ( + mixed_mode == CUDAGraphMode.PIECEWISE and self.use_breakable_cg + ): + num_reqs = min(num_tokens, self.max_num_reqs) desc = BatchExecutionDescriptor( cg_mode=mixed_mode, num_tokens=num_tokens, @@ -301,12 +301,10 @@ def capture( Args: create_forward_fn: Factory that prepares inputs (OUTSIDE graph) and - returns a forward_fn. For FULL cudagraph mode, it is invoked - once with warmup=True for the warmup pass, and again with - warmup=False for the captured pass. For attention backends - that perform lazy metadata initialization (e.g. FlashMLA), - FULL cudagraph capture requires distinct metadatas for warmup - and capture. + returns a forward_fn. For FULL and breakable PIECEWISE modes, + it is invoked once with warmup=True and again with warmup=False + because attention backends may mutate or lazily initialize + metadata during warmup. """ with graph_capture(device=self.device): # Capture in order: PIECEWISE first, then FULL. PIECEWISE has larger @@ -330,11 +328,17 @@ def capture( logger.debug( "CG Capture: mode=%s, batch_desc=%s", desc.cg_mode.name, desc ) - if desc.cg_mode == CUDAGraphMode.PIECEWISE: + if ( + desc.cg_mode == CUDAGraphMode.PIECEWISE + and not self.use_breakable_cg + ): forward_fn(CUDAGraphMode.PIECEWISE) else: # Capture with fresh attention state. forward_fn = create_forward_fn(desc, warmup=False) + if desc.cg_mode == CUDAGraphMode.PIECEWISE: + forward_fn(CUDAGraphMode.PIECEWISE) + continue assert desc not in self.graphs, ( f"Graph already captured for {desc}" ) @@ -489,7 +493,10 @@ def create_forward_fn( block_tables, attn_groups, kv_cache_config, - skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), + skip_attn=( + desc.cg_mode == CUDAGraphMode.PIECEWISE + and not self.use_breakable_cg + ), ) # Capture with dummy rows marked as padding. @@ -498,7 +505,7 @@ def create_forward_fn( def forward_fn(cg_mode: CUDAGraphMode) -> None: batch_descriptor = None if cg_mode == CUDAGraphMode.PIECEWISE: - assert attn_metadata is None + assert (attn_metadata is not None) == self.use_breakable_cg batch_descriptor = BatchDescriptor( num_tokens=num_tokens, has_lora=has_lora, diff --git a/vllm/v1/worker/gpu/mm/encoder_runner.py b/vllm/v1/worker/gpu/mm/encoder_runner.py index 48a3af25053f..c337ddd4aa1a 100644 --- a/vllm/v1/worker/gpu/mm/encoder_runner.py +++ b/vllm/v1/worker/gpu/mm/encoder_runner.py @@ -5,7 +5,11 @@ from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime from vllm.multimodal.inputs import MultiModalKwargsItem -from vllm.multimodal.utils import get_mm_features_in_window, group_and_batch_mm_kwargs +from vllm.multimodal.utils import ( + get_mm_features_in_window, + group_and_batch_mm_kwargs, + set_mm_embedding_modality, +) from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs @@ -136,6 +140,9 @@ def gather_mm_embeddings( else: mm_embeds_item = encoder_output[start_idx:end_idx] + # Attach modality for Omni interleaved merge (collected on demand). + set_mm_embedding_modality(mm_embeds_item, mm_feature.modality) + req_start_pos = query_start_loc[i] + start_pos - cur_query_start is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] |= ( True if is_embed is None else is_embed diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 854b71b69fcb..18eb40640add 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -5,6 +5,7 @@ import torch import torch.nn as nn +from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.v1.core.sched.output import NewRequestData @@ -139,7 +140,10 @@ def prepare_attn( kv_cache_config: KVCacheConfig, for_capture: bool = False, ) -> dict[str, Any]: - if cudagraph_mode == CUDAGraphMode.FULL: + if cudagraph_mode == CUDAGraphMode.FULL or ( + cudagraph_mode == CUDAGraphMode.PIECEWISE + and is_breakable_cudagraph_enabled() + ): # Use padded sizes - padding is handled by model_runner.prepare_attn. num_reqs = input_batch.num_reqs_after_padding num_tokens = input_batch.num_tokens_after_padding diff --git a/vllm/v1/worker/gpu/model_states/mm_pruning.py b/vllm/v1/worker/gpu/model_states/mm_pruning.py index 781baa6d15f7..8bdb85d4a17f 100644 --- a/vllm/v1/worker/gpu/model_states/mm_pruning.py +++ b/vllm/v1/worker/gpu/model_states/mm_pruning.py @@ -5,7 +5,10 @@ from vllm.config import ModelConfig from vllm.model_executor.models.interfaces import supports_multimodal_pruning -from vllm.multimodal.utils import get_mm_features_in_window +from vllm.multimodal.utils import ( + copy_mm_embedding_modality, + get_mm_features_in_window, +) from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.mm.rope import RopeState @@ -42,7 +45,11 @@ def strip(self, mm_embeds: list[torch.Tensor]) -> list[torch.Tensor]: speculator reuses the target's already-recomputed positions, hence there is no position write-back here. """ - return [mm[:, : self.inputs_embeds_size] for mm in mm_embeds] + stripped: list[torch.Tensor] = [] + for mm in mm_embeds: + out = mm[:, : self.inputs_embeds_size] + stripped.append(copy_mm_embedding_modality(mm, out)) + return stripped def recompute( self, @@ -82,7 +89,10 @@ def recompute( num_computed_tokens=num_computed, ) self.rope_state.update_prefill_positions(req_idx, new_positions, delta) - cleaned.extend(req_cleaned) + cleaned.extend( + copy_mm_embedding_modality(src, dst) + for src, dst in zip(req_embeds, req_cleaned) + ) assert pos == len(mm_embeds) return cleaned diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py index b30712b5a5d5..19919043c831 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py @@ -56,7 +56,10 @@ def create_forward_fn( block_tables, attn_groups, kv_cache_config, - skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), + skip_attn=( + desc.cg_mode == CUDAGraphMode.PIECEWISE + and not self.use_breakable_cg + ), ) return lambda cg_mode: forward_fn( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 998a28b6aee0..478274a05671 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -27,6 +27,14 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), backend=speculative_config.attention_backend, ), + cache_config=( + replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ) + if speculative_config.kv_cache_dtype is not None + else vllm_config.cache_config + ), ) with set_model_tag("dflash_head"): dflash_model = get_model( diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index bbee614a1aeb..9ea2ff0f7363 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -27,6 +27,14 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), backend=speculative_config.attention_backend, ), + cache_config=( + replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ) + if speculative_config.kv_cache_dtype is not None + else vllm_config.cache_config + ), ) with set_model_tag("dspark_head"): diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py index bdd588e5786d..579652bc7d64 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/utils.py @@ -3,7 +3,7 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig +from vllm.config import VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group from vllm.lora.layers.base import BaseLayerWithLoRA from vllm.model_executor.model_loader import get_model @@ -39,6 +39,14 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod speculative_config = vllm_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config + if speculative_config.kv_cache_dtype is not None: + vllm_config = replace( + vllm_config, + cache_config=replace( + vllm_config.cache_config, + cache_dtype=speculative_config.kv_cache_dtype, + ), + ) with set_model_tag("eagle_head"): eagle_model = get_model( vllm_config=vllm_config, model_config=draft_model_config diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index b104e4c5cae0..79a905adada6 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -105,7 +105,12 @@ MultiModalKwargsItem, PlaceholderRange, ) -from vllm.multimodal.utils import get_mm_features_in_window, group_and_batch_mm_kwargs +from vllm.multimodal.utils import ( + copy_mm_embedding_modality, + get_mm_features_in_window, + group_and_batch_mm_kwargs, + set_mm_embedding_modality, +) from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingType @@ -120,6 +125,7 @@ from vllm.utils.torch_utils import ( PIN_MEMORY, async_tensor_h2d, + current_stream, get_dtype_size, is_quantized_kv_cache, kv_cache_dtype_str_to_dtype, @@ -3244,11 +3250,13 @@ def _gather_mm_embeddings( is_mm_embed[ req_start_pos + start_idx : req_start_pos + end_idx ] |= is_embed + set_mm_embedding_modality(mm_embeds_item, mm_feature.modality) mm_embeds_req.append(mm_embeds_item) if self.is_multimodal_pruning_enabled and self.uses_mrope: assert req_state.mrope_positions is not None should_sync_mrope_positions = True + old_mm_embeds_req = mm_embeds_req mm_embeds_req, new_mrope_positions, new_delta = ( self.model.recompute_mrope_positions( input_ids=req_state.prompt_token_ids, @@ -3257,6 +3265,10 @@ def _gather_mm_embeddings( num_computed_tokens=req_state.num_computed_tokens, ) ) + mm_embeds_req = [ + copy_mm_embedding_modality(src, dst) + for src, dst in zip(old_mm_embeds_req, mm_embeds_req) + ] req_state.mrope_positions.copy_(new_mrope_positions) req_state.mrope_position_delta = new_delta @@ -3436,6 +3448,7 @@ def _pool( ) if raw_pooler_output is None or not any(finished_mask): + self._sync_device() model_runner_output.pooler_output = [None] * num_reqs return model_runner_output @@ -6608,17 +6621,19 @@ def profile_cudagraph_memory(self) -> int: per_graph_estimate = {} encoder_memory_estimate = 0 - # On ROCm, capture these throwaway profiling graphs on the current stream - # instead of the fresh side stream graph_capture() allocates by default. - # torch's allocator pools free blocks per stream, so a side-stream forward - # strands a persistent aiter scratch buffer in a separate pool, shifting - # the physical placement of the real KV cache allocated afterward and - # slowing bandwidth-bound decode ~20%. The graphs are discarded, so a - # side stream is unnecessary here. - # cap_ctx=None keeps the side-stream path on CUDA, where the current - # stream is the legacy default stream, on which capture cannot begin. + # On ROCm, capture these throwaway profiling graphs on vLLM's dedicated + # compute stream instead of the fresh side stream graph_capture() + # allocates by default. torch's allocator pools free blocks per stream, + # so a side-stream forward strands a persistent aiter scratch buffer in + # a separate pool, shifting the physical placement of the real KV cache + # allocated afterward and slowing bandwidth-bound decode ~20%. The + # graphs are discarded, so a side stream is unnecessary here. + # Use current_stream(), not torch.cuda.current_stream(): before vLLM + # initializes its dedicated stream, torch returns the per-thread default + # stream (cuda_stream=0), which cannot be used for cudagraph capture. + # cap_ctx=None keeps the side-stream path on CUDA. cap_ctx = ( - GraphCaptureContext(torch.cuda.current_stream(self.device)) + GraphCaptureContext(current_stream()) if current_platform.is_rocm() else None ) diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index bef811f6e61d..fc4948d0e9ff 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -388,7 +388,7 @@ def flash_attn_varlen_func( from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd - out, softmax_lse = _flash_attn_fwd( + out, softmax_lse, _, _ = _flash_attn_fwd( q, k, v,