forked from Karylab-cklius/vllm
Compare commits
51
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a02155c787 | ||
|
|
97a98006b0 | ||
|
|
0a684ab0c0 | ||
|
|
0d9210a502 | ||
|
|
1d874867ea | ||
|
|
2e2e626b40 | ||
|
|
af91f4b3e4 | ||
|
|
2396a61108 | ||
|
|
97a668152b | ||
|
|
58b2012aa2 | ||
|
|
b7c20d0cfa | ||
|
|
a2b1f9fc3b | ||
|
|
642076d26c | ||
|
|
5feb3950e5 | ||
|
|
4ec199b66a | ||
|
|
7ca017778f | ||
|
|
fbfe58133d | ||
|
|
9dd62d80ab | ||
|
|
f878367898 | ||
|
|
bd091079cb | ||
|
|
b23bd73f54 | ||
|
|
e2d7adeb64 | ||
|
|
15cb8e140d | ||
|
|
f007cceb42 | ||
|
|
0a5069e4e3 | ||
|
|
8ce53a616e | ||
|
|
ae10e855ab | ||
|
|
530ee36a0d | ||
|
|
d835ad572c | ||
|
|
47d0597ca2 | ||
|
|
818cf61e91 | ||
|
|
c01618fdc8 | ||
|
|
823eaf667d | ||
|
|
f1f1259692 | ||
|
|
df13b5aef5 | ||
|
|
4938d44a3b | ||
|
|
37bf988c2f | ||
|
|
9459fc6471 | ||
|
|
5245c80564 | ||
|
|
9bc266d923 | ||
|
|
5c9f6557d7 | ||
|
|
319db65b68 | ||
|
|
83a7669827 | ||
|
|
401bed48ad | ||
|
|
d492d1e697 | ||
|
|
f37e113590 | ||
|
|
a10e369f06 | ||
|
|
aa3f2efe42 | ||
|
|
a32d1bff95 | ||
|
|
ee47a21fcd | ||
|
|
0f471a3088 |
@@ -18,6 +18,8 @@ steps:
|
||||
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
- tests/kernels/mamba/test_cpu_short_conv.py
|
||||
- tests/kernels/mamba/test_causal_conv1d.py
|
||||
- tests/kernels/mamba/test_mamba_ssm.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
@@ -28,7 +30,9 @@ steps:
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
|
||||
|
||||
# Note: SDE can't be downloaded from CI host because of AWS WAF
|
||||
# - label: CPU-Compatibility Tests
|
||||
|
||||
@@ -7,9 +7,6 @@
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# The macmini queue uses persistent checkouts, so refresh tags for setuptools-scm.
|
||||
git fetch --tags --force origin
|
||||
|
||||
# The Rust frontend build needs protoc.
|
||||
if ! command -v protoc >/dev/null 2>&1; then
|
||||
brew install protobuf
|
||||
|
||||
@@ -40,7 +40,9 @@ function cpu_tests() {
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
|
||||
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
|
||||
|
||||
# skip tests requiring model downloads if HF_TOKEN is not set
|
||||
# due to rate-limits
|
||||
@@ -97,3 +99,4 @@ function cpu_tests() {
|
||||
# All of CPU tests are expected to be finished less than 40 mins.
|
||||
export -f cpu_tests
|
||||
timeout 2h bash -c cpu_tests
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Basic Correctness
|
||||
key: basic-correctness
|
||||
timeout_in_minutes: 68
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Benchmarks CLI Test
|
||||
key: benchmarks-cli-test
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -51,7 +51,7 @@ steps:
|
||||
|
||||
- label: e2e Scheduling (1 GPU)
|
||||
key: e2e-scheduling-1-gpu
|
||||
timeout_in_minutes: 53
|
||||
timeout_in_minutes: 35
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
|
||||
@@ -39,7 +39,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server)
|
||||
key: entrypoints-integration-api-server
|
||||
device: h200_35gb
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -59,7 +59,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 68
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -78,7 +78,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 83
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -156,7 +156,7 @@ steps:
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
key: v1-sample-logits
|
||||
timeout_in_minutes: 83
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
|
||||
@@ -5,12 +5,14 @@ steps:
|
||||
- label: Model Executor
|
||||
device: h200_35gb
|
||||
key: model-executor
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- vllm/model_executor/warmup
|
||||
- tests/model_executor
|
||||
- tests/model_executor/test_jit_warmup.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
@@ -34,7 +36,9 @@ steps:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- vllm/model_executor/warmup
|
||||
- tests/model_executor
|
||||
- tests/model_executor/test_jit_warmup.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
@@ -137,7 +137,7 @@ steps:
|
||||
|
||||
- label: Language Models Test (MTEB)
|
||||
key: language-models-test-mteb
|
||||
timeout_in_minutes: 68
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2"
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
timeout_in_minutes: 68
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -20,7 +20,7 @@ steps:
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -54,7 +54,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
device: h200_35gb
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 75
|
||||
timeout_in_minutes: 50
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
@@ -85,7 +85,7 @@ steps:
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
key: multi-modal-processor
|
||||
timeout_in_minutes: 98
|
||||
timeout_in_minutes: 65
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: PyTorch Compilation Unit Tests
|
||||
device: h200_35gb
|
||||
key: pytorch-compilation-unit-tests
|
||||
timeout_in_minutes: 150
|
||||
timeout_in_minutes: 90
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
|
||||
+86
-42
@@ -227,6 +227,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
cuda_archs_loose_intersection(CUDA_ARCHS
|
||||
"${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
|
||||
message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
|
||||
|
||||
set(VLLM_COMPILED_CUDA_ARCHS)
|
||||
foreach(_ARCH ${CUDA_ARCHS})
|
||||
set(_COMPILED_ARCH "${_ARCH}")
|
||||
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
|
||||
if(_ARCH MATCHES "^(10|11|12)\\.0$")
|
||||
set(_COMPILED_ARCH "${_ARCH}f")
|
||||
endif()
|
||||
elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
|
||||
if(_ARCH MATCHES "^(10\\.(0|1|3)|12\\.(0|1))$")
|
||||
set(_COMPILED_ARCH "${_ARCH}a")
|
||||
endif()
|
||||
endif()
|
||||
list(APPEND VLLM_COMPILED_CUDA_ARCHS "${_COMPILED_ARCH}")
|
||||
endforeach()
|
||||
list(JOIN VLLM_COMPILED_CUDA_ARCHS "," VLLM_COMPILED_CUDA_ARCHS_STR)
|
||||
else()
|
||||
#
|
||||
# For other GPU targets override the GPU architectures detected by cmake/torch
|
||||
@@ -385,8 +401,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
#
|
||||
# _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
|
||||
#
|
||||
# Shared entry sources are part of the base extension source list, but some
|
||||
# optional kernel families below append source-local feature macros to them.
|
||||
set(VLLM_STABLE_TORCH_BINDINGS_SRC
|
||||
"csrc/libtorch_stable/torch_bindings.cpp")
|
||||
set(CACHE_KERNELS_SRC "csrc/libtorch_stable/cache_kernels.cu")
|
||||
set(SCALED_MM_ENTRY_SRC
|
||||
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
|
||||
set(NVFP4_QUANT_ENTRY_SRC
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu")
|
||||
set(NVFP4_SCALED_MM_ENTRY_SRC
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
|
||||
|
||||
set(VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/torch_bindings.cpp"
|
||||
"${VLLM_STABLE_TORCH_BINDINGS_SRC}"
|
||||
"csrc/libtorch_stable/cuda_view.cu"
|
||||
"csrc/libtorch_stable/cuda_utils_kernels.cu"
|
||||
"csrc/libtorch_stable/activation_kernels.cu"
|
||||
@@ -409,7 +437,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/sampler.cu"
|
||||
"csrc/libtorch_stable/topk.cu"
|
||||
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
|
||||
"csrc/libtorch_stable/cache_kernels.cu"
|
||||
"${CACHE_KERNELS_SRC}"
|
||||
"csrc/libtorch_stable/cache_kernels_fused.cu"
|
||||
"csrc/libtorch_stable/custom_all_reduce.cu"
|
||||
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
|
||||
@@ -425,11 +453,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
|
||||
"9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
|
||||
if(COOPERATIVE_TOPK_ARCHS)
|
||||
list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
|
||||
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
@@ -467,11 +490,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/cutlass_extensions/common.cpp"
|
||||
"csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
|
||||
"${SCALED_MM_ENTRY_SRC}"
|
||||
"${NVFP4_QUANT_ENTRY_SRC}"
|
||||
"${NVFP4_SCALED_MM_ENTRY_SRC}"
|
||||
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
|
||||
"csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC}"
|
||||
DEFINITIONS VLLM_COMPILED_CUDA_ARCHS=\"${VLLM_COMPILED_CUDA_ARCHS_STR}\")
|
||||
|
||||
#
|
||||
# Machete kernels
|
||||
@@ -547,11 +573,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
CUDA_ARCHS "${CUDA_ARCHS}")
|
||||
|
||||
if(COOPERATIVE_TOPK_ARCHS)
|
||||
set(COOPERATIVE_TOPK_SRC "csrc/libtorch_stable/cooperative_topk.cu")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/cooperative_topk.cu")
|
||||
"${COOPERATIVE_TOPK_SRC}")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "csrc/libtorch_stable/cooperative_topk.cu"
|
||||
SRCS "${COOPERATIVE_TOPK_SRC}"
|
||||
CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC};${COOPERATIVE_TOPK_SRC}"
|
||||
DEFINITIONS VLLM_ENABLE_COOPERATIVE_TOPK=1)
|
||||
endif()
|
||||
|
||||
# Only build Marlin kernels if we are building for at least some compatible archs.
|
||||
@@ -760,8 +790,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SCALED_MM_SM90_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM90_SRCS}"
|
||||
DEFINITIONS ENABLE_SCALED_MM_SM90=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM90_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1")
|
||||
# Let scaled_mm_c2x know it doesn't need to build these arches
|
||||
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
|
||||
message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
|
||||
@@ -794,8 +826,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SCALED_MM_SM120_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM120_SRCS}"
|
||||
DEFINITIONS ENABLE_SCALED_MM_SM120=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM120_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
|
||||
# Let scaled_mm_c2x know it doesn't need to build these arches
|
||||
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
|
||||
message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
|
||||
@@ -828,8 +862,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SCALED_MM_SM100_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM100_SRCS}"
|
||||
DEFINITIONS ENABLE_SCALED_MM_SM100=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM100_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1")
|
||||
# Let scaled_mm_c2x know it doesn't need to build these arches
|
||||
list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
|
||||
message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}")
|
||||
@@ -858,8 +894,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SCALED_MM_C2X_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_C2X_SRCS}"
|
||||
DEFINITIONS ENABLE_SCALED_MM_C2X=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_C2X_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
|
||||
message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
|
||||
else()
|
||||
if (SCALED_MM_3X_ARCHS)
|
||||
@@ -884,8 +922,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${CUTLASS_MOE_SM90_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM90_SRCS}"
|
||||
DEFINITIONS ENABLE_CUTLASS_MOE_SM90=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM90_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
|
||||
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
|
||||
else()
|
||||
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
|
||||
@@ -908,8 +948,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${CUTLASS_MOE_SM100_SRCS}"
|
||||
CUDA_ARCHS "${SCALED_MM_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM100_SRCS}"
|
||||
DEFINITIONS ENABLE_CUTLASS_MOE_SM10X_OR_SM11X=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM100_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
|
||||
# The implementation is named sm100 historically, but it is built for the
|
||||
# SM10x/SM11x family: CUDA 12 Thor reports SM101, CUDA 13 Thor reports
|
||||
# SM110. Keep the compile-time macro aligned with runtime dispatch.
|
||||
message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
|
||||
else()
|
||||
if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
|
||||
@@ -950,10 +995,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
|
||||
#
|
||||
|
||||
# SM12x FP4 kernels. These share some generic NVFP4 quantization entry
|
||||
# sources with the SM10x/11x block below; set_gencode_flags_for_srcs appends
|
||||
# per-source flags, so shared files accumulate both SM12x and SM10x/11x
|
||||
# gencodes when both families are requested.
|
||||
# Shared FP4 implementation sources live next to the FP4 arch logic because
|
||||
# both SM12x and SM10x/11x append family-specific gencodes and feature macros
|
||||
# to them.
|
||||
set(FP4_SHARED_SRCS
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
|
||||
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
@@ -961,18 +1012,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS)
|
||||
set(FP4_SM120_SRCS
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
|
||||
${FP4_SHARED_SRCS}
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
|
||||
)
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${FP4_SM120_SRCS}"
|
||||
CUDA_ARCHS "${FP4_SM120_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM120_SRCS}"
|
||||
DEFINITIONS ENABLE_NVFP4_SM120=1)
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${SCALED_MM_ENTRY_SRC}"
|
||||
DEFINITIONS ENABLE_CUTLASS_MOE_SM120=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
|
||||
message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.")
|
||||
@@ -987,14 +1039,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS)
|
||||
set(FP4_SM100_SRCS
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
|
||||
${FP4_SHARED_SRCS}
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu"
|
||||
"csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
|
||||
"csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
|
||||
if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
|
||||
message(STATUS
|
||||
"Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).")
|
||||
@@ -1002,9 +1050,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${FP4_SM100_SRCS}"
|
||||
CUDA_ARCHS "${FP4_SM100_ARCHS}")
|
||||
set_compile_definitions_for_srcs(
|
||||
SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM100_SRCS}"
|
||||
DEFINITIONS ENABLE_NVFP4_SM100=1)
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
|
||||
message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building SM10x/11x NVFP4/MXFP4 as no compatible archs were found.")
|
||||
@@ -1058,7 +1107,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
SRCS "${CUTLASS_MLA_SRCS}"
|
||||
CUDA_ARCHS "${MLA_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MLA_SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
|
||||
# Add MLA-specific include directories only to MLA source files
|
||||
set_source_files_properties(${CUTLASS_MLA_SRCS}
|
||||
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
|
||||
@@ -1106,10 +1154,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
# Needed to use cuda/hip APIs from C-shim
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
|
||||
if(COOPERATIVE_TOPK_ARCHS)
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
VLLM_ENABLE_COOPERATIVE_TOPK=1)
|
||||
endif()
|
||||
# Needed by CUTLASS kernels
|
||||
target_compile_definitions(_C_stable_libtorch PRIVATE
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
|
||||
@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
|
||||
"csrc/cpu/layernorm.cpp"
|
||||
"csrc/cpu/mla_decode.cpp"
|
||||
"csrc/cpu/pos_encoding.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp")
|
||||
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp"
|
||||
# TODO: Remove these files
|
||||
"csrc/cpu/activation.cpp"
|
||||
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/utils.cpp"
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp"
|
||||
# TODO: Remove these files
|
||||
|
||||
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
|
||||
set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
|
||||
else()
|
||||
set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
|
||||
set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
|
||||
set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
|
||||
|
||||
set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
|
||||
if(NOT _qutlass_fc_root)
|
||||
@@ -125,8 +125,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
CUDA_ARCHS "${QUTLASS_ARCHS}"
|
||||
)
|
||||
|
||||
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
|
||||
# Keep it as its own extension (registers torch.ops._qutlass_C).
|
||||
define_extension_target(
|
||||
_qutlass_C
|
||||
DESTINATION vllm
|
||||
@@ -139,9 +137,11 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
WITH_SOABI)
|
||||
|
||||
target_compile_definitions(_qutlass_C PRIVATE
|
||||
QUTLASS_DISABLE_PYBIND=1
|
||||
QUTLASS_MINIMAL_BUILD=1
|
||||
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL
|
||||
USE_CUDA)
|
||||
|
||||
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
|
||||
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
@@ -344,6 +344,28 @@ macro(set_gencode_flags_for_srcs)
|
||||
endif()
|
||||
endmacro()
|
||||
|
||||
#
|
||||
# For a list of source files append preprocessor definitions to file-specific
|
||||
# compile options. Use this for optional kernel feature macros so toggling one
|
||||
# kernel family does not perturb the compile command for every source in the
|
||||
# extension target.
|
||||
#
|
||||
macro(set_compile_definitions_for_srcs)
|
||||
set(options)
|
||||
set(oneValueArgs)
|
||||
set(multiValueArgs SRCS DEFINITIONS)
|
||||
cmake_parse_arguments(arg "${options}" "${oneValueArgs}"
|
||||
"${multiValueArgs}" ${ARGN})
|
||||
|
||||
foreach(_DEF ${arg_DEFINITIONS})
|
||||
set_property(
|
||||
SOURCE ${arg_SRCS}
|
||||
APPEND PROPERTY
|
||||
COMPILE_DEFINITIONS "${_DEF}"
|
||||
)
|
||||
endforeach()
|
||||
endmacro()
|
||||
|
||||
#
|
||||
# For the given `SRC_CUDA_ARCHS` list of gencode versions in the form
|
||||
# `<major>.<minor>[letter]` compute the "loose intersection" with the
|
||||
|
||||
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
reg.val[1] = fp16_to_fp32_bits(raw_lo);
|
||||
}
|
||||
float reduce_sum() const {
|
||||
AliasReg ar;
|
||||
ar.reg = reg;
|
||||
float result = 0;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&result, &ar](int i) { result += ar.values[i]; });
|
||||
|
||||
return result;
|
||||
// VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
|
||||
// Step 1: pairwise sum of the two 4-wide halves
|
||||
__vector float s = vec_add(reg.val[0], reg.val[1]);
|
||||
// Step 2: rotate by 8 bytes (2 floats) and add
|
||||
s = vec_add(s, vec_sld(s, s, 8));
|
||||
// Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total
|
||||
s = vec_add(s, vec_sld(s, s, 4));
|
||||
return vec_extract(s, 0);
|
||||
}
|
||||
FP32Vec8 exp() const {
|
||||
f32x4x2_t out;
|
||||
|
||||
@@ -0,0 +1,285 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
|
||||
// mamba_kernels.hpp.
|
||||
|
||||
#include "cpu/mamba_kernels.hpp"
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <torch/library.h>
|
||||
#include <c10/util/Optional.h>
|
||||
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// causal_conv1d_update
|
||||
// ---------------------------------------------------------------------------
|
||||
at::Tensor causal_conv1d_update_cpu_impl(
|
||||
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
|
||||
const c10::optional<at::Tensor>& bias,
|
||||
const c10::optional<std::string>& activation,
|
||||
const c10::optional<at::Tensor>& conv_state_indices,
|
||||
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
|
||||
bool do_silu = false;
|
||||
if (activation.has_value()) {
|
||||
const std::string& act = activation.value();
|
||||
do_silu = (act == "silu" || act == "swish");
|
||||
}
|
||||
|
||||
at::ScalarType dtype = x.scalar_type();
|
||||
|
||||
// Input x: contiguous in native dtype.
|
||||
at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
|
||||
|
||||
// conv_state: NEVER copy the full paged tensor just for layout reasons.
|
||||
// If the dtype matches we work directly on conv_state (contiguous or not)
|
||||
// by extracting strides and passing them to the kernel.
|
||||
// Only a dtype-conversion copy is made when types differ (rare for BF16).
|
||||
bool state_type_ok = (conv_state.scalar_type() == dtype);
|
||||
at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
|
||||
// state_c and conv_state may be non-contiguous — that is intentional.
|
||||
|
||||
// Weight: coerce to same dtype if needed (should match in practice)
|
||||
at::Tensor w_c =
|
||||
(weight.scalar_type() != dtype)
|
||||
? weight.to(dtype).contiguous()
|
||||
: (weight.is_contiguous() ? weight : weight.contiguous());
|
||||
|
||||
// Bias stays float32 (small scalar, used only for fp32 accumulation)
|
||||
at::Tensor bias_f32;
|
||||
if (bias.has_value() && bias.value().defined())
|
||||
bias_f32 = bias.value().to(at::kFloat).contiguous();
|
||||
|
||||
int64_t batch = x_c.size(0);
|
||||
int64_t dim = x_c.size(1);
|
||||
int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
|
||||
int64_t width = w_c.size(1);
|
||||
int64_t state_len = state_c.size(2);
|
||||
|
||||
// Extract strides — works for contiguous AND non-contiguous (transposed)
|
||||
// state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
|
||||
// contiguous) stride(1): between conv channels (dim stride) stride(2):
|
||||
// between state elements (=1 when contiguous, =dim when transposed)
|
||||
int64_t stride_s_slot = state_c.stride(0);
|
||||
int64_t stride_s_dim = state_c.stride(1);
|
||||
int64_t stride_s_state = state_c.stride(2);
|
||||
|
||||
at::Tensor out = x_c.clone(); // native dtype, no float32 alloc
|
||||
|
||||
const int32_t* cache_idx_ptr = nullptr;
|
||||
at::Tensor cache_idx_int;
|
||||
if (conv_state_indices.has_value()) {
|
||||
cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
|
||||
cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
|
||||
mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
|
||||
x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
|
||||
stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
|
||||
bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
|
||||
out.data_ptr<scalar_t>(), cache_idx_ptr,
|
||||
static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
|
||||
do_silu);
|
||||
});
|
||||
|
||||
// Write back only when a type-conversion copy was made.
|
||||
// Layout-only non-contiguity is handled via strides above — no copy needed.
|
||||
if (!state_type_ok) conv_state.copy_(state_c);
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// selective_state_update
|
||||
// ---------------------------------------------------------------------------
|
||||
void selective_state_update_cpu_impl(
|
||||
at::Tensor& state, // (nstates, nheads, dim, dstate)
|
||||
const at::Tensor& x, // (N, nheads, dim)
|
||||
const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
|
||||
const at::Tensor& C, const c10::optional<at::Tensor>& D,
|
||||
const c10::optional<at::Tensor>& z,
|
||||
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
|
||||
const c10::optional<at::Tensor>& state_batch_indices,
|
||||
const c10::optional<at::Tensor>& dst_state_batch_indices,
|
||||
int64_t null_block_id, at::Tensor& out,
|
||||
const c10::optional<at::Tensor>& num_accepted_tokens,
|
||||
const c10::optional<at::Tensor>& cu_seqlens) {
|
||||
at::ScalarType state_type = state.scalar_type();
|
||||
at::ScalarType input_type = x.scalar_type();
|
||||
|
||||
// x, B, C must be contiguous and match input_type
|
||||
auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor x_in = ensure_input(x);
|
||||
at::Tensor B_in = ensure_input(B);
|
||||
at::Tensor C_in = ensure_input(C);
|
||||
at::Tensor z_in;
|
||||
if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
|
||||
|
||||
// A, D, dt_bias are float32 model parameters that arrive here as expanded
|
||||
// tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
|
||||
// We need just the scalar value per head as a (nheads,) 1-D array so that
|
||||
// A_ptr[h] in the kernel correctly reads head h's value.
|
||||
//
|
||||
// Strategy: peel trailing expanded (stride=0) dims via .select(), which is
|
||||
// a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0)
|
||||
// → .select(2,0) → (nheads, head_dim) strides (1,0)
|
||||
// → .select(1,0) → (nheads,) stride (1,) ← contiguous, free.
|
||||
// No allocation, no type conversion (A is already float32).
|
||||
auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = t;
|
||||
// Peel trailing dimensions that are broadcast (stride=0 or size=1)
|
||||
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
|
||||
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
|
||||
at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32
|
||||
at::Tensor D_f32, dt_bias_f32;
|
||||
if (D.has_value() && D.value().defined())
|
||||
D_f32 = to_per_head_1d_f32(D.value());
|
||||
if (dt_bias.has_value() && dt_bias.value().defined())
|
||||
dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
|
||||
|
||||
// dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
|
||||
// the type conversion so we convert head_dim x fewer elements.
|
||||
at::Tensor dt_f32;
|
||||
{
|
||||
// If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
|
||||
// take a zero-copy view of index 0 along that dim first.
|
||||
at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads)
|
||||
at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
|
||||
dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
|
||||
}
|
||||
|
||||
int64_t nheads = state.size(1);
|
||||
int64_t dim = state.size(2);
|
||||
int64_t dstate = state.size(3);
|
||||
int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
|
||||
? cu_seqlens.value().size(0) - 1
|
||||
: x_in.size(0);
|
||||
int64_t ngroups = B_in.size(1);
|
||||
|
||||
// Strides
|
||||
int64_t stride_state_n = state.stride(0);
|
||||
int64_t stride_state_h = state.stride(1);
|
||||
int64_t stride_state_d = state.stride(2);
|
||||
int64_t stride_x_n = x_in.stride(0);
|
||||
int64_t stride_x_h = x_in.stride(1);
|
||||
int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads)
|
||||
int64_t stride_BC_n = B_in.stride(0);
|
||||
int64_t stride_BC_g = B_in.stride(1);
|
||||
int64_t stride_out_n = out.stride(0);
|
||||
int64_t stride_out_h = out.stride(1);
|
||||
|
||||
// Optional index pointers
|
||||
auto get_int32_ptr =
|
||||
[](const c10::optional<at::Tensor>& opt) -> const int32_t* {
|
||||
return (opt.has_value() && opt.value().defined())
|
||||
? opt.value().data_ptr<int32_t>()
|
||||
: nullptr;
|
||||
};
|
||||
const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
|
||||
const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
|
||||
const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
|
||||
const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
|
||||
|
||||
// Dispatch on (state_t, input_t, out_t): write directly into `out`
|
||||
// without any intermediate float32 buffer.
|
||||
VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
|
||||
using state_t = scalar_t;
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
|
||||
using input_t = scalar_t;
|
||||
VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
|
||||
using out_t = scalar_t;
|
||||
mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
|
||||
state.data_ptr<state_t>(), stride_state_n, stride_state_h,
|
||||
stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
|
||||
dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
|
||||
B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
|
||||
stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
|
||||
z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
|
||||
dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
|
||||
out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
|
||||
dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
|
||||
nheads, ngroups, dim, dstate, dt_softplus);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mamba_chunk_scan_fwd_cpu
|
||||
// ---------------------------------------------------------------------------
|
||||
void mamba_chunk_scan_fwd_cpu_impl(
|
||||
at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller
|
||||
at::Tensor&
|
||||
final_states, // [batch, nheads, headdim, dstate] float32 contiguous
|
||||
const at::Tensor& x, // [seqlen, nheads, headdim]
|
||||
const at::Tensor&
|
||||
dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
|
||||
const at::Tensor& A, // [nheads] float32
|
||||
const at::Tensor& B, // [seqlen, ngroups, dstate]
|
||||
const at::Tensor& C, // [seqlen, ngroups, dstate]
|
||||
const c10::optional<at::Tensor>& D, // [nheads] float32 (optional)
|
||||
const c10::optional<at::Tensor>& z, // [seqlen, nheads, headdim] (optional)
|
||||
const at::Tensor& cu_seqlens // [batch+1] int32
|
||||
) {
|
||||
const at::ScalarType input_type = x.scalar_type();
|
||||
|
||||
auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor x_in = ensure_contig(x);
|
||||
at::Tensor B_in = ensure_contig(B);
|
||||
at::Tensor C_in = ensure_contig(C);
|
||||
at::Tensor z_in;
|
||||
if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
|
||||
|
||||
// A and D are float32 model parameters, potentially broadcast-expanded.
|
||||
// Strip trailing broadcast dims to get a contiguous (nheads,) array.
|
||||
auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = t;
|
||||
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
|
||||
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor A_f32 = to_per_head_f32(A);
|
||||
at::Tensor D_f32;
|
||||
if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
|
||||
|
||||
// dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
|
||||
// Python.
|
||||
at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
|
||||
if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
|
||||
|
||||
at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
|
||||
|
||||
const int64_t batch = final_states.size(0);
|
||||
const int64_t nheads = final_states.size(1);
|
||||
const int64_t headdim = final_states.size(2);
|
||||
const int64_t dstate = final_states.size(3);
|
||||
const int64_t ngroups = B_in.size(1);
|
||||
|
||||
TORCH_CHECK(final_states.is_contiguous(),
|
||||
"mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
|
||||
TORCH_CHECK(out.is_contiguous(),
|
||||
"mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
|
||||
"raw data_ptr)");
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
|
||||
mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
|
||||
final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
|
||||
dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
|
||||
B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
|
||||
D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
|
||||
z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
|
||||
out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
|
||||
ngroups, headdim, dstate);
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,382 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Fused CPU vector kernels for Mamba decode-step hotspots:
|
||||
// - causal_conv1d_update (depthwise 1-D conv state roll + compute)
|
||||
// - selective_state_update (SSM recurrence, single-step)
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cpu_types.hpp"
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <cstdint>
|
||||
#include <algorithm>
|
||||
|
||||
namespace mamba_cpu {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// causal_conv1d_update — templated for native BF16/FP32
|
||||
//
|
||||
// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
|
||||
// Explicit strides are passed so the kernel writes directly into the
|
||||
// correct memory locations without making a contiguous copy of the full
|
||||
// paged tensor (which was the source of the 34-41% direct_copy_kernel).
|
||||
//
|
||||
// stride_s_slot = state.stride(0) — between cache slots
|
||||
// stride_s_dim = state.stride(1) — between conv_dim channels
|
||||
// stride_s_state = state.stride(2) — between state elements
|
||||
//
|
||||
// When stride_s_state == 1 (contiguous), the memmove fast path is used.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename scalar_t>
|
||||
inline void causal_conv1d_update_kernel(
|
||||
const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
|
||||
int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
|
||||
const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
|
||||
scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
|
||||
int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
|
||||
int64_t width, int64_t state_len, bool do_silu) {
|
||||
#pragma omp parallel for
|
||||
for (int64_t b = 0; b < batch; ++b) {
|
||||
int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
|
||||
if (cache_idx == pad_slot_id) continue;
|
||||
|
||||
for (int64_t t = 0; t < seqlen; ++t) {
|
||||
const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
|
||||
scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
|
||||
// Base of this slot in the (possibly non-contiguous) paged state
|
||||
scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
|
||||
|
||||
for (int64_t d = 0; d < dim; ++d) {
|
||||
float x_val = static_cast<float>(x_b[d * seqlen]);
|
||||
scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state
|
||||
const scalar_t* w = weight_ptr + d * width;
|
||||
|
||||
// Accumulate in float32 for precision
|
||||
float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
|
||||
for (int64_t k = 0; k < state_len; ++k) {
|
||||
acc += static_cast<float>(w[k]) *
|
||||
static_cast<float>(sd[k * stride_s_state]);
|
||||
}
|
||||
acc += static_cast<float>(w[state_len]) * x_val;
|
||||
|
||||
// Shift state left and append new input.
|
||||
// Use memmove when contiguous (stride==1); element loop otherwise.
|
||||
if (stride_s_state == 1) {
|
||||
if (state_len > 1)
|
||||
std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
|
||||
if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
|
||||
} else {
|
||||
for (int64_t k = 0; k < state_len - 1; ++k)
|
||||
sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
|
||||
if (state_len > 0)
|
||||
sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
|
||||
}
|
||||
|
||||
if (do_silu) {
|
||||
float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
|
||||
: std::exp(acc) / (1.0f + std::exp(acc));
|
||||
acc *= sigmoid;
|
||||
}
|
||||
out_b[d * seqlen] = static_cast<scalar_t>(acc);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// selective_state_update
|
||||
//
|
||||
// Template parameters:
|
||||
// state_t - dtype of ssm_state cache (typically BFloat16)
|
||||
// input_t - dtype of x, B, C (typically BFloat16)
|
||||
// out_t - dtype of output tensor (typically BFloat16)
|
||||
// Write directly — no float32 intermediate buffer needed.
|
||||
//
|
||||
// A, D, dt_bias are accepted as const float* (they are always float32
|
||||
// model parameters in Mamba2). This eliminates the per-call float32→BF16
|
||||
// conversion and the .contiguous() materialisation of the broadcast-expand.
|
||||
//
|
||||
// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
|
||||
// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename state_t, typename input_t, typename out_t = float>
|
||||
inline void selective_state_update_kernel(
|
||||
state_t* __restrict__ state_ptr, int64_t stride_state_n,
|
||||
int64_t stride_state_h, int64_t stride_state_d,
|
||||
const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
|
||||
// dt: (N, nheads) — scalar per head, NOT expanded to head_dim
|
||||
const float* __restrict__ dt_ptr, int64_t stride_dt_n,
|
||||
// A: (nheads,) float32 — scalar per head
|
||||
const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
|
||||
const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
|
||||
// D: (nheads,) float32 — scalar per head (nullptr if not used)
|
||||
const float* __restrict__ D_ptr,
|
||||
// z: same shape as x (optional)
|
||||
const input_t* __restrict__ z_ptr,
|
||||
// dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
|
||||
const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
|
||||
int64_t stride_out_n, int64_t stride_out_h,
|
||||
const int32_t* __restrict__ state_batch_indices,
|
||||
const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
|
||||
const int32_t* __restrict__ num_accepted_tokens,
|
||||
const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
|
||||
int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
|
||||
using state_vec_t = vec_op::vec_t<state_t>;
|
||||
using input_vec_t = vec_op::vec_t<input_t>;
|
||||
constexpr int VEC_ELEM_NUM = 8;
|
||||
|
||||
int64_t nheads_per_group = nheads / ngroups;
|
||||
|
||||
for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
|
||||
int64_t bos, seq_len;
|
||||
if (cu_seqlens != nullptr) {
|
||||
bos = cu_seqlens[seq_idx];
|
||||
seq_len = cu_seqlens[seq_idx + 1] - bos;
|
||||
} else {
|
||||
bos = seq_idx;
|
||||
seq_len = 1;
|
||||
}
|
||||
|
||||
int64_t state_read_idx = (state_batch_indices != nullptr)
|
||||
? state_batch_indices[seq_idx]
|
||||
: seq_idx;
|
||||
if (state_read_idx == null_block_id) continue;
|
||||
|
||||
int64_t state_write_idx = (num_accepted_tokens == nullptr)
|
||||
? ((dst_state_batch_indices != nullptr)
|
||||
? dst_state_batch_indices[seq_idx]
|
||||
: state_read_idx)
|
||||
: -1;
|
||||
|
||||
state_t* s = state_ptr + state_read_idx * stride_state_n;
|
||||
|
||||
for (int64_t t = 0; t < seq_len; ++t) {
|
||||
int64_t token_idx = bos + t;
|
||||
const input_t* x_tok = x_ptr + token_idx * stride_x_n;
|
||||
// dt: (N, nheads) — one float per head per token
|
||||
const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
|
||||
const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
|
||||
const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
|
||||
out_t* out_tok = out_ptr + token_idx * stride_out_n;
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t h = 0; h < nheads; ++h) {
|
||||
int64_t g = h / nheads_per_group;
|
||||
const input_t* x_h = x_tok + h * stride_x_h;
|
||||
const input_t* B_g = B_tok + g * stride_BC_g;
|
||||
const input_t* C_g = C_tok + g * stride_BC_g;
|
||||
out_t* out_h = out_tok + h * stride_out_h;
|
||||
state_t* s_h = s + h * stride_state_h;
|
||||
|
||||
// Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
|
||||
float dt_val = dt_tok[h];
|
||||
if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
|
||||
if (dt_softplus) {
|
||||
dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
|
||||
}
|
||||
const float A_val = A_ptr[h]; // scalar: same for all dim, dstate
|
||||
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
|
||||
|
||||
const input_t* z_h =
|
||||
(z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
|
||||
: nullptr;
|
||||
|
||||
vec_op::FP32Vec8 dt_vec(dt_val);
|
||||
// dA = exp(A * dt): A and dt are SCALARS per head, so compute once
|
||||
// and broadcast. This saves 7 redundant std::exp() calls that
|
||||
// FP32Vec8::exp() would otherwise make on the broadcast vector.
|
||||
const float dA_scalar = std::exp(A_val * dt_val);
|
||||
vec_op::FP32Vec8 dA(dA_scalar); // broadcast
|
||||
|
||||
for (int64_t d = 0; d < dim; ++d) {
|
||||
float x_val = static_cast<float>(x_h[d]);
|
||||
|
||||
vec_op::FP32Vec8 out_vec(0.0f);
|
||||
state_t* s_hd = s_h + d * stride_state_d;
|
||||
const input_t* B_g_base = B_g;
|
||||
const input_t* C_g_base = C_g;
|
||||
|
||||
vec_op::FP32Vec8 x_vec(x_val);
|
||||
// dBx = B * x * dt — same dA for all dstate (A is scalar)
|
||||
// s_new = s * dA + B * x * dt
|
||||
|
||||
int64_t n = 0;
|
||||
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
|
||||
vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
|
||||
vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
|
||||
vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
|
||||
|
||||
vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
|
||||
vec_op::FP32Vec8 s_new = s_v * dA + dBx;
|
||||
|
||||
state_vec_t(s_new).save(s_hd + n);
|
||||
out_vec = out_vec + s_new * C_v;
|
||||
}
|
||||
|
||||
float out_val = out_vec.reduce_sum();
|
||||
for (; n < dstate; ++n) {
|
||||
// Reuse dA_scalar computed once per head — no exp() re-call
|
||||
float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
|
||||
float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
|
||||
s_hd[n] = static_cast<state_t>(s_new);
|
||||
out_val += s_new * static_cast<float>(C_g[n]);
|
||||
}
|
||||
|
||||
if (D_ptr != nullptr) out_val += x_val * D_val;
|
||||
if (z_h != nullptr) {
|
||||
float z_val = static_cast<float>(z_h[d]);
|
||||
float sigmoid = (z_val >= 0)
|
||||
? 1.0f / (1.0f + std::exp(-z_val))
|
||||
: std::exp(z_val) / (1.0f + std::exp(z_val));
|
||||
out_val *= z_val * sigmoid;
|
||||
}
|
||||
out_h[d] = static_cast<out_t>(out_val);
|
||||
}
|
||||
}
|
||||
|
||||
if (num_accepted_tokens != nullptr &&
|
||||
dst_state_batch_indices != nullptr) {
|
||||
int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
|
||||
if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
|
||||
state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
|
||||
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
|
||||
state_write_idx != state_read_idx) {
|
||||
state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
|
||||
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mamba_chunk_scan_fwd
|
||||
//
|
||||
// Prefill SSM recurrence for Mamba2 / SSD models.
|
||||
//
|
||||
// Key difference from selective_state_update_kernel (decode path):
|
||||
// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
|
||||
// Each thread owns a (batch, head) slice and runs the entire token
|
||||
// sequence without any per-token OpenMP synchronisation overhead.
|
||||
// For seqlen=256, this eliminates 256 thread-barrier launches per batch.
|
||||
//
|
||||
// `dt` arrives already processed (float32, after bias + softplus + clamp)
|
||||
// to keep this kernel simple. Preprocessing is done in the Python wrapper.
|
||||
//
|
||||
// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
|
||||
// tensor, pre-initialised by the caller (zero or from initial_states).
|
||||
// Each (b, h) slice is private to exactly one thread via collapse(2), so
|
||||
// there are no write conflicts.
|
||||
//
|
||||
// D is treated as a scalar per head ([nheads] float32).
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename input_t>
|
||||
inline void mamba_chunk_scan_fwd_kernel(
|
||||
float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32
|
||||
const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim]
|
||||
const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed)
|
||||
const float* __restrict__ A_ptr, // [nheads] f32
|
||||
const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate]
|
||||
const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate]
|
||||
const float* __restrict__ D_ptr, // [nheads] f32 (nullable)
|
||||
const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable)
|
||||
input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim]
|
||||
const int32_t* __restrict__ cu_seqlens, // [batch+1] int32
|
||||
int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
|
||||
int64_t dstate) {
|
||||
using input_vec_t = vec_op::vec_t<input_t>;
|
||||
constexpr int VEC_ELEM_NUM = 8;
|
||||
|
||||
const int64_t nheads_per_group = nheads / ngroups;
|
||||
// states layout: [batch, nheads, headdim, dstate] contiguous (caller
|
||||
// guarantee)
|
||||
const int64_t stride_s_b = nheads * headdim * dstate;
|
||||
const int64_t stride_s_h = headdim * dstate;
|
||||
// stride_s_d = dstate, stride_s_n = 1
|
||||
|
||||
#pragma omp parallel for collapse(2) schedule(static)
|
||||
for (int64_t b = 0; b < batch; ++b) {
|
||||
for (int64_t h = 0; h < nheads; ++h) {
|
||||
const int64_t seq_start = cu_seqlens[b];
|
||||
const int64_t seq_end = cu_seqlens[b + 1];
|
||||
const int64_t g = h / nheads_per_group;
|
||||
|
||||
const float A_val = A_ptr[h];
|
||||
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
|
||||
|
||||
// Working state slice: states[b, h, :, :] — float32, headdim * dstate.
|
||||
// Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
|
||||
float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
|
||||
|
||||
for (int64_t t = seq_start; t < seq_end; ++t) {
|
||||
const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
|
||||
const float* dt_h = dt_ptr + t * nheads + h;
|
||||
const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
|
||||
const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
|
||||
const input_t* z_h = (z_ptr != nullptr)
|
||||
? z_ptr + t * nheads * headdim + h * headdim
|
||||
: nullptr;
|
||||
input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
|
||||
|
||||
const float dt_val = *dt_h;
|
||||
const float dA_val = std::exp(A_val * dt_val);
|
||||
const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar
|
||||
const vec_op::FP32Vec8 dt_vec(dt_val);
|
||||
|
||||
for (int64_t d = 0; d < headdim; ++d) {
|
||||
const float x_val = static_cast<float>(x_h[d]);
|
||||
float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32
|
||||
|
||||
// Vectorised SSM update + readout over dstate:
|
||||
// s_new = s * dA + x * dt * B
|
||||
// y += s_new * C
|
||||
int64_t n = 0;
|
||||
vec_op::FP32Vec8 y_vec(0.0f);
|
||||
const vec_op::FP32Vec8 x_vec(x_val);
|
||||
|
||||
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
|
||||
const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
|
||||
const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
|
||||
const vec_op::FP32Vec8 s_v(s_bhd + n);
|
||||
|
||||
const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
|
||||
s_new.save(s_bhd + n);
|
||||
y_vec = y_vec + s_new * C_v;
|
||||
}
|
||||
|
||||
float y_val = y_vec.reduce_sum();
|
||||
|
||||
// Scalar tail for remaining dstate elements
|
||||
for (; n < dstate; ++n) {
|
||||
const float B_n = static_cast<float>(B_g[n]);
|
||||
const float C_n = static_cast<float>(C_g[n]);
|
||||
const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
|
||||
s_bhd[n] = s_new;
|
||||
y_val += s_new * C_n;
|
||||
}
|
||||
|
||||
// D skip connection (scalar per head)
|
||||
if (D_ptr != nullptr) y_val += x_val * D_val;
|
||||
|
||||
// z gating: out = y * z * sigmoid(z) (SiLU)
|
||||
if (z_h != nullptr) {
|
||||
const float z_val = static_cast<float>(z_h[d]);
|
||||
const float sigmoid =
|
||||
(z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
|
||||
: std::exp(z_val) / (1.0f + std::exp(z_val));
|
||||
y_val *= z_val * sigmoid;
|
||||
}
|
||||
|
||||
out_h[d] = static_cast<input_t>(y_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mamba_cpu
|
||||
@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
|
||||
torch::Tensor slot_mapping,
|
||||
const int64_t block_size);
|
||||
|
||||
at::Tensor causal_conv1d_update_cpu_impl(
|
||||
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
|
||||
const c10::optional<at::Tensor>& bias,
|
||||
const c10::optional<std::string>& activation,
|
||||
const c10::optional<at::Tensor>& conv_state_indices,
|
||||
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
|
||||
|
||||
void selective_state_update_cpu_impl(
|
||||
at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
|
||||
const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
|
||||
const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
|
||||
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
|
||||
const c10::optional<at::Tensor>& state_batch_indices,
|
||||
const c10::optional<at::Tensor>& dst_state_batch_indices,
|
||||
int64_t null_block_id, at::Tensor& out,
|
||||
const c10::optional<at::Tensor>& num_accepted_tokens,
|
||||
const c10::optional<at::Tensor>& cu_seqlens);
|
||||
|
||||
void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
|
||||
const at::Tensor& x, const at::Tensor& dt,
|
||||
const at::Tensor& A, const at::Tensor& B,
|
||||
const at::Tensor& C,
|
||||
const c10::optional<at::Tensor>& D,
|
||||
const c10::optional<at::Tensor>& z,
|
||||
const at::Tensor& cu_seqlens);
|
||||
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids);
|
||||
|
||||
namespace cpu_utils {
|
||||
@@ -595,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"block_size) -> ()",
|
||||
&compute_slot_mapping_kernel_impl);
|
||||
|
||||
// Mamba CPU kernels
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu_vec("
|
||||
"Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
|
||||
"Tensor? bias, str? activation, Tensor? conv_state_indices, "
|
||||
"Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
|
||||
&causal_conv1d_update_cpu_impl);
|
||||
|
||||
ops.def(
|
||||
"selective_state_update_cpu("
|
||||
"Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
|
||||
"Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
|
||||
"Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
|
||||
"SymInt null_block_id, Tensor(a13!) out, "
|
||||
"Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
|
||||
&selective_state_update_cpu_impl);
|
||||
|
||||
ops.def(
|
||||
"mamba_chunk_scan_fwd_cpu("
|
||||
"Tensor(a0!) out, Tensor(a1!) final_states, "
|
||||
"Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
|
||||
"Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
|
||||
&mamba_chunk_scan_fwd_cpu_impl);
|
||||
|
||||
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
|
||||
|
||||
// Speculative decoding kernels
|
||||
|
||||
@@ -47,6 +47,8 @@ torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
|
||||
torch::stable::Tensor const& perm);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
std::string get_compiled_cuda_archs();
|
||||
|
||||
bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
|
||||
bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
|
||||
bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
|
||||
|
||||
@@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
|
||||
// AWQ zeros: (size_k // group_size, size_n // 8)
|
||||
const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
|
||||
int32_t group_size) {
|
||||
int32_t val =
|
||||
qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
|
||||
int32_t zero =
|
||||
qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
|
||||
blockIdx.y];
|
||||
// Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
|
||||
// blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
|
||||
// same row (stride 1) instead of striding across rows (stride size_n/8).
|
||||
int col = blockIdx.y * 32 + threadIdx.x;
|
||||
if (col >= size_n / 8) return;
|
||||
(void)size_k;
|
||||
|
||||
int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
|
||||
int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
|
||||
int32_t new_val = 0;
|
||||
|
||||
#pragma unroll
|
||||
@@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
|
||||
zero >>= 4;
|
||||
}
|
||||
|
||||
output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
|
||||
output[blockIdx.x * (size_n / 8) + col] = new_val;
|
||||
}
|
||||
|
||||
torch::stable::Tensor marlin_int4_fp8_preprocess(
|
||||
@@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
|
||||
"qweight.size(0) % qzeros.size(0) != 0");
|
||||
STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
|
||||
|
||||
dim3 blocks(size_k / 32, size_n / 8);
|
||||
dim3 blocks(size_k, (size_n / 8 + 31) / 32);
|
||||
marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
|
||||
reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
|
||||
|
||||
@@ -51,7 +51,8 @@ void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,
|
||||
|
||||
#endif
|
||||
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
|
||||
void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
|
||||
torch::stable::Tensor const& a_tensors,
|
||||
torch::stable::Tensor const& b_tensors,
|
||||
@@ -83,8 +84,9 @@ void cutlass_scaled_mm_sm100(torch::stable::Tensor& c,
|
||||
std::optional<torch::stable::Tensor> const& bias);
|
||||
#endif
|
||||
|
||||
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined(ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
(defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
|
||||
void get_cutlass_moe_mm_data_caller(
|
||||
const torch::stable::Tensor& topk_ids,
|
||||
@@ -175,11 +177,14 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
|
||||
|
||||
bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
|
||||
// CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
|
||||
// or CUDA 12.8 and SM100 (Blackwell). Only report archs that have an
|
||||
// actual cutlass_moe_mm dispatch compiled into this file.
|
||||
// or CUDA 12.8 and SM10x/SM11x (Blackwell / Thor). CUDA 12 reports Thor as
|
||||
// SM101 while CUDA 13 reports it as SM110, but both use this sm100-named
|
||||
// implementation. Only report archs that have an actual cutlass_moe_mm
|
||||
// dispatch compiled into this file.
|
||||
|
||||
#if defined CUDA_VERSION
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
|
||||
if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
|
||||
return CUDA_VERSION >= 12080;
|
||||
}
|
||||
@@ -281,8 +286,11 @@ void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
|
||||
torch::stable::Tensor const& c_strides, bool per_act_token,
|
||||
bool per_out_ch) {
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
|
||||
if (version_num >= 100 && version_num < 110) {
|
||||
#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
|
||||
// Keep runtime dispatch aligned with the CMake arch list and support query:
|
||||
// CUDA 12 Thor is SM101 and CUDA 13 Thor is SM110.
|
||||
if (version_num >= 100 && version_num < 120) {
|
||||
cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
|
||||
expert_offsets, problem_sizes, a_strides, b_strides,
|
||||
c_strides, per_act_token, per_out_ch);
|
||||
@@ -316,8 +324,9 @@ void get_cutlass_moe_mm_data(
|
||||
// This function currently gets compiled only if we have a valid cutlass moe
|
||||
// mm to run it for.
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
|
||||
get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
|
||||
problem_sizes2, input_permutation,
|
||||
@@ -338,8 +347,9 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
|
||||
torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
|
||||
const bool swap_ab) {
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
|
||||
get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
|
||||
expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
|
||||
@@ -362,8 +372,9 @@ void get_cutlass_batched_moe_mm_data(
|
||||
// This function currently gets compiled only if we have a valid cutlass moe
|
||||
// mm to run it for.
|
||||
int32_t version_num = get_sm_version_num();
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
|
||||
#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
|
||||
ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) || \
|
||||
(defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
|
||||
get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
|
||||
problem_sizes2, expert_num_tokens,
|
||||
|
||||
@@ -4,6 +4,10 @@
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
|
||||
#ifndef USE_ROCM
|
||||
std::string get_compiled_cuda_archs() { return VLLM_COMPILED_CUDA_ARCHS; }
|
||||
#endif
|
||||
|
||||
// Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
|
||||
// Note: We register under namespace "_C" so ops are accessible as
|
||||
// torch.ops._C.<op_name> for compatibility with existing code.
|
||||
@@ -30,6 +34,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");
|
||||
|
||||
ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
|
||||
ops.def("get_compiled_cuda_archs() -> str");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
|
||||
@@ -763,6 +768,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
|
||||
// ops.impl("op_name", &func) without a dispatch key in the non-stable API.
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
|
||||
#ifndef USE_ROCM
|
||||
ops.impl("get_compiled_cuda_archs", TORCH_BOX(&get_compiled_cuda_archs));
|
||||
ops.impl("cutlass_scaled_mm_supports_fp8",
|
||||
TORCH_BOX(&cutlass_scaled_mm_supports_fp8));
|
||||
ops.impl("cutlass_group_gemm_supported",
|
||||
|
||||
@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
|
||||
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
|
||||
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
|
||||
|
||||
@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
|
||||
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
|
||||
- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu)
|
||||
- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
|
||||
|
||||
### SiLU+Mul + Quantization (`fuse_act_quant`)
|
||||
|
||||
@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
|
||||
- Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
|
||||
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
|
||||
- CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
|
||||
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
|
||||
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
|
||||
|
||||
### RMSNorm + Padding (`fuse_act_padding`)
|
||||
|
||||
|
||||
@@ -68,13 +68,14 @@ vllm serve <model> \
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
|
||||
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
|
||||
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
|
||||
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
|
||||
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
|
||||
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
|
||||
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
|
||||
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
|
||||
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
|
||||
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
|
||||
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
|
||||
|
||||
## Secondary Tiers
|
||||
@@ -179,7 +180,7 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
|
||||
|
||||
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
|
||||
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
|
||||
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
|
||||
- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
|
||||
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
|
||||
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
|
||||
|
||||
|
||||
@@ -46,6 +46,26 @@ export CPU_ARCH=$(uname -m) # x86_64 or aarch64
|
||||
uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
|
||||
```
|
||||
|
||||
!!! warning "CUDA architecture coverage"
|
||||
|
||||
Pre-built CUDA wheels are compiled for the CUDA architectures selected by
|
||||
vLLM's release and build pipelines. This list is intentionally smaller than
|
||||
every architecture CMake can build, because each additional architecture
|
||||
increases wheel size.
|
||||
|
||||
In particular, CUDA 12.9 wheels do not use CUDA 13 family-specific targets.
|
||||
To keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
|
||||
architecture-specific Blackwell/Thor targets, such as `sm_103` or
|
||||
`sm_121`, even though vLLM can build them from source. CUDA 13 wheels use
|
||||
family-specific targets such as `sm_100f`, `sm_110f`, and `sm_120f`, which
|
||||
cover the corresponding major-version GPU family.
|
||||
|
||||
If vLLM logs a warning that your visible CUDA device is not covered by the
|
||||
wheel's compiled CUDA architectures, or if you see a CUDA error such as
|
||||
`no kernel image is available for execution on the device`, install a CUDA
|
||||
13 wheel when possible, or build from source with a `TORCH_CUDA_ARCH_LIST`
|
||||
that includes your GPU.
|
||||
|
||||
#### Install the latest code
|
||||
|
||||
LLM inference is a fast-evolving field, and the latest code may contain bug fixes, performance improvements, and new features that are not released yet. To allow users to try the latest code without waiting for the next release, vLLM provides wheels for every commit since `v0.5.3` on <https://wheels.vllm.ai/nightly>. There are multiple indices that could be used:
|
||||
|
||||
@@ -31,10 +31,8 @@
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | |
|
||||
| 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | |
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | |
|
||||
| meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | |
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | |
|
||||
| Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | |
|
||||
| RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | |
|
||||
|
||||
@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
|
||||
- `response_format`: Format of the response ("json", "text") (optional)
|
||||
- `temperature`: Sampling temperature between 0 and 1 (optional)
|
||||
|
||||
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182).
|
||||
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).
|
||||
|
||||
**Response Format:**
|
||||
|
||||
|
||||
@@ -174,6 +174,33 @@ If the script runs successfully, you should see the message `sanity check is suc
|
||||
|
||||
If the test script hangs or crashes, usually it means the hardware/drivers are broken in some sense. You should try to contact your system administrator or hardware vendor for further assistance. As a common workaround, you can try to tune some NCCL environment variables, such as `export NCCL_P2P_DISABLE=1` to see if it helps. Please check [their documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) for more information. Please only use these environment variables as a temporary workaround, as they might affect the performance of the system. The best solution is still to fix the hardware/drivers so that the test script can run successfully.
|
||||
|
||||
## CUDA architecture not covered by the wheel
|
||||
|
||||
If vLLM logs a warning that the current wheel was built for a set of CUDA
|
||||
architectures but one of your visible CUDA devices is not covered, the installed
|
||||
wheel may not contain native CUDA kernels for that GPU. The same issue may also
|
||||
surface later as a CUDA runtime error such as `no kernel image is available for
|
||||
execution on the device`.
|
||||
|
||||
The architecture list in a pre-built wheel is determined by the vLLM release and
|
||||
build pipelines, then filtered by CMake before individual kernels choose their
|
||||
own per-kernel architectures. It is not the same as the full set of
|
||||
architectures that vLLM can build from source.
|
||||
|
||||
CUDA 12.9 wheels use architecture-specific targets for newer NVIDIA GPUs. To
|
||||
keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
|
||||
Blackwell/Thor targets such as `sm_103` or `sm_121`, even though vLLM can build
|
||||
them from source. CUDA 13 wheels can use family-specific targets such as
|
||||
`sm_100f`, `sm_110f`, and `sm_120f`, which cover the corresponding
|
||||
major-version GPU family.
|
||||
|
||||
If you see this warning or error, install a CUDA 13 wheel when possible.
|
||||
Otherwise, build vLLM from source and set `TORCH_CUDA_ARCH_LIST` to include
|
||||
your GPU's compute capability. See the CUDA installation guide's
|
||||
[pre-built wheels](../getting_started/installation/gpu.md#pre-built-wheels) and
|
||||
[build wheel from source](../getting_started/installation/gpu.md#build-wheel-from-source)
|
||||
sections for details.
|
||||
|
||||
## Python multiprocessing
|
||||
|
||||
### `RuntimeError` Exception
|
||||
|
||||
Generated
+3
-2
@@ -5560,6 +5560,7 @@ dependencies = [
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"uuid",
|
||||
"vllm-bench",
|
||||
"vllm-chat",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-managed-engine",
|
||||
@@ -6320,9 +6321,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
|
||||
|
||||
[[package]]
|
||||
name = "xgrammar-structural-tag"
|
||||
version = "0.1.0+xgrammar.0.2.2.4d145cc"
|
||||
version = "0.2.0+xgrammar.0.2.4.dd729e7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6"
|
||||
checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
|
||||
dependencies = [
|
||||
"auto_impl",
|
||||
"serde",
|
||||
|
||||
+2
-1
@@ -132,6 +132,7 @@ trait-set = "0.3.0"
|
||||
url = "2.5.7"
|
||||
uuid = { version = "1.22.0", features = ["v4"] }
|
||||
validator = { version = "0.20.0", features = ["derive"] }
|
||||
vllm-bench = { path = "src/bench" }
|
||||
vllm-chat = { path = "src/chat" }
|
||||
vllm-engine-core-client = { path = "src/engine-core-client" }
|
||||
vllm-llm = { path = "src/llm" }
|
||||
@@ -142,7 +143,7 @@ vllm-server = { path = "src/server" }
|
||||
vllm-text = { path = "src/text" }
|
||||
vllm-tokenizer = { path = "src/tokenizer" }
|
||||
winnow = { version = "1.0.2", features = ["simd"] }
|
||||
xgrammar-structural-tag = "0.1.0"
|
||||
xgrammar-structural-tag = "0.2.0"
|
||||
zeromq = { version = "0.6.0", default-features = false, features = [
|
||||
"tokio-runtime",
|
||||
"all-transport",
|
||||
|
||||
@@ -3,8 +3,6 @@
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use clap::Parser;
|
||||
|
||||
/// Backend type for the benchmark endpoint.
|
||||
#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum BackendKind {
|
||||
@@ -77,7 +75,7 @@ pub enum DatasetName {
|
||||
ShareGpt,
|
||||
#[value(name = "sonnet")]
|
||||
Sonnet,
|
||||
#[value(name = "speed-bench")]
|
||||
#[value(name = "speed-bench", alias = "speed_bench")]
|
||||
SpeedBench,
|
||||
#[value(name = "hf")]
|
||||
Hf,
|
||||
@@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig {
|
||||
}
|
||||
|
||||
/// High-performance benchmark client for vLLM serving endpoints.
|
||||
#[derive(Parser, Debug, Clone)]
|
||||
#[command(
|
||||
name = "vllm-bench",
|
||||
about = "Benchmark online serving throughput",
|
||||
version
|
||||
)]
|
||||
pub struct Cli {
|
||||
#[derive(clap::Args, Debug, Clone)]
|
||||
pub struct BenchServeArgs {
|
||||
/// The type of backend or endpoint to use for the benchmark.
|
||||
#[arg(long, default_value = "openai")]
|
||||
pub backend: BackendKind,
|
||||
@@ -659,7 +652,7 @@ pub struct Cli {
|
||||
pub lora_assignment: LoraAssignment,
|
||||
}
|
||||
|
||||
impl Cli {
|
||||
impl BenchServeArgs {
|
||||
/// Resolve the base URL from explicit --base-url or from --host/--port.
|
||||
pub fn resolve_base_url(&self) -> String {
|
||||
if let Some(ref base) = self.base_url {
|
||||
|
||||
+212
-188
@@ -4,7 +4,9 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
|
||||
use crate::cli::{
|
||||
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
|
||||
};
|
||||
use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
|
||||
use crate::error::{BenchError, Result};
|
||||
|
||||
@@ -215,63 +217,63 @@ pub struct BenchConfig {
|
||||
}
|
||||
|
||||
impl BenchConfig {
|
||||
pub fn from_cli(cli: &Cli) -> Result<Self> {
|
||||
if cli.burstiness <= 0.0 {
|
||||
pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
|
||||
if args.burstiness <= 0.0 {
|
||||
return Err(BenchError::Config("Burstiness must be positive".into()));
|
||||
}
|
||||
|
||||
if cli.num_prompts == 0 {
|
||||
if args.num_prompts == 0 {
|
||||
return Err(BenchError::Config(
|
||||
"--num-prompts must be at least 1".into(),
|
||||
));
|
||||
}
|
||||
|
||||
if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
|
||||
if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
|
||||
return Err(BenchError::Config(
|
||||
"--request-rate must be positive (or inf)".into(),
|
||||
));
|
||||
}
|
||||
if cli.max_model_len == Some(0) {
|
||||
if args.max_model_len == Some(0) {
|
||||
return Err(BenchError::Config(
|
||||
"--max-model-len must be at least 1".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let base_url = cli.resolve_base_url();
|
||||
let api_url = cli.resolve_api_url();
|
||||
let base_url = args.resolve_base_url();
|
||||
let api_url = args.resolve_api_url();
|
||||
|
||||
let extra_headers = cli.parse_headers()?;
|
||||
let mut extra_body = cli.parse_extra_body()?;
|
||||
let extra_headers = args.parse_headers()?;
|
||||
let mut extra_body = args.parse_extra_body()?;
|
||||
|
||||
// Merge sampling parameters into extra_body (matches Python behavior).
|
||||
// Python collects non-None sampling params and merges them UNDER extra_body,
|
||||
// meaning extra_body keys take precedence over sampling params.
|
||||
{
|
||||
let mut sampling_params = serde_json::Map::new();
|
||||
if let Some(v) = cli.top_p {
|
||||
if let Some(v) = args.top_p {
|
||||
sampling_params.insert("top_p".into(), serde_json::json!(v));
|
||||
}
|
||||
if let Some(v) = cli.top_k {
|
||||
if let Some(v) = args.top_k {
|
||||
sampling_params.insert("top_k".into(), serde_json::json!(v));
|
||||
}
|
||||
if let Some(v) = cli.min_p {
|
||||
if let Some(v) = args.min_p {
|
||||
sampling_params.insert("min_p".into(), serde_json::json!(v));
|
||||
}
|
||||
if let Some(v) = cli.temperature {
|
||||
if let Some(v) = args.temperature {
|
||||
sampling_params.insert("temperature".into(), serde_json::json!(v));
|
||||
}
|
||||
if let Some(v) = cli.frequency_penalty {
|
||||
if let Some(v) = args.frequency_penalty {
|
||||
sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
|
||||
}
|
||||
if let Some(v) = cli.presence_penalty {
|
||||
if let Some(v) = args.presence_penalty {
|
||||
sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
|
||||
}
|
||||
if let Some(v) = cli.repetition_penalty {
|
||||
if let Some(v) = args.repetition_penalty {
|
||||
sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
|
||||
}
|
||||
|
||||
if !sampling_params.is_empty() {
|
||||
if !cli.backend.is_openai_compatible() {
|
||||
if !args.backend.is_openai_compatible() {
|
||||
return Err(BenchError::Config(
|
||||
"Sampling parameters are only supported by openai-compatible backends."
|
||||
.into(),
|
||||
@@ -299,7 +301,7 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
// Parse metadata
|
||||
let metadata = match &cli.metadata {
|
||||
let metadata = match &args.metadata {
|
||||
None => None,
|
||||
Some(items) => {
|
||||
let mut pairs = Vec::new();
|
||||
@@ -314,24 +316,24 @@ impl BenchConfig {
|
||||
};
|
||||
|
||||
// Parse goodput SLOs
|
||||
let goodput = parse_goodput(&cli.goodput)?;
|
||||
let goodput = parse_goodput(&args.goodput)?;
|
||||
|
||||
// Parse ramp-up config
|
||||
let ramp_up = parse_ramp_up(cli)?;
|
||||
let ramp_up = parse_ramp_up(args)?;
|
||||
|
||||
// Default percentile metrics based on backend type
|
||||
let default_percentile_metrics = if cli.backend.is_pooling() {
|
||||
let default_percentile_metrics = if args.backend.is_pooling() {
|
||||
"e2el"
|
||||
} else {
|
||||
"ttft,tpot,itl,e2el"
|
||||
};
|
||||
let percentile_metrics_str =
|
||||
cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
|
||||
args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
|
||||
let selected_percentile_metrics: Vec<String> =
|
||||
percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
|
||||
|
||||
let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
|
||||
let sweep_summary_percentiles = cli
|
||||
let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
|
||||
let sweep_summary_percentiles = args
|
||||
.sweep_summary_percentiles
|
||||
.as_deref()
|
||||
.map(|raw| parse_percentiles(raw, true))
|
||||
@@ -344,38 +346,38 @@ impl BenchConfig {
|
||||
selected_percentiles.push(90.0);
|
||||
}
|
||||
|
||||
let tokenizer_id = if cli.skip_tokenizer_init {
|
||||
let tokenizer_id = if args.skip_tokenizer_init {
|
||||
None
|
||||
} else {
|
||||
Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
|
||||
args.tokenizer.clone().or_else(|| args.model.clone())
|
||||
};
|
||||
|
||||
// Resolve input/output lengths
|
||||
let random_input_len = cli.resolved_random_input_len();
|
||||
let random_output_len = cli.resolved_random_output_len();
|
||||
let per_turn_input_len = cli.resolved_per_turn_input_len();
|
||||
let random_input_len = args.resolved_random_input_len();
|
||||
let random_output_len = args.resolved_random_output_len();
|
||||
let per_turn_input_len = args.resolved_per_turn_input_len();
|
||||
|
||||
// Normalized multi-turn turn counts (computed in validation block below, defaults
|
||||
// to num_turns if multi-turn mode is not active)
|
||||
let mut multi_turn_min_turns = cli.multi_turn_num_turns;
|
||||
let mut multi_turn_max_turns = cli.multi_turn_num_turns;
|
||||
let mut multi_turn_min_turns = args.multi_turn_num_turns;
|
||||
let mut multi_turn_max_turns = args.multi_turn_num_turns;
|
||||
|
||||
// For random datasets with openai-compatible backends, default to ignore_eos.
|
||||
// Exception: multi-turn mode, where ignore_eos causes unbounded context growth
|
||||
// across turns. Multi-turn uses min_tokens instead for output length control.
|
||||
// Pooling backends don't generate tokens, so ignore_eos is irrelevant.
|
||||
let ignore_eos = if cli.backend.is_pooling() {
|
||||
let ignore_eos = if args.backend.is_pooling() {
|
||||
false
|
||||
} else {
|
||||
cli.ignore_eos
|
||||
|| ((cli.dataset_name == DatasetName::Random
|
||||
|| cli.dataset_name == DatasetName::RandomMm)
|
||||
&& cli.backend.is_openai_compatible()
|
||||
&& !cli.multi_turn)
|
||||
args.ignore_eos
|
||||
|| ((args.dataset_name == DatasetName::Random
|
||||
|| args.dataset_name == DatasetName::RandomMm)
|
||||
&& args.backend.is_openai_compatible()
|
||||
&& !args.multi_turn)
|
||||
};
|
||||
|
||||
// Pooling backends don't support multi-turn
|
||||
if cli.backend.is_pooling() && cli.multi_turn {
|
||||
if args.backend.is_pooling() && args.multi_turn {
|
||||
return Err(BenchError::Config(
|
||||
"Pooling/embedding backends do not support --multi-turn".into(),
|
||||
));
|
||||
@@ -383,7 +385,7 @@ impl BenchConfig {
|
||||
|
||||
// LoRA validation. Adapter names must be non-empty after trim; pooling
|
||||
// backends are out of scope (vLLM LoRA routing is for generative paths).
|
||||
let lora_modules = match cli.lora_modules.as_ref() {
|
||||
let lora_modules = match args.lora_modules.as_ref() {
|
||||
None => None,
|
||||
Some(names) => {
|
||||
if names.is_empty() {
|
||||
@@ -391,7 +393,7 @@ impl BenchConfig {
|
||||
"--lora-modules requires at least one adapter name".into(),
|
||||
));
|
||||
}
|
||||
if cli.backend.is_pooling() {
|
||||
if args.backend.is_pooling() {
|
||||
return Err(BenchError::Config(
|
||||
"--lora-modules is not supported for pooling/embedding backends".into(),
|
||||
));
|
||||
@@ -411,18 +413,18 @@ impl BenchConfig {
|
||||
};
|
||||
|
||||
// Random-MM validation and config parsing
|
||||
let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
|
||||
if cli.backend != BackendKind::OpenaiChat {
|
||||
let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
|
||||
if args.backend != BackendKind::OpenaiChat {
|
||||
return Err(BenchError::Config(
|
||||
"Multi-modal content (images) is only supported on 'openai-chat' backend."
|
||||
.into(),
|
||||
));
|
||||
}
|
||||
let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
|
||||
&cli.random_mm_limit_mm_per_prompt,
|
||||
&args.random_mm_limit_mm_per_prompt,
|
||||
)?;
|
||||
let buckets =
|
||||
crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
|
||||
crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
|
||||
(limit, buckets)
|
||||
} else {
|
||||
(MmLimitPerPrompt::default(), Vec::new())
|
||||
@@ -432,18 +434,18 @@ impl BenchConfig {
|
||||
// sonnet (uses built-in Shakespeare's sonnets).
|
||||
|
||||
// Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
|
||||
let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
|
||||
let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;
|
||||
|
||||
// Batched inputs only make sense for pooling backends (the generation
|
||||
// backends send one prompt per request).
|
||||
if cli.random_batch_size == 0 {
|
||||
if args.random_batch_size == 0 {
|
||||
return Err(BenchError::Config(
|
||||
"--random-batch-size must be at least 1".into(),
|
||||
));
|
||||
}
|
||||
if cli.random_batch_size > 1
|
||||
&& !cli.backend.is_pooling()
|
||||
&& cli.dataset_name != DatasetName::RandomRerank
|
||||
if args.random_batch_size > 1
|
||||
&& !args.backend.is_pooling()
|
||||
&& args.dataset_name != DatasetName::RandomRerank
|
||||
{
|
||||
return Err(BenchError::Config(
|
||||
"--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
|
||||
@@ -451,16 +453,16 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
// random-rerank validation (mirrors Python RandomDatasetForReranking)
|
||||
let is_reranker = !cli.no_reranker;
|
||||
if cli.dataset_name == DatasetName::RandomRerank {
|
||||
if !cli.backend.is_pooling() {
|
||||
let is_reranker = !args.no_reranker;
|
||||
if args.dataset_name == DatasetName::RandomRerank {
|
||||
if !args.backend.is_pooling() {
|
||||
return Err(BenchError::Config(
|
||||
"--dataset-name random-rerank requires an embeddings/pooling backend \
|
||||
(e.g. --backend vllm-rerank)"
|
||||
.into(),
|
||||
));
|
||||
}
|
||||
if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
|
||||
if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
|
||||
return Err(BenchError::Config(
|
||||
"--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
|
||||
(the query is folded into the first batch slot)"
|
||||
@@ -470,8 +472,8 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
// Custom dataset validation
|
||||
if cli.dataset_name == DatasetName::Custom {
|
||||
match cli.dataset_path.as_deref() {
|
||||
if args.dataset_name == DatasetName::Custom {
|
||||
match args.dataset_path.as_deref() {
|
||||
None => {
|
||||
return Err(BenchError::Config(
|
||||
"--dataset-path is required for --dataset-name custom \
|
||||
@@ -486,7 +488,7 @@ impl BenchConfig {
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
if !cli.skip_chat_template {
|
||||
if !args.skip_chat_template {
|
||||
eprintln!(
|
||||
"NOTE: client-side chat template rendering is not supported; custom \
|
||||
dataset prompts are sent raw (equivalent to --skip-chat-template)."
|
||||
@@ -495,29 +497,29 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
// Prefix repetition validation
|
||||
if cli.dataset_name == DatasetName::PrefixRepetition {
|
||||
if cli.prefix_repetition_num_prefixes == 0 {
|
||||
if args.dataset_name == DatasetName::PrefixRepetition {
|
||||
if args.prefix_repetition_num_prefixes == 0 {
|
||||
return Err(BenchError::Config(
|
||||
"--prefix-repetition-num-prefixes must be at least 1".into(),
|
||||
));
|
||||
}
|
||||
if cli.num_prompts < cli.prefix_repetition_num_prefixes {
|
||||
if args.num_prompts < args.prefix_repetition_num_prefixes {
|
||||
return Err(BenchError::Config(format!(
|
||||
"--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
|
||||
cli.num_prompts, cli.prefix_repetition_num_prefixes
|
||||
args.num_prompts, args.prefix_repetition_num_prefixes
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
// HF dataset validation
|
||||
if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
|
||||
if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
|
||||
return Err(BenchError::Config(
|
||||
"--dataset-path is required for --dataset-name hf \
|
||||
(set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
|
||||
.into(),
|
||||
));
|
||||
}
|
||||
if let Some(len) = cli.hf_output_len
|
||||
if let Some(len) = args.hf_output_len
|
||||
&& len == 0
|
||||
{
|
||||
return Err(BenchError::Config(
|
||||
@@ -526,13 +528,13 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
// Multi-turn validation
|
||||
if cli.multi_turn {
|
||||
if cli.backend != BackendKind::OpenaiChat {
|
||||
if args.multi_turn {
|
||||
if args.backend != BackendKind::OpenaiChat {
|
||||
return Err(BenchError::Config(
|
||||
"--multi-turn requires --backend openai-chat".into(),
|
||||
));
|
||||
}
|
||||
if cli.multi_turn_num_turns == 0 {
|
||||
if args.multi_turn_num_turns == 0 {
|
||||
return Err(BenchError::Config(
|
||||
"--multi-turn-num-turns must be at least 1".into(),
|
||||
));
|
||||
@@ -541,18 +543,18 @@ impl BenchConfig {
|
||||
// Normalize and validate min/max turns. ShareGPT only consumes max_turns
|
||||
// (the loader walks all available turns up to the cap), so the
|
||||
// min/num/max coupling used for synthetic generation does not apply.
|
||||
if cli.dataset_name == DatasetName::ShareGpt {
|
||||
if cli.multi_turn_max_turns == 1 {
|
||||
if args.dataset_name == DatasetName::ShareGpt {
|
||||
if args.multi_turn_max_turns == 1 {
|
||||
return Err(BenchError::Config(
|
||||
"--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
|
||||
));
|
||||
}
|
||||
} else {
|
||||
(multi_turn_min_turns, multi_turn_max_turns) =
|
||||
match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
|
||||
(0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
|
||||
(m, 0) => (m, cli.multi_turn_num_turns),
|
||||
(0, x) => (cli.multi_turn_num_turns, x),
|
||||
match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
|
||||
(0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
|
||||
(m, 0) => (m, args.multi_turn_num_turns),
|
||||
(0, x) => (args.multi_turn_num_turns, x),
|
||||
(m, x) => (m, x),
|
||||
};
|
||||
if multi_turn_min_turns < 1 {
|
||||
@@ -575,8 +577,8 @@ impl BenchConfig {
|
||||
}
|
||||
|
||||
// Validate prefix sharing ratios
|
||||
let pg = cli.multi_turn_prefix_global_ratio;
|
||||
let pc = cli.multi_turn_prefix_conversation_ratio;
|
||||
let pg = args.multi_turn_prefix_global_ratio;
|
||||
let pc = args.multi_turn_prefix_conversation_ratio;
|
||||
if !(0.0..=1.0).contains(&pg) {
|
||||
return Err(BenchError::Config(
|
||||
"--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
|
||||
@@ -592,20 +594,20 @@ impl BenchConfig {
|
||||
"--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
|
||||
));
|
||||
}
|
||||
if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
|
||||
if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
|
||||
return Err(BenchError::Config(
|
||||
"Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
|
||||
if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
|
||||
return Err(BenchError::Config(format!(
|
||||
"--steady-state-threshold must be in (0.0, 1.0], got {}",
|
||||
cli.steady_state_threshold
|
||||
args.steady_state_threshold
|
||||
)));
|
||||
}
|
||||
if let Some(mw) = cli.steady_state_min_window
|
||||
if let Some(mw) = args.steady_state_min_window
|
||||
&& mw < 0.0
|
||||
{
|
||||
return Err(BenchError::Config(format!(
|
||||
@@ -613,122 +615,122 @@ impl BenchConfig {
|
||||
)));
|
||||
}
|
||||
|
||||
if cli.profile_batch_threshold.is_some() && !cli.profile {
|
||||
if args.profile_batch_threshold.is_some() && !args.profile {
|
||||
return Err(BenchError::Config(
|
||||
"--profile-batch-threshold requires --profile".into(),
|
||||
));
|
||||
}
|
||||
if cli.profile_duration <= 0.0 {
|
||||
if args.profile_duration <= 0.0 {
|
||||
return Err(BenchError::Config(
|
||||
"--profile-duration must be positive".into(),
|
||||
));
|
||||
}
|
||||
if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
|
||||
if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
|
||||
return Err(BenchError::Config(
|
||||
"--profile-duration requires --profile-batch-threshold".into(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(BenchConfig {
|
||||
backend: cli.backend,
|
||||
backend: args.backend,
|
||||
base_url,
|
||||
api_url,
|
||||
model: cli.model.clone(),
|
||||
model_name: cli.served_model_name.clone(),
|
||||
model: args.model.clone(),
|
||||
model_name: args.served_model_name.clone(),
|
||||
tokenizer_id,
|
||||
tokenizer_mode: cli.tokenizer_mode.clone(),
|
||||
trust_remote_code: cli.trust_remote_code,
|
||||
skip_tokenizer_init: cli.skip_tokenizer_init,
|
||||
dataset_name: cli.dataset_name,
|
||||
dataset_path: cli.dataset_path.clone(),
|
||||
max_model_len: cli.max_model_len,
|
||||
tokenizer_mode: args.tokenizer_mode.clone(),
|
||||
trust_remote_code: args.trust_remote_code,
|
||||
skip_tokenizer_init: args.skip_tokenizer_init,
|
||||
dataset_name: args.dataset_name,
|
||||
dataset_path: args.dataset_path.clone(),
|
||||
max_model_len: args.max_model_len,
|
||||
random_input_len,
|
||||
random_output_len,
|
||||
random_prefix_len: cli.random_prefix_len,
|
||||
random_prefix_len: args.random_prefix_len,
|
||||
random_range_ratio,
|
||||
random_batch_size: cli.random_batch_size,
|
||||
random_batch_size: args.random_batch_size,
|
||||
is_reranker,
|
||||
custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
|
||||
prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
|
||||
prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
|
||||
prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
|
||||
prefix_repetition_output_len: cli
|
||||
custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
|
||||
prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
|
||||
prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
|
||||
prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
|
||||
prefix_repetition_output_len: args
|
||||
.output_len
|
||||
.unwrap_or(cli.prefix_repetition_output_len),
|
||||
random_cache_hit_fraction: cli.random_cache_hit_fraction,
|
||||
random_cache_ratio: cli.random_cache_ratio,
|
||||
sharegpt_output_len: cli.sharegpt_output_len,
|
||||
sonnet_input_len: cli.sonnet_input_len,
|
||||
sonnet_output_len: cli.sonnet_output_len,
|
||||
sonnet_prefix_len: cli.sonnet_prefix_len,
|
||||
no_oversample: cli.no_oversample,
|
||||
disable_shuffle: cli.disable_shuffle,
|
||||
num_prompts: cli.num_prompts,
|
||||
request_rate: cli.request_rate,
|
||||
burstiness: cli.burstiness,
|
||||
max_concurrency: cli.max_concurrency,
|
||||
steady_state_threshold: cli.steady_state_threshold,
|
||||
steady_state_min_window: cli.steady_state_min_window,
|
||||
no_steady_state: cli.no_steady_state,
|
||||
disable_tqdm: cli.disable_tqdm,
|
||||
num_warmups: cli.num_warmups,
|
||||
profile: cli.profile,
|
||||
profile_batch_threshold: cli.profile_batch_threshold,
|
||||
profile_duration: cli.profile_duration,
|
||||
save_result: cli.save_result,
|
||||
save_detailed: cli.save_detailed,
|
||||
append_result: cli.append_result,
|
||||
result_dir: cli.result_dir.clone(),
|
||||
result_filename: cli.result_filename.clone(),
|
||||
seed: cli.seed,
|
||||
.unwrap_or(args.prefix_repetition_output_len),
|
||||
random_cache_hit_fraction: args.random_cache_hit_fraction,
|
||||
random_cache_ratio: args.random_cache_ratio,
|
||||
sharegpt_output_len: args.sharegpt_output_len,
|
||||
sonnet_input_len: args.sonnet_input_len,
|
||||
sonnet_output_len: args.sonnet_output_len,
|
||||
sonnet_prefix_len: args.sonnet_prefix_len,
|
||||
no_oversample: args.no_oversample,
|
||||
disable_shuffle: args.disable_shuffle,
|
||||
num_prompts: args.num_prompts,
|
||||
request_rate: args.request_rate,
|
||||
burstiness: args.burstiness,
|
||||
max_concurrency: args.max_concurrency,
|
||||
steady_state_threshold: args.steady_state_threshold,
|
||||
steady_state_min_window: args.steady_state_min_window,
|
||||
no_steady_state: args.no_steady_state,
|
||||
disable_tqdm: args.disable_tqdm,
|
||||
num_warmups: args.num_warmups,
|
||||
profile: args.profile,
|
||||
profile_batch_threshold: args.profile_batch_threshold,
|
||||
profile_duration: args.profile_duration,
|
||||
save_result: args.save_result,
|
||||
save_detailed: args.save_detailed,
|
||||
append_result: args.append_result,
|
||||
result_dir: args.result_dir.clone(),
|
||||
result_filename: args.result_filename.clone(),
|
||||
seed: args.seed,
|
||||
ignore_eos,
|
||||
insecure: cli.insecure,
|
||||
insecure: args.insecure,
|
||||
selected_percentile_metrics,
|
||||
selected_percentiles,
|
||||
sweep_summary_percentiles,
|
||||
label: cli.label.clone(),
|
||||
logprobs: cli.logprobs,
|
||||
request_id_prefix: cli.get_request_id_prefix(),
|
||||
ready_check_timeout_sec: cli.ready_check_timeout_sec,
|
||||
label: args.label.clone(),
|
||||
logprobs: args.logprobs,
|
||||
request_id_prefix: args.get_request_id_prefix(),
|
||||
ready_check_timeout_sec: args.ready_check_timeout_sec,
|
||||
extra_headers,
|
||||
extra_body,
|
||||
metadata,
|
||||
dry_run: cli.dry_run,
|
||||
dry_run: args.dry_run,
|
||||
goodput,
|
||||
ramp_up,
|
||||
multi_turn: cli.multi_turn,
|
||||
multi_turn_num_turns: cli.multi_turn_num_turns,
|
||||
multi_turn: args.multi_turn,
|
||||
multi_turn_num_turns: args.multi_turn_num_turns,
|
||||
multi_turn_min_turns,
|
||||
multi_turn_max_turns,
|
||||
sharegpt_multi_turn_max_turns: if cli.multi_turn
|
||||
&& cli.dataset_name == DatasetName::ShareGpt
|
||||
&& cli.multi_turn_max_turns != 0
|
||||
sharegpt_multi_turn_max_turns: if args.multi_turn
|
||||
&& args.dataset_name == DatasetName::ShareGpt
|
||||
&& args.multi_turn_max_turns != 0
|
||||
{
|
||||
Some(cli.multi_turn_max_turns)
|
||||
Some(args.multi_turn_max_turns)
|
||||
} else {
|
||||
None
|
||||
},
|
||||
per_turn_input_len,
|
||||
multi_turn_concurrency: cli.multi_turn_concurrency,
|
||||
multi_turn_delay_ms: cli.multi_turn_delay_ms,
|
||||
multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
|
||||
multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
|
||||
speed_bench_config: cli.speed_bench_config,
|
||||
speed_bench_category: cli.speed_bench_category.clone(),
|
||||
speed_bench_max_input_len: cli.speed_bench_max_input_len,
|
||||
hf_split: cli.hf_split.clone(),
|
||||
hf_subset: cli.hf_subset.clone(),
|
||||
hf_output_len: cli.hf_output_len,
|
||||
hf_text_column: cli.hf_text_column.clone(),
|
||||
reset_prefix_cache: cli.reset_prefix_cache,
|
||||
prompt_token_ids: cli.prompt_token_ids,
|
||||
random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
|
||||
random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
|
||||
multi_turn_concurrency: args.multi_turn_concurrency,
|
||||
multi_turn_delay_ms: args.multi_turn_delay_ms,
|
||||
multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
|
||||
multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
|
||||
speed_bench_config: args.speed_bench_config,
|
||||
speed_bench_category: args.speed_bench_category.clone(),
|
||||
speed_bench_max_input_len: args.speed_bench_max_input_len,
|
||||
hf_split: args.hf_split.clone(),
|
||||
hf_subset: args.hf_subset.clone(),
|
||||
hf_output_len: args.hf_output_len,
|
||||
hf_text_column: args.hf_text_column.clone(),
|
||||
reset_prefix_cache: args.reset_prefix_cache,
|
||||
prompt_token_ids: args.prompt_token_ids,
|
||||
random_mm_base_items_per_request: args.random_mm_base_items_per_request,
|
||||
random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
|
||||
random_mm_limit,
|
||||
random_mm_buckets,
|
||||
enable_multimodal_chat: cli.enable_multimodal_chat,
|
||||
enable_multimodal_chat: args.enable_multimodal_chat,
|
||||
lora_modules,
|
||||
lora_assignment: cli.lora_assignment,
|
||||
lora_assignment: args.lora_assignment,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -811,17 +813,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
|
||||
let strategy = match cli.ramp_up_strategy {
|
||||
fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
|
||||
let strategy = match args.ramp_up_strategy {
|
||||
None => return Ok(None),
|
||||
Some(s) => s,
|
||||
};
|
||||
|
||||
let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
|
||||
let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
|
||||
BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
|
||||
})?;
|
||||
|
||||
let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
|
||||
let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
|
||||
BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
|
||||
})?;
|
||||
|
||||
@@ -843,7 +845,21 @@ mod tests {
|
||||
use clap::Parser;
|
||||
|
||||
use super::*;
|
||||
use crate::cli::Cli;
|
||||
use crate::cli::BenchServeArgs;
|
||||
|
||||
#[derive(Parser)]
|
||||
struct TestCli {
|
||||
#[command(flatten)]
|
||||
args: BenchServeArgs,
|
||||
}
|
||||
|
||||
fn parse_args<I, T>(args: I) -> BenchServeArgs
|
||||
where
|
||||
I: IntoIterator<Item = T>,
|
||||
T: Into<std::ffi::OsString> + Clone,
|
||||
{
|
||||
TestCli::parse_from(args).args
|
||||
}
|
||||
|
||||
fn base_multi_turn_args() -> Vec<&'static str> {
|
||||
vec![
|
||||
@@ -859,8 +875,8 @@ mod tests {
|
||||
#[test]
|
||||
fn test_prefix_sharing_defaults_to_zero() {
|
||||
let args = base_multi_turn_args();
|
||||
let cli = Cli::parse_from(args);
|
||||
let config = BenchConfig::from_cli(&cli).unwrap();
|
||||
let args = parse_args(args);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
|
||||
assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
|
||||
}
|
||||
@@ -874,8 +890,8 @@ mod tests {
|
||||
"--multi-turn-prefix-conversation-ratio",
|
||||
"0.8",
|
||||
]);
|
||||
let cli = Cli::parse_from(args);
|
||||
let config = BenchConfig::from_cli(&cli).unwrap();
|
||||
let args = parse_args(args);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
|
||||
assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
|
||||
}
|
||||
@@ -889,8 +905,8 @@ mod tests {
|
||||
"--multi-turn-prefix-conversation-ratio",
|
||||
"0.6",
|
||||
]);
|
||||
let cli = Cli::parse_from(args);
|
||||
assert!(BenchConfig::from_cli(&cli).is_err());
|
||||
let args = parse_args(args);
|
||||
assert!(BenchConfig::from_args(&args).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -902,16 +918,16 @@ mod tests {
|
||||
"--multi-turn-prefix-conversation-ratio",
|
||||
"0.5",
|
||||
]);
|
||||
let cli = Cli::parse_from(args);
|
||||
assert!(BenchConfig::from_cli(&cli).is_err());
|
||||
let args = parse_args(args);
|
||||
assert!(BenchConfig::from_args(&args).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prefix_sharing_out_of_range_fails() {
|
||||
let mut args = base_multi_turn_args();
|
||||
args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
|
||||
let cli = Cli::parse_from(args);
|
||||
assert!(BenchConfig::from_cli(&cli).is_err());
|
||||
let args = parse_args(args);
|
||||
assert!(BenchConfig::from_args(&args).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -928,8 +944,8 @@ mod tests {
|
||||
"--multi-turn-prefix-global-ratio",
|
||||
"0.1",
|
||||
];
|
||||
let cli = Cli::parse_from(args);
|
||||
assert!(BenchConfig::from_cli(&cli).is_err());
|
||||
let args = parse_args(args);
|
||||
assert!(BenchConfig::from_args(&args).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -944,8 +960,8 @@ mod tests {
|
||||
"--dataset-name",
|
||||
"sharegpt",
|
||||
];
|
||||
let cli = Cli::parse_from(args);
|
||||
let config = BenchConfig::from_cli(&cli).unwrap();
|
||||
let args = parse_args(args);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
|
||||
assert_eq!(config.multi_turn_max_turns, 3);
|
||||
assert_eq!(config.sharegpt_multi_turn_max_turns, None);
|
||||
@@ -968,8 +984,8 @@ mod tests {
|
||||
"--multi-turn-max-turns",
|
||||
"2",
|
||||
];
|
||||
let cli = Cli::parse_from(args);
|
||||
let config = BenchConfig::from_cli(&cli).unwrap();
|
||||
let args = parse_args(args);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
|
||||
}
|
||||
|
||||
@@ -987,8 +1003,8 @@ mod tests {
|
||||
"--multi-turn-max-turns",
|
||||
"1",
|
||||
];
|
||||
let cli = Cli::parse_from(args);
|
||||
let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
|
||||
let args = parse_args(args);
|
||||
let err = BenchConfig::from_args(&args).unwrap_err().to_string();
|
||||
assert!(
|
||||
err.contains("at least 2 for ShareGPT"),
|
||||
"expected ShareGPT-specific error, got: {err}"
|
||||
@@ -1009,8 +1025,8 @@ mod tests {
|
||||
"--multi-turn-max-turns",
|
||||
"20",
|
||||
];
|
||||
let cli = Cli::parse_from(args);
|
||||
let config = BenchConfig::from_cli(&cli).unwrap();
|
||||
let args = parse_args(args);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
|
||||
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
|
||||
}
|
||||
@@ -1018,8 +1034,8 @@ mod tests {
|
||||
#[test]
|
||||
fn test_sweep_summary_percentiles_default_empty() {
|
||||
let args = base_multi_turn_args();
|
||||
let cli = Cli::parse_from(args);
|
||||
let config = BenchConfig::from_cli(&cli).unwrap();
|
||||
let args = parse_args(args);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
|
||||
assert!(config.sweep_summary_percentiles.is_empty());
|
||||
assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
|
||||
@@ -1034,8 +1050,8 @@ mod tests {
|
||||
"--sweep-summary-percentiles",
|
||||
"90,95,90",
|
||||
]);
|
||||
let cli = Cli::parse_from(args);
|
||||
let config = BenchConfig::from_cli(&cli).unwrap();
|
||||
let args = parse_args(args);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
|
||||
assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
|
||||
assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
|
||||
@@ -1045,8 +1061,8 @@ mod tests {
|
||||
fn test_invalid_sweep_summary_percentile_fails() {
|
||||
let mut args = base_multi_turn_args();
|
||||
args.extend(["--sweep-summary-percentiles", "101"]);
|
||||
let cli = Cli::parse_from(args);
|
||||
assert!(BenchConfig::from_cli(&cli).is_err());
|
||||
let args = parse_args(args);
|
||||
assert!(BenchConfig::from_args(&args).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1058,12 +1074,20 @@ mod tests {
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
];
|
||||
let cli = Cli::parse_from(args);
|
||||
let config = BenchConfig::from_cli(&cli).unwrap();
|
||||
let args = parse_args(args);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
|
||||
assert_eq!(config.max_model_len, Some(4096));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tokenizer_id_deferred_when_model_is_unspecified() {
|
||||
let args = parse_args(["vllm-bench"]);
|
||||
let config = BenchConfig::from_args(&args).unwrap();
|
||||
|
||||
assert_eq!(config.tokenizer_id, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_zero_max_model_len_fails() {
|
||||
let args = vec![
|
||||
@@ -1073,9 +1097,9 @@ mod tests {
|
||||
"--max-model-len",
|
||||
"0",
|
||||
];
|
||||
let cli = Cli::parse_from(args);
|
||||
let args = parse_args(args);
|
||||
|
||||
assert!(BenchConfig::from_cli(&cli).is_err());
|
||||
assert!(BenchConfig::from_args(&args).is_err());
|
||||
}
|
||||
#[test]
|
||||
fn test_range_ratio_parse_float() {
|
||||
|
||||
@@ -40,8 +40,11 @@ impl HubRepo {
|
||||
.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}"))?;
|
||||
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
|
||||
if let Ok(token) = std::env::var("HF_TOKEN") {
|
||||
builder = builder.with_token(Some(token));
|
||||
}
|
||||
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
|
||||
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
|
||||
})
|
||||
})
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
mod backends;
|
||||
mod benchmark;
|
||||
mod cli;
|
||||
mod compare;
|
||||
mod config;
|
||||
mod datasets;
|
||||
mod error;
|
||||
mod hub;
|
||||
mod metrics;
|
||||
mod multi_run;
|
||||
mod multi_turn;
|
||||
mod output;
|
||||
mod rate_control;
|
||||
mod ready_checker;
|
||||
mod sweep;
|
||||
mod tiktoken;
|
||||
mod tokenizer;
|
||||
|
||||
use anyhow::Context;
|
||||
|
||||
pub use cli::{
|
||||
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
|
||||
};
|
||||
use config::BenchConfig;
|
||||
|
||||
/// Prepare process-wide resources for a benchmark run.
|
||||
pub fn prepare_process() {
|
||||
// Raise the open-file soft limit to the hard limit. High-concurrency
|
||||
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
|
||||
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
|
||||
&& new > 1024
|
||||
{
|
||||
eprintln!("Open-file limit: {new}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the online serving benchmark.
|
||||
pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
|
||||
// --- Compare mode: no server needed, just diff two JSON files ---
|
||||
if let Some(ref files) = args.compare {
|
||||
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
|
||||
}
|
||||
|
||||
let config = BenchConfig::from_args(&args).context("Configuration error")?;
|
||||
|
||||
async {
|
||||
if config.multi_turn {
|
||||
if let Some(ref sweep_mc) = args.sweep_max_concurrency {
|
||||
// --- Sweep over concurrency in multi-turn mode ---
|
||||
let values = sweep::parse_concurrency_values(sweep_mc)
|
||||
.context("Invalid --sweep-max-concurrency")?;
|
||||
sweep::run_multi_turn_concurrency_sweep(
|
||||
&config,
|
||||
&values,
|
||||
args.sweep_num_prompts_factor,
|
||||
)
|
||||
.await?;
|
||||
} else {
|
||||
// --- Single multi-turn conversation benchmark ---
|
||||
multi_turn::run_multi_turn_benchmark(&config).await?;
|
||||
}
|
||||
} else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
|
||||
// --- Sweep over max-concurrency ---
|
||||
let values = sweep::parse_concurrency_values(sweep_mc)
|
||||
.context("Invalid --sweep-max-concurrency")?;
|
||||
sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
|
||||
} else if let Some(ref sweep_rate) = args.sweep_request_rate {
|
||||
// --- Sweep over request-rate ---
|
||||
let values =
|
||||
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
|
||||
sweep::run_rate_sweep(&config, &values).await?;
|
||||
} else if args.num_runs > 1 {
|
||||
// --- Multi-run with statistical aggregation ---
|
||||
multi_run::run_multi(&config, args.num_runs).await?;
|
||||
} else {
|
||||
// --- Normal single benchmark ---
|
||||
benchmark::run_benchmark(&config).await?;
|
||||
}
|
||||
anyhow::Ok(())
|
||||
}
|
||||
.await
|
||||
.context("Benchmark failed")
|
||||
}
|
||||
+14
-74
@@ -1,92 +1,32 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
mod backends;
|
||||
mod benchmark;
|
||||
mod cli;
|
||||
mod compare;
|
||||
mod config;
|
||||
mod datasets;
|
||||
mod error;
|
||||
mod hub;
|
||||
mod metrics;
|
||||
mod multi_run;
|
||||
mod multi_turn;
|
||||
mod output;
|
||||
mod rate_control;
|
||||
mod ready_checker;
|
||||
mod sweep;
|
||||
mod tiktoken;
|
||||
mod tokenizer;
|
||||
|
||||
#[cfg(not(target_env = "msvc"))]
|
||||
#[global_allocator]
|
||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||
|
||||
use anyhow::Context;
|
||||
use clap::Parser;
|
||||
use cli::Cli;
|
||||
use config::BenchConfig;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(
|
||||
name = "vllm-bench",
|
||||
about = "Benchmark online serving throughput",
|
||||
version
|
||||
)]
|
||||
struct Cli {
|
||||
#[command(flatten)]
|
||||
args: vllm_bench::BenchServeArgs,
|
||||
}
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
// Raise the open-file soft limit to the hard limit. High-concurrency
|
||||
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
|
||||
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
|
||||
&& new > 1024
|
||||
{
|
||||
eprintln!("Open-file limit: {new}");
|
||||
}
|
||||
|
||||
let cli = Cli::parse();
|
||||
|
||||
// --- Compare mode: no server needed, just diff two JSON files ---
|
||||
if let Some(ref files) = cli.compare {
|
||||
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
|
||||
}
|
||||
|
||||
let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
|
||||
vllm_bench::prepare_process();
|
||||
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("Failed to build tokio runtime");
|
||||
.context("Failed to build tokio runtime")?;
|
||||
|
||||
runtime
|
||||
.block_on(async {
|
||||
if config.multi_turn {
|
||||
if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
|
||||
// --- Sweep over concurrency in multi-turn mode ---
|
||||
let values = sweep::parse_concurrency_values(sweep_mc)
|
||||
.context("Invalid --sweep-max-concurrency")?;
|
||||
sweep::run_multi_turn_concurrency_sweep(
|
||||
&config,
|
||||
&values,
|
||||
cli.sweep_num_prompts_factor,
|
||||
)
|
||||
.await?;
|
||||
} else {
|
||||
// --- Single multi-turn conversation benchmark ---
|
||||
multi_turn::run_multi_turn_benchmark(&config).await?;
|
||||
}
|
||||
} else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
|
||||
// --- Sweep over max-concurrency ---
|
||||
let values = sweep::parse_concurrency_values(sweep_mc)
|
||||
.context("Invalid --sweep-max-concurrency")?;
|
||||
sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
|
||||
.await?;
|
||||
} else if let Some(ref sweep_rate) = cli.sweep_request_rate {
|
||||
// --- Sweep over request-rate ---
|
||||
let values =
|
||||
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
|
||||
sweep::run_rate_sweep(&config, &values).await?;
|
||||
} else if cli.num_runs > 1 {
|
||||
// --- Multi-run with statistical aggregation ---
|
||||
multi_run::run_multi(&config, cli.num_runs).await?;
|
||||
} else {
|
||||
// --- Normal single benchmark ---
|
||||
benchmark::run_benchmark(&config).await?;
|
||||
}
|
||||
anyhow::Ok(())
|
||||
})
|
||||
.context("Benchmark failed")
|
||||
runtime.block_on(vllm_bench::run(cli.args))
|
||||
}
|
||||
|
||||
@@ -38,7 +38,7 @@ pub(super) fn build_batched_items(
|
||||
let keep_on_cpu = spec.keep_on_cpu_keys.contains(key);
|
||||
let (value, field) = match spec.field_layout_for(key) {
|
||||
Some(FieldLayout::Batched) => (
|
||||
tensor.batched_value_at(index)?,
|
||||
tensor.batched_wire_value_at(index)?,
|
||||
MmField::Batched(MmBatchedField { keep_on_cpu }),
|
||||
),
|
||||
Some(FieldLayout::Flat { sizes_key }) => {
|
||||
@@ -47,7 +47,7 @@ pub(super) fn build_batched_items(
|
||||
})?;
|
||||
let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
|
||||
(
|
||||
tensor.flat_value_range(start, end)?,
|
||||
tensor.flat_wire_value_range(start, end)?,
|
||||
MmField::Flat(MmFlatField {
|
||||
slices: vec![MmSlice::Slice(SliceSpec {
|
||||
start: Some(0),
|
||||
@@ -60,7 +60,7 @@ pub(super) fn build_batched_items(
|
||||
)
|
||||
}
|
||||
None => (
|
||||
tensor.clone(),
|
||||
tensor.try_into()?,
|
||||
MmField::Shared(MmSharedField {
|
||||
batch_size: len,
|
||||
keep_on_cpu,
|
||||
@@ -71,7 +71,7 @@ pub(super) fn build_batched_items(
|
||||
data.insert(
|
||||
key.clone(),
|
||||
MmFieldElem {
|
||||
data: Some(value.try_into()?),
|
||||
data: Some(value),
|
||||
field,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -12,7 +12,7 @@ use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
|
||||
use crate::error::{Error, Result, bail_multimodal, multimodal};
|
||||
|
||||
/// Representation for multimodal kwarg values for transformation.
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Debug)]
|
||||
pub(super) enum KwargValue {
|
||||
/// Float tensor with row-major flat data and shape.
|
||||
F32Tensor { data: Vec<f32>, shape: Vec<usize> },
|
||||
@@ -107,28 +107,19 @@ impl KwargValue {
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<KwargValue> for ProtocolKwargValue {
|
||||
impl TryFrom<&KwargValue> for ProtocolKwargValue {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(value: KwargValue) -> Result<Self> {
|
||||
match value {
|
||||
KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor(
|
||||
WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?,
|
||||
)),
|
||||
KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor(
|
||||
WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?,
|
||||
)),
|
||||
KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor(
|
||||
WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?,
|
||||
)),
|
||||
KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
|
||||
WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
|
||||
)),
|
||||
KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
|
||||
WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
|
||||
)),
|
||||
KwargValue::Passthrough(value) => Ok(value),
|
||||
}
|
||||
fn try_from(value: &KwargValue) -> Result<Self> {
|
||||
let tensor = match value {
|
||||
KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
|
||||
KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
|
||||
KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
|
||||
KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
|
||||
KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
|
||||
KwargValue::Passthrough(value) => return Ok(value.clone()),
|
||||
};
|
||||
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -145,63 +136,55 @@ impl KwargValue {
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract one media item from a batched tensor field.
|
||||
/// Convert one media item from a batched tensor field to wire bytes.
|
||||
///
|
||||
/// Batched fields use their first axis as media-item index and drop that
|
||||
/// axis in the per-feature value, matching vLLM's batched-field semantics.
|
||||
pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> {
|
||||
match self {
|
||||
Self::F32Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
|
||||
Ok(Self::F32Tensor { data, shape })
|
||||
}
|
||||
Self::F16Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
|
||||
Ok(Self::F16Tensor { data, shape })
|
||||
}
|
||||
Self::Bf16Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
|
||||
Ok(Self::Bf16Tensor { data, shape })
|
||||
}
|
||||
Self::I64Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
|
||||
Ok(Self::I64Tensor { data, shape })
|
||||
}
|
||||
Self::U32Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
|
||||
Ok(Self::U32Tensor { data, shape })
|
||||
}
|
||||
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
|
||||
}
|
||||
pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
|
||||
self.wire_value_range(index, index + 1, true)
|
||||
}
|
||||
|
||||
/// Extract one media item's variable-length range from a flat tensor field.
|
||||
/// Convert one media item's flat tensor range directly to wire bytes.
|
||||
///
|
||||
/// Flat fields keep the first axis as the sliced length for this item.
|
||||
pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> {
|
||||
match self {
|
||||
pub(super) fn flat_wire_value_range(
|
||||
&self,
|
||||
start: usize,
|
||||
end: usize,
|
||||
) -> Result<ProtocolKwargValue> {
|
||||
self.wire_value_range(start, end, false)
|
||||
}
|
||||
|
||||
fn wire_value_range(
|
||||
&self,
|
||||
start: usize,
|
||||
end: usize,
|
||||
drop_axis: bool,
|
||||
) -> Result<ProtocolKwargValue> {
|
||||
let tensor = match self {
|
||||
Self::F32Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
|
||||
Ok(Self::F32Tensor { data, shape })
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
|
||||
WireTensor::from_f32(shape, data)
|
||||
}
|
||||
Self::F16Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
|
||||
Ok(Self::F16Tensor { data, shape })
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
|
||||
WireTensor::from_f16(shape, data)
|
||||
}
|
||||
Self::Bf16Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
|
||||
Ok(Self::Bf16Tensor { data, shape })
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
|
||||
WireTensor::from_bf16(shape, data)
|
||||
}
|
||||
Self::I64Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
|
||||
Ok(Self::I64Tensor { data, shape })
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
|
||||
WireTensor::from_i64(shape, data)
|
||||
}
|
||||
Self::U32Tensor { data, shape } => {
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
|
||||
Ok(Self::U32Tensor { data, shape })
|
||||
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
|
||||
WireTensor::from_u32(shape, data)
|
||||
}
|
||||
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
|
||||
}
|
||||
Self::Passthrough(value) => return Ok(value.clone()),
|
||||
};
|
||||
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -240,13 +223,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
|
||||
}
|
||||
|
||||
/// Slice a flat row-major tensor along its first axis.
|
||||
fn slice_first_axis_range<T: Clone>(
|
||||
fn slice_first_axis_range<'a, T>(
|
||||
shape: &[usize],
|
||||
data: &[T],
|
||||
data: &'a [T],
|
||||
start: usize,
|
||||
end: usize,
|
||||
drop_axis: bool,
|
||||
) -> Result<(Vec<usize>, Vec<T>)> {
|
||||
) -> Result<(Vec<usize>, &'a [T])> {
|
||||
let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
|
||||
if start > end || end > first_dim {
|
||||
bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
|
||||
@@ -270,7 +253,7 @@ fn slice_first_axis_range<T: Clone>(
|
||||
shape[0] = end - start;
|
||||
shape
|
||||
};
|
||||
Ok((out_shape, data[data_start..data_end].to_vec()))
|
||||
Ok((out_shape, &data[data_start..data_end]))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -278,35 +261,39 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn batched_value_at_drops_first_axis() {
|
||||
fn batched_wire_value_at_drops_first_axis() {
|
||||
let value = KwargValue::F32Tensor {
|
||||
data: vec![1.0, 2.0, 3.0, 4.0],
|
||||
shape: vec![2, 2],
|
||||
};
|
||||
|
||||
let value = value.batched_value_at(1).unwrap();
|
||||
let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
|
||||
panic!("expected tensor");
|
||||
};
|
||||
|
||||
assert!(matches!(
|
||||
value,
|
||||
KwargValue::F32Tensor { data, shape }
|
||||
if shape == vec![2] && data == vec![3.0, 4.0]
|
||||
));
|
||||
assert_eq!(tensor.shape, vec![2]);
|
||||
assert_eq!(
|
||||
tensor.data.into_raw_view().unwrap(),
|
||||
[3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flat_value_range_keeps_first_axis() {
|
||||
fn flat_wire_value_range_keeps_first_axis() {
|
||||
let value = KwargValue::U32Tensor {
|
||||
data: (0..10).collect(),
|
||||
shape: vec![5, 2],
|
||||
};
|
||||
|
||||
let value = value.flat_value_range(1, 3).unwrap();
|
||||
let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
|
||||
panic!("expected tensor");
|
||||
};
|
||||
|
||||
assert!(matches!(
|
||||
value,
|
||||
KwargValue::U32Tensor { data, shape }
|
||||
if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
|
||||
));
|
||||
assert_eq!(tensor.shape, vec![2, 2]);
|
||||
assert_eq!(
|
||||
tensor.data.into_raw_view().unwrap(),
|
||||
[2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::<Vec<_>>()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -336,7 +323,7 @@ mod tests {
|
||||
let value =
|
||||
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
|
||||
|
||||
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
|
||||
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
|
||||
else {
|
||||
panic!("expected tensor");
|
||||
};
|
||||
@@ -351,7 +338,7 @@ mod tests {
|
||||
let value =
|
||||
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
|
||||
|
||||
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
|
||||
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
|
||||
else {
|
||||
panic!("expected tensor");
|
||||
};
|
||||
|
||||
@@ -130,7 +130,7 @@ fn build_video_item(
|
||||
let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key);
|
||||
let (value, field) = match support.spec.field_layout_for(&key) {
|
||||
Some(FieldLayout::Batched) => (
|
||||
tensor.batched_value_at(0)?,
|
||||
tensor.batched_wire_value_at(0)?,
|
||||
MmField::Batched(MmBatchedField { keep_on_cpu }),
|
||||
),
|
||||
Some(FieldLayout::Flat { .. }) => {
|
||||
@@ -138,7 +138,7 @@ fn build_video_item(
|
||||
.first_dim()
|
||||
.ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?;
|
||||
(
|
||||
tensor,
|
||||
(&tensor).try_into()?,
|
||||
MmField::Flat(MmFlatField {
|
||||
slices: vec![MmSlice::Slice(SliceSpec {
|
||||
start: Some(0),
|
||||
@@ -151,7 +151,7 @@ fn build_video_item(
|
||||
)
|
||||
}
|
||||
None => (
|
||||
tensor,
|
||||
(&tensor).try_into()?,
|
||||
MmField::Shared(MmSharedField {
|
||||
batch_size: 1,
|
||||
keep_on_cpu,
|
||||
@@ -162,7 +162,7 @@ fn build_video_item(
|
||||
data.insert(
|
||||
key,
|
||||
MmFieldElem {
|
||||
data: Some(value.try_into()?),
|
||||
data: Some(value),
|
||||
field,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor {
|
||||
Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
|
||||
};
|
||||
|
||||
apply_structural_tag_constraint(request, parser.structural_tag_model())?;
|
||||
apply_structural_tag_constraint(request, parser.structural_tag_builder())?;
|
||||
|
||||
if parser.preserve_special_tokens() {
|
||||
request.decode_options.skip_special_tokens = false;
|
||||
|
||||
@@ -7,7 +7,8 @@ use thiserror_ext::AsReport;
|
||||
use vllm_engine_core_client::protocol::structured_outputs::{
|
||||
StructuredOutputBackend, StructuredOutputsParams,
|
||||
};
|
||||
use vllm_parser::tool::StructuralTagModel;
|
||||
use vllm_parser::tool::StructuralTagBuilder;
|
||||
use xgrammar_structural_tag::builders::StructuralTagOptions;
|
||||
use xgrammar_structural_tag::{
|
||||
FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
|
||||
build_structural_tag,
|
||||
@@ -20,9 +21,9 @@ use crate::{Error, Result as ChatResult};
|
||||
/// support and the request's tool choice.
|
||||
pub(super) fn apply_structural_tag_constraint(
|
||||
request: &mut ChatRequest,
|
||||
model: Option<StructuralTagModel>,
|
||||
builder: Option<&dyn StructuralTagBuilder>,
|
||||
) -> ChatResult<()> {
|
||||
let Some(model) = model else {
|
||||
let Some(builder) = builder else {
|
||||
return Ok(());
|
||||
};
|
||||
let Some(tool_choice) = structural_tag_tool_choice(request) else {
|
||||
@@ -42,11 +43,16 @@ pub(super) fn apply_structural_tag_constraint(
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let structural_tag = build_structural_tag(model, &tools, tool_choice, false)
|
||||
.and_then(|tag| tag.to_json_string())
|
||||
.map_err(|error| Error::StructuralTag {
|
||||
message: error.to_report_string(),
|
||||
})?;
|
||||
let structural_tag = build_structural_tag(
|
||||
builder,
|
||||
&tools,
|
||||
tool_choice,
|
||||
StructuralTagOptions::default().with_reasoning(false),
|
||||
)
|
||||
.and_then(|tag| tag.to_json_string())
|
||||
.map_err(|error| Error::StructuralTag {
|
||||
message: error.to_report_string(),
|
||||
})?;
|
||||
|
||||
// Overwrite any existing structured output settings with the structural tag constraint.
|
||||
request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
|
||||
@@ -141,7 +147,7 @@ mod tests {
|
||||
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request);
|
||||
@@ -154,7 +160,7 @@ mod tests {
|
||||
let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
assert!(request.sampling_params.structured_outputs.is_none());
|
||||
@@ -169,7 +175,7 @@ mod tests {
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
@@ -184,7 +190,7 @@ mod tests {
|
||||
let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request);
|
||||
@@ -201,7 +207,7 @@ mod tests {
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
@@ -221,7 +227,7 @@ mod tests {
|
||||
);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
.expect("structural tag should build");
|
||||
|
||||
let tag = structural_tag_value(&request).to_string();
|
||||
@@ -234,7 +240,7 @@ mod tests {
|
||||
let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
assert!(request.sampling_params.structured_outputs.is_none());
|
||||
@@ -249,7 +255,7 @@ mod tests {
|
||||
});
|
||||
let parser = qwen3_coder_parser(&request.tools);
|
||||
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
|
||||
apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
|
||||
.expect("structural tag decision should succeed");
|
||||
|
||||
let params = structured_outputs(&request);
|
||||
|
||||
@@ -236,9 +236,10 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool {
|
||||
|
||||
loops.into_iter().any(|loop_ast| {
|
||||
matches!(loop_ast.target, Expr::Var(_))
|
||||
&& message_varnames
|
||||
.iter()
|
||||
.any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content")))
|
||||
&& (is_var_access(&loop_ast.iter, "content")
|
||||
|| message_varnames.iter().any(|varname| {
|
||||
is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -315,6 +316,16 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_openai_template_with_content_parameter_loop() {
|
||||
assert_eq!(
|
||||
detect(
|
||||
"{% macro render(content) %}{% for item in content %}{{ item }}{% endfor %}{% endmacro %}{% for message in messages %}{{ render(message.content) }}{% endfor %}"
|
||||
),
|
||||
ChatTemplateContentFormat::OpenAi
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_openai_template_with_messages_alias() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -1309,6 +1309,26 @@ mod tests {
|
||||
.assert_eq(&rendered);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen35_template_auto_detects_openai_multimodal_content() {
|
||||
let mut request = image_request();
|
||||
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
|
||||
|
||||
let rendered = render_mm(
|
||||
QWEN3_5_0_8B_TEMPLATE,
|
||||
&request,
|
||||
ChatTemplateContentFormatOption::Auto,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
expect![[r#"
|
||||
Text(
|
||||
"<|im_start|>user\na<|vision_start|><|image_pad|><|vision_end|>b<|im_end|>\n",
|
||||
)
|
||||
"#]]
|
||||
.assert_debug_eq(&rendered.prompt);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() {
|
||||
let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
|
||||
|
||||
@@ -29,6 +29,7 @@ tokio-util.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
uuid.workspace = true
|
||||
vllm-bench.workspace = true
|
||||
vllm-chat.workspace = true
|
||||
vllm-engine-core-client.workspace = true
|
||||
vllm-managed-engine.workspace = true
|
||||
|
||||
+11
-1
@@ -79,13 +79,23 @@ impl Cli {
|
||||
}
|
||||
|
||||
/// Supported top-level CLI commands.
|
||||
#[derive(Debug, Subcommand, PartialEq, Eq)]
|
||||
#[derive(Debug, Subcommand)]
|
||||
pub enum Command {
|
||||
/// Run the Rust OpenAI frontend as a Python-supervised worker.
|
||||
Frontend(FrontendArgs),
|
||||
/// Launch a managed Python headless engine, then run the Rust OpenAI
|
||||
/// frontend.
|
||||
Serve(ServeArgs),
|
||||
/// Run vLLM benchmarks.
|
||||
#[command(subcommand)]
|
||||
Bench(BenchCommand),
|
||||
}
|
||||
|
||||
/// Supported benchmark commands.
|
||||
#[derive(Debug, Subcommand)]
|
||||
pub enum BenchCommand {
|
||||
/// Benchmark online serving throughput.
|
||||
Serve(vllm_bench::BenchServeArgs),
|
||||
}
|
||||
|
||||
/// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
|
||||
|
||||
@@ -5,7 +5,27 @@ use expect_test::expect;
|
||||
use vllm_engine_core_client::TransportMode;
|
||||
use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
|
||||
|
||||
use super::{Cli, Command};
|
||||
use super::{BenchCommand, Cli, Command};
|
||||
|
||||
#[test]
|
||||
fn bench_serve_args_parse_without_managed_engine_repartition() {
|
||||
let cli = Cli::try_parse_from([
|
||||
"vllm-rs",
|
||||
"bench",
|
||||
"serve",
|
||||
"--backend",
|
||||
"openai-chat",
|
||||
"--request-rate",
|
||||
"inf",
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
|
||||
panic!("expected bench serve args");
|
||||
};
|
||||
assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
|
||||
assert!(args.request_rate.is_infinite());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serve_args_forward_python_flags_with_separator() {
|
||||
|
||||
@@ -12,7 +12,7 @@ use tokio_util::sync::CancellationToken;
|
||||
use tracing::{info, warn};
|
||||
use vllm_managed_engine::ManagedEngineHandle;
|
||||
|
||||
use crate::cli::{Cli, Command};
|
||||
use crate::cli::{BenchCommand, Cli, Command};
|
||||
|
||||
#[global_allocator]
|
||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||
@@ -100,6 +100,10 @@ fn main() -> Result<()> {
|
||||
async fn async_main(cli: Cli) -> Result<()> {
|
||||
match cli.command {
|
||||
Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
|
||||
Command::Bench(BenchCommand::Serve(bench_args)) => {
|
||||
vllm_bench::prepare_process();
|
||||
vllm_bench::run(bench_args).await
|
||||
}
|
||||
Command::Serve(args) => {
|
||||
let handshake_port = args.managed_engine.resolve_handshake_port()?;
|
||||
|
||||
|
||||
@@ -55,52 +55,57 @@ pub struct WireNdArray {
|
||||
|
||||
impl WireNdArray {
|
||||
/// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
|
||||
pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
|
||||
pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
|
||||
let data = data.as_ref();
|
||||
validate_element_count(&shape, data.len())?;
|
||||
Ok(Self {
|
||||
dtype: "float32".to_string(),
|
||||
shape,
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(&data)),
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
|
||||
pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
|
||||
pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
|
||||
let data = data.as_ref();
|
||||
validate_element_count(&shape, data.len())?;
|
||||
Ok(Self {
|
||||
dtype: "float16".to_string(),
|
||||
shape,
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(&data)),
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
|
||||
pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
|
||||
pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
|
||||
let data = data.as_ref();
|
||||
validate_element_count(&shape, data.len())?;
|
||||
Ok(Self {
|
||||
dtype: "bfloat16".to_string(),
|
||||
shape,
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(&data)),
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
|
||||
pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
|
||||
pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
|
||||
let data = data.as_ref();
|
||||
validate_element_count(&shape, data.len())?;
|
||||
Ok(Self {
|
||||
dtype: "int64".to_string(),
|
||||
shape,
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(&data)),
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
|
||||
pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
|
||||
pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
|
||||
let data = data.as_ref();
|
||||
validate_element_count(&shape, data.len())?;
|
||||
Ok(Self {
|
||||
dtype: "uint32".to_string(),
|
||||
shape,
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(&data)),
|
||||
data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use vllm_parser::tool::{
|
||||
Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput,
|
||||
Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput,
|
||||
};
|
||||
use vllm_parser::unified::{
|
||||
UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
|
||||
@@ -85,8 +85,8 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
|
||||
self.inner.preserve_special_tokens()
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
self.inner.structural_tag_model()
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
self.inner.structural_tag_builder()
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{DeepSeekDsmlToolParser, DsmlTokens};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3.2 models.
|
||||
///
|
||||
@@ -47,8 +47,8 @@ impl ToolParser for DeepSeekV32ToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::DeepSeekV32)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::DeepSeekV32.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{DeepSeekDsmlToolParser, DsmlTokens};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V4 models.
|
||||
///
|
||||
@@ -50,8 +50,8 @@ impl ToolParser for DeepSeekV4ToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::DeepSeekV4)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::DeepSeekV4.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
@@ -73,7 +73,7 @@ mod tests {
|
||||
|
||||
use super::DeepSeekV4ToolParser;
|
||||
use crate::tool::test_utils::{collect_stream, test_tools};
|
||||
use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _};
|
||||
use crate::tool::{ToolParser, ToolParserTestExt as _};
|
||||
|
||||
fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
|
||||
let params = params
|
||||
@@ -91,13 +91,10 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deepseek_v4_exposes_structural_tag_model() {
|
||||
fn deepseek_v4_exposes_structural_tag_builder() {
|
||||
let parser = DeepSeekV4ToolParser::new(&test_tools());
|
||||
|
||||
assert_eq!(
|
||||
parser.structural_tag_model(),
|
||||
Some(StructuralTagModel::DeepSeekV4)
|
||||
);
|
||||
assert!(parser.structural_tag_builder().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3 JSON-fenced tool calls.
|
||||
///
|
||||
@@ -35,8 +35,8 @@ impl ToolParser for DeepSeekV3ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::DeepSeekR1)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::DeepSeekR1.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for DeepSeek V3.1 raw JSON tool calls.
|
||||
///
|
||||
@@ -31,8 +31,8 @@ impl ToolParser for DeepSeekV31ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::DeepSeekV31)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::DeepSeekV31.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{GlmXmlToolParser, Separator};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Tool parser for GLM-4.7 MoE XML-style tool calls.
|
||||
///
|
||||
@@ -25,8 +25,8 @@ impl ToolParser for Glm47MoeToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Glm47)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Glm47.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
use crate::tool::{StructuralTagBuilder, Tool};
|
||||
|
||||
const TOOL_CALLS_START: &str = "<tool_calls>";
|
||||
const TOOL_CALLS_END: &str = "</tool_calls>";
|
||||
@@ -116,8 +116,8 @@ impl ToolParser for HyV3ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::HyV3)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::HyV3.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Hermes",
|
||||
@@ -48,8 +48,8 @@ impl ToolParser for HermesToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Hermes)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Hermes.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -12,7 +12,9 @@ use super::{
|
||||
argument_delta_event, tool_call_header_event,
|
||||
};
|
||||
use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{
|
||||
Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum LlamaJsonMode {
|
||||
@@ -136,8 +138,8 @@ impl ToolParser for Llama3JsonToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Llama)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Llama.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
|
||||
use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
|
||||
parser_name: "Qwen XML",
|
||||
@@ -50,8 +50,8 @@ impl ToolParser for Qwen3XmlToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Qwen3)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Qwen3.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -11,7 +11,7 @@ use winnow::token::{literal, rest, take_until, take_while};
|
||||
|
||||
use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
use crate::tool::{StructuralTagBuilder, Tool};
|
||||
|
||||
const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
|
||||
const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
|
||||
@@ -150,8 +150,8 @@ impl ToolParser for KimiK2ToolParser {
|
||||
true
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Kimi)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Kimi.builder())
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
|
||||
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagModel, Tool};
|
||||
use crate::tool::{StructuralTagBuilder, Tool};
|
||||
|
||||
const TOOL_CALL_START: &str = "<minimax:tool_call>";
|
||||
const TOOL_CALL_END: &str = "</minimax:tool_call>";
|
||||
@@ -115,8 +115,8 @@ impl ToolParser for MinimaxM2ToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Minimax)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Minimax.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
|
||||
@@ -36,7 +36,7 @@ pub use qwen_coder::Qwen3CoderToolParser;
|
||||
pub use seed_oss::SeedOssToolParser;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
pub use xgrammar_structural_tag::Model as StructuralTagModel;
|
||||
pub use xgrammar_structural_tag::builders::StructuralTagBuilder;
|
||||
|
||||
use crate::utils;
|
||||
|
||||
@@ -187,8 +187,8 @@ pub trait ToolParser: Send {
|
||||
false
|
||||
}
|
||||
|
||||
/// Return the xgrammar structural-tag model used for strict tool calling.
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
/// Return the xgrammar structural-tag builder used for strict tool calling.
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
None
|
||||
}
|
||||
|
||||
|
||||
@@ -9,8 +9,8 @@ use winnow::token::{literal, take_until};
|
||||
|
||||
use super::parameters::ToolSchemas;
|
||||
use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
|
||||
use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::Tool;
|
||||
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagBuilder, Tool};
|
||||
|
||||
const TOOL_CALL_START: &str = "<tool_call>";
|
||||
const TOOL_CALL_END: &str = "</tool_call>";
|
||||
@@ -146,8 +146,8 @@ impl ToolParser for Qwen3CoderToolParser {
|
||||
Ok(Box::new(Self::new(tools)))
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
Some(StructuralTagModel::Qwen3Coder)
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
Some(xgrammar_structural_tag::Model::Qwen3Coder.builder())
|
||||
}
|
||||
|
||||
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
|
||||
@@ -294,7 +294,7 @@ mod tests {
|
||||
use serde_json::{Value, json};
|
||||
use thiserror_ext::AsReport;
|
||||
|
||||
use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser};
|
||||
use super::{Qwen3CoderToolParser, ToolParser};
|
||||
use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
|
||||
use crate::tool::{ToolParserOutput, ToolParserTestExt as _};
|
||||
|
||||
@@ -308,13 +308,10 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qwen_coder_exposes_structural_tag_model() {
|
||||
fn qwen_coder_exposes_structural_tag_builder() {
|
||||
let parser = Qwen3CoderToolParser::new(&test_tools());
|
||||
|
||||
assert_eq!(
|
||||
parser.structural_tag_model(),
|
||||
Some(StructuralTagModel::Qwen3Coder)
|
||||
);
|
||||
assert!(parser.structural_tag_builder().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -7,7 +7,7 @@ use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
|
||||
use crate::reasoning::ReasoningParser;
|
||||
use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput};
|
||||
use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};
|
||||
|
||||
/// Unified parser that composes existing reasoning and tool parsers.
|
||||
pub struct CombinedParser {
|
||||
@@ -79,8 +79,8 @@ impl UnifiedParser for CombinedParser {
|
||||
|| self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens())
|
||||
}
|
||||
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
self.tool.as_ref().and_then(|parser| parser.structural_tag_model())
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
self.tool.as_ref().and_then(|parser| parser.structural_tag_builder())
|
||||
}
|
||||
|
||||
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
|
||||
@@ -269,10 +269,7 @@ mod tests {
|
||||
fn combined_parser_emits_tool_calls_from_visible_content() {
|
||||
let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap();
|
||||
let mut parser = CombinedParser::new(None, Some(tool));
|
||||
assert!(matches!(
|
||||
parser.structural_tag_model(),
|
||||
Some(crate::tool::StructuralTagModel::Qwen3)
|
||||
));
|
||||
assert!(parser.structural_tag_builder().is_some());
|
||||
|
||||
let output = collect(
|
||||
&mut parser,
|
||||
|
||||
@@ -16,7 +16,7 @@ use vllm_tokenizer::DynTokenizer;
|
||||
|
||||
use crate::reasoning::ReasoningError;
|
||||
use crate::tool::{
|
||||
StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
|
||||
StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
|
||||
};
|
||||
|
||||
/// Result alias for unified parser operations.
|
||||
@@ -171,8 +171,8 @@ pub trait UnifiedParser: Send {
|
||||
false
|
||||
}
|
||||
|
||||
/// Return the xgrammar structural-tag model used for strict tool calling.
|
||||
fn structural_tag_model(&self) -> Option<StructuralTagModel> {
|
||||
/// Return the xgrammar structural-tag builder used for strict tool calling.
|
||||
fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
|
||||
None
|
||||
}
|
||||
|
||||
|
||||
@@ -238,13 +238,18 @@ fn collect_generate(
|
||||
None
|
||||
};
|
||||
let prompt_logprobs = if include_prompt_logprobs {
|
||||
let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
|
||||
ApiError::server_error(
|
||||
"raw generate response requested prompt_logprobs but generation returned none"
|
||||
.to_string(),
|
||||
)
|
||||
})?;
|
||||
Some(raw_prompt_logprobs_to_maps(prompt_logprobs))
|
||||
match collected.prompt_logprobs.as_ref() {
|
||||
Some(prompt_logprobs) => Some(raw_prompt_logprobs_to_maps(prompt_logprobs)),
|
||||
// A single-token prompt has no scored positions; same mapping
|
||||
// as /v1/completions.
|
||||
None if collected.prompt_token_ids.len() == 1 => Some(vec![None]),
|
||||
None => {
|
||||
return Err(ApiError::server_error(
|
||||
"raw generate response requested prompt_logprobs but generation returned none"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -472,4 +477,48 @@ mod tests {
|
||||
Some(2)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
|
||||
let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
|
||||
request_id: "raw-1".to_string(),
|
||||
prompt_logprobs: None,
|
||||
token_ids: vec![3],
|
||||
logprobs: None,
|
||||
finish_reason: FinishReason::stop_eos(),
|
||||
usage: vllm_llm::TokenUsage {
|
||||
prompt_token_count: prompt_token_ids.len(),
|
||||
output_token_count: 1,
|
||||
cached_token_count: 0,
|
||||
},
|
||||
kv_transfer_params: None,
|
||||
ec_transfer_params: None,
|
||||
prompt_token_ids,
|
||||
};
|
||||
|
||||
let response = collect_generate(
|
||||
output_without_payload(vec![9707]),
|
||||
"raw-1".to_string(),
|
||||
ApiServerOptions::default(),
|
||||
ResponseOptions {
|
||||
include_prompt_logprobs: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.expect("single-token prompt without payload maps to [None]");
|
||||
let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
|
||||
assert_eq!(prompt_logprobs.len(), 1);
|
||||
assert!(prompt_logprobs[0].is_none());
|
||||
|
||||
collect_generate(
|
||||
output_without_payload(vec![9707, 11]),
|
||||
"raw-2".to_string(),
|
||||
ApiServerOptions::default(),
|
||||
ResponseOptions {
|
||||
include_prompt_logprobs: true,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.expect_err("multi-token prompt without payload is an engine failure");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -35,7 +35,7 @@ use crate::routes::openai::chat_completions::types::{
|
||||
ChatMessageDelta,
|
||||
};
|
||||
use crate::routes::openai::utils::logprobs::{
|
||||
decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps,
|
||||
decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
|
||||
};
|
||||
use crate::routes::openai::utils::types::{
|
||||
ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
|
||||
@@ -181,14 +181,11 @@ async fn collect_chat_completion(
|
||||
None
|
||||
};
|
||||
let prompt_logprobs = if include_prompt_logprobs {
|
||||
Some(decoded_prompt_logprobs_to_maps(
|
||||
prompt_logprobs.as_ref().ok_or_else(|| {
|
||||
server_error!(
|
||||
"chat response requested prompt_logprobs but generation returned none"
|
||||
)
|
||||
})?,
|
||||
Some(prompt_logprobs_to_maps(
|
||||
prompt_logprobs.as_ref(),
|
||||
&prompt_token_ids,
|
||||
return_tokens_as_token_ids,
|
||||
))
|
||||
)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
@@ -5,7 +5,6 @@ mod convert;
|
||||
mod types;
|
||||
mod validate;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::convert::Infallible;
|
||||
use std::result::Result;
|
||||
use std::sync::Arc;
|
||||
@@ -29,8 +28,8 @@ use vllm_text::{
|
||||
|
||||
use self::convert::{ResponseOptions, prepare_completion_request};
|
||||
use super::utils::logprobs::{
|
||||
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
|
||||
decoded_prompt_logprobs_to_openai, text_len,
|
||||
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai,
|
||||
prompt_logprobs_to_maps, text_len,
|
||||
};
|
||||
use super::utils::types::Usage;
|
||||
use crate::config::ApiServerOptions;
|
||||
@@ -505,27 +504,6 @@ fn prompt_only_logprobs_to_openai(
|
||||
))
|
||||
}
|
||||
|
||||
fn prompt_logprobs_to_maps(
|
||||
prompt_logprobs: Option<&DecodedPromptLogprobs>,
|
||||
prompt_token_ids: &[u32],
|
||||
return_tokens_as_token_ids: bool,
|
||||
) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
|
||||
if let Some(prompt_logprobs) = prompt_logprobs {
|
||||
return Ok(decoded_prompt_logprobs_to_maps(
|
||||
prompt_logprobs,
|
||||
return_tokens_as_token_ids,
|
||||
));
|
||||
}
|
||||
|
||||
if let [_token_id] = prompt_token_ids {
|
||||
return Ok(vec![None]);
|
||||
}
|
||||
|
||||
Err(server_error!(
|
||||
"completion response requested prompt_logprobs but generation returned none"
|
||||
))
|
||||
}
|
||||
|
||||
fn usage_chunk(
|
||||
request_id: &str,
|
||||
response_model: &str,
|
||||
|
||||
@@ -100,20 +100,31 @@ pub fn decoded_prompt_logprobs_to_openai(
|
||||
})
|
||||
}
|
||||
|
||||
/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response
|
||||
/// shape.
|
||||
pub fn decoded_prompt_logprobs_to_maps(
|
||||
prompt_logprobs: &DecodedPromptLogprobs,
|
||||
/// Map decoded prompt logprobs into vLLM-style per-position maps, treating a
|
||||
/// missing single-token payload as `[None]`.
|
||||
pub fn prompt_logprobs_to_maps(
|
||||
prompt_logprobs: Option<&DecodedPromptLogprobs>,
|
||||
prompt_token_ids: &[u32],
|
||||
return_tokens_as_token_ids: bool,
|
||||
) -> Vec<Option<HashMap<String, f32>>> {
|
||||
std::iter::once(None)
|
||||
.chain(prompt_logprobs.scored_positions.iter().map(|position| {
|
||||
Some(position_top_logprobs_map(
|
||||
position,
|
||||
return_tokens_as_token_ids,
|
||||
))
|
||||
}))
|
||||
.collect()
|
||||
) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
|
||||
if let Some(prompt_logprobs) = prompt_logprobs {
|
||||
return Ok(std::iter::once(None)
|
||||
.chain(prompt_logprobs.scored_positions.iter().map(|position| {
|
||||
Some(position_top_logprobs_map(
|
||||
position,
|
||||
return_tokens_as_token_ids,
|
||||
))
|
||||
}))
|
||||
.collect());
|
||||
}
|
||||
|
||||
if let [_token_id] = prompt_token_ids {
|
||||
return Ok(vec![None]);
|
||||
}
|
||||
|
||||
Err(server_error!(
|
||||
"prompt_logprobs were requested but generation returned none"
|
||||
))
|
||||
}
|
||||
|
||||
/// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
|
||||
@@ -275,7 +286,13 @@ pub fn clamp_logprob(logprob: f32) -> f32 {
|
||||
mod tests {
|
||||
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
|
||||
|
||||
use super::decoded_logprobs_to_openai_chat;
|
||||
use super::{decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps};
|
||||
|
||||
#[test]
|
||||
fn prompt_logprobs_maps_reject_missing_multi_token_payload() {
|
||||
prompt_logprobs_to_maps(None, &[9707, 11], false)
|
||||
.expect_err("multi-token prompt without payload is an engine failure");
|
||||
}
|
||||
|
||||
fn sample_logprobs() -> DecodedLogprobs {
|
||||
DecodedLogprobs {
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cuda_platform_base(monkeypatch: pytest.MonkeyPatch) -> Any:
|
||||
stable_libtorch_module = ModuleType("vllm._C_stable_libtorch")
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"vllm._C_stable_libtorch",
|
||||
stable_libtorch_module,
|
||||
)
|
||||
from vllm.platforms.cuda import CudaPlatformBase
|
||||
|
||||
return CudaPlatformBase
|
||||
|
||||
|
||||
def test_compiled_arch_covers_device(cuda_platform_base: Any) -> None:
|
||||
assert cuda_platform_base._compiled_arch_covers_device(
|
||||
"12.1a", DeviceCapability(12, 1)
|
||||
)
|
||||
assert not cuda_platform_base._compiled_arch_covers_device(
|
||||
"12.0a", DeviceCapability(12, 1)
|
||||
)
|
||||
|
||||
assert cuda_platform_base._compiled_arch_covers_device(
|
||||
"12.0f", DeviceCapability(12, 1)
|
||||
)
|
||||
assert not cuda_platform_base._compiled_arch_covers_device(
|
||||
"10.0f", DeviceCapability(12, 1)
|
||||
)
|
||||
|
||||
|
||||
def test_warn_if_device_arch_not_compiled(
|
||||
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
|
||||
) -> None:
|
||||
def device_count(cls: type[Any]) -> int:
|
||||
return 2
|
||||
|
||||
def get_device_capability(
|
||||
cls: type[Any], device_id: int = 0
|
||||
) -> DeviceCapability | None:
|
||||
capabilities = {
|
||||
0: DeviceCapability(12, 1),
|
||||
1: DeviceCapability(10, 3),
|
||||
}
|
||||
return capabilities[device_id]
|
||||
|
||||
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
|
||||
return f"GPU {device_id}"
|
||||
|
||||
warnings: list[tuple[str, str, str]] = []
|
||||
|
||||
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
|
||||
warnings.append((message, compiled_archs, devices))
|
||||
|
||||
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_base,
|
||||
"get_device_capability",
|
||||
classmethod(get_device_capability),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_base, "get_device_name", classmethod(get_device_name)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.platforms.cuda.torch.ops",
|
||||
SimpleNamespace(
|
||||
_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "12.0f,10.0a")
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.platforms.cuda.logger",
|
||||
SimpleNamespace(warning_once=warning_once),
|
||||
)
|
||||
|
||||
cuda_platform_base._warn_if_device_arch_not_compiled()
|
||||
|
||||
assert len(warnings) == 1
|
||||
assert warnings[0][1] == "12.0f, 10.0a"
|
||||
assert "1: GPU 1 (compute capability 10.3)" in warnings[0][2]
|
||||
|
||||
|
||||
def test_warn_if_device_arch_not_compiled_no_warning(
|
||||
monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
|
||||
) -> None:
|
||||
def device_count(cls: type[Any]) -> int:
|
||||
return 1
|
||||
|
||||
def get_device_capability(
|
||||
cls: type[Any], device_id: int = 0
|
||||
) -> DeviceCapability | None:
|
||||
return DeviceCapability(10, 3)
|
||||
|
||||
def get_device_name(cls: type[Any], device_id: int = 0) -> str:
|
||||
raise AssertionError("get_device_name should not be called for covered devices")
|
||||
|
||||
warnings: list[tuple[str, str, str]] = []
|
||||
|
||||
def warning_once(message: str, compiled_archs: str, devices: str) -> None:
|
||||
warnings.append((message, compiled_archs, devices))
|
||||
|
||||
monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_base,
|
||||
"get_device_capability",
|
||||
classmethod(get_device_capability),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_base, "get_device_name", classmethod(get_device_name)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.platforms.cuda.torch.ops",
|
||||
SimpleNamespace(_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "10.0f")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm.platforms.cuda.logger",
|
||||
SimpleNamespace(warning_once=warning_once),
|
||||
)
|
||||
|
||||
cuda_platform_base._warn_if_device_arch_not_compiled()
|
||||
|
||||
assert warnings == []
|
||||
@@ -60,6 +60,8 @@ def _make_quick_allreduce(
|
||||
qar.use_fp16_kernels = use_fp16_kernels
|
||||
qar.qr_quant_level = QuickReduceRegime[quant_level]
|
||||
qar.qr_max_size = qr_max_size
|
||||
qar.qr_min_size = None
|
||||
qar.qr_quantization_min_size = None
|
||||
return qar
|
||||
|
||||
|
||||
@@ -511,13 +513,21 @@ def test_quick_reduce_regime_values():
|
||||
assert QuickReduceRegime.INT8.value == 1
|
||||
assert QuickReduceRegime.INT6.value == 2
|
||||
assert QuickReduceRegime.INT4.value == 3
|
||||
assert QuickReduceRegime.NONE.value == 4
|
||||
assert QuickReduceRegime.INT3.value == 4
|
||||
assert QuickReduceRegime.NONE.value == 5
|
||||
|
||||
|
||||
def test_quick_reduce_regime_names():
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
|
||||
|
||||
assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
|
||||
assert set(QuickReduceRegime.__members__) == {
|
||||
"FP",
|
||||
"INT8",
|
||||
"INT6",
|
||||
"INT4",
|
||||
"INT3",
|
||||
"NONE",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
|
||||
@@ -693,7 +703,7 @@ def test_quick_allreduce_min_size_table():
|
||||
for dtype in [torch.float16, torch.bfloat16]:
|
||||
for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
|
||||
min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
|
||||
assert len(min_sizes) == 4
|
||||
assert len(min_sizes) == 5
|
||||
assert all(size > 0 for size in min_sizes)
|
||||
|
||||
|
||||
|
||||
@@ -515,3 +515,15 @@ def test_structured_outputs_structural_tag_invalid(structural_tag):
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
structured_outputs={"structural_tag": structural_tag},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"])
|
||||
def test_non_numeric_logprobs_rejected(field_name):
|
||||
"""A non-numeric logprobs value must be a clean 400 validation error, not a
|
||||
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
|
||||
with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
|
||||
ChatCompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
**{field_name: "2"},
|
||||
)
|
||||
|
||||
@@ -610,3 +610,16 @@ class TestCompletionPromptListLimit:
|
||||
max_tokens=1,
|
||||
)
|
||||
assert len(request.prompt_embeds) == 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"])
|
||||
def test_non_numeric_logprobs_rejected(field_name):
|
||||
"""A non-numeric logprobs value must be a clean 400 validation error, not a
|
||||
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
|
||||
with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
|
||||
CompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
prompt="Test prompt",
|
||||
max_tokens=10,
|
||||
**{field_name: "2"},
|
||||
)
|
||||
|
||||
@@ -6,7 +6,6 @@ max_concurrency: 100
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--max-num-batched-tokens 32768
|
||||
--safetensors-load-strategy prefetch
|
||||
--moe-backend flashinfer_cutlass
|
||||
--prefill-context-parallel-size 4
|
||||
|
||||
@@ -6,7 +6,6 @@ max_concurrency: 100
|
||||
server_args: >-
|
||||
--enforce-eager
|
||||
--max-model-len 4096
|
||||
--max-num-batched-tokens 32768
|
||||
--safetensors-load-strategy prefetch
|
||||
--moe-backend flashinfer_cutlass
|
||||
--tensor-parallel-size 2
|
||||
|
||||
@@ -11,9 +11,6 @@ from vllm._custom_ops import (
|
||||
scaled_fp8_quant,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.ops.triton_merge_attn_states import (
|
||||
mask_empty_context,
|
||||
)
|
||||
from vllm.v1.attention.ops.triton_merge_attn_states import (
|
||||
merge_attn_states as merge_attn_states_triton,
|
||||
)
|
||||
@@ -76,59 +73,6 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]
|
||||
all_case_info: list[tuple] = []
|
||||
|
||||
|
||||
def test_mask_empty_context() -> None:
|
||||
query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
|
||||
query_start_loc = torch.cat(
|
||||
(torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
|
||||
).cuda()
|
||||
context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
|
||||
context_start_loc = torch.cat(
|
||||
(torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
|
||||
).cuda()
|
||||
num_heads, num_tokens, head_dim = 4, 165, 16
|
||||
lse = torch.randn(num_heads, num_tokens, device="cuda")
|
||||
output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
|
||||
# Empty-context rows carry undefined (possibly non-finite) attention output.
|
||||
output[33:164] = float("nan")
|
||||
|
||||
expected_lse = lse.clone()
|
||||
expected_lse[:, 33:164] = float("-inf")
|
||||
expected_output = output.clone()
|
||||
expected_output[33:164] = 0.0
|
||||
|
||||
mask_empty_context(lse, output, query_start_loc, context_start_loc)
|
||||
|
||||
torch.testing.assert_close(lse, expected_lse)
|
||||
torch.testing.assert_close(output, expected_output)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
|
||||
@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
|
||||
def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
|
||||
"""When a token is empty on both sides (both LSE -inf), the 0/0 softmax
|
||||
scales must not surface as NaN in the merged output."""
|
||||
num_tokens, num_heads, head_size = 6, 8, 128
|
||||
prefix_output = torch.zeros(
|
||||
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
|
||||
)
|
||||
prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")
|
||||
suffix_output = torch.zeros(
|
||||
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
|
||||
)
|
||||
suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
|
||||
|
||||
# Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
|
||||
# their outputs and set both LSEs to -inf).
|
||||
empty = slice(2, 4)
|
||||
prefix_lse[:, empty] = float("-inf")
|
||||
suffix_lse[:, empty] = float("-inf")
|
||||
|
||||
output = torch.empty_like(prefix_output)
|
||||
merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
|
||||
|
||||
assert not output.isnan().any()
|
||||
|
||||
|
||||
def generate_markdown_table():
|
||||
global all_case_info
|
||||
table_header = (
|
||||
|
||||
@@ -425,7 +425,7 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No
|
||||
match the single-call result.
|
||||
"""
|
||||
from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
|
||||
causal_conv1d_torch,
|
||||
causal_conv1d_fn_cpu as causal_conv1d_torch,
|
||||
)
|
||||
|
||||
x, weight, bias = _conv_inputs(total_tokens)
|
||||
|
||||
@@ -18,8 +18,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
DEVICE = current_platform.device_type
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
|
||||
reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
|
||||
not (
|
||||
current_platform.is_cuda_alike()
|
||||
or current_platform.is_xpu()
|
||||
or current_platform.is_cpu()
|
||||
),
|
||||
reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU",
|
||||
)
|
||||
|
||||
|
||||
@@ -284,7 +288,8 @@ def test_causal_conv1d_varlen(
|
||||
batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
|
||||
):
|
||||
device = DEVICE
|
||||
torch.accelerator.empty_cache()
|
||||
if not current_platform.is_cpu():
|
||||
torch.accelerator.empty_cache()
|
||||
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
|
||||
if itype == torch.bfloat16:
|
||||
rtol, atol = 1e-2, 5e-2
|
||||
|
||||
@@ -20,8 +20,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
DEVICE = current_platform.device_type
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
|
||||
reason="mamba_ssm kernels require CUDA-alike or XPU",
|
||||
not (
|
||||
current_platform.is_cuda_alike()
|
||||
or current_platform.is_xpu()
|
||||
or current_platform.is_cpu()
|
||||
),
|
||||
reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU",
|
||||
)
|
||||
|
||||
# selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op,
|
||||
@@ -342,6 +346,13 @@ def test_selective_scan(
|
||||
@pytest.mark.parametrize("has_z", [False, True])
|
||||
@pytest.mark.parametrize("dstate", [16, 64])
|
||||
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cpu(),
|
||||
reason=(
|
||||
"CPU kernel for selective_state_update only supports "
|
||||
"Mamba 2 (scalar A/dt), not Mamba 1."
|
||||
),
|
||||
)
|
||||
def test_selective_state_update(dim, dstate, has_z, itype):
|
||||
device = DEVICE
|
||||
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
|
||||
@@ -436,6 +447,13 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r
|
||||
@pytest.mark.parametrize("dstate", [16, 64])
|
||||
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
|
||||
@pytest.mark.parametrize("max_seq_len", [1, 2, 4])
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cpu(),
|
||||
reason=(
|
||||
"CPU kernel for selective_state_update only supports "
|
||||
"Mamba 2 (scalar A/dt), not Mamba 1."
|
||||
),
|
||||
)
|
||||
def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
|
||||
device = DEVICE
|
||||
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
|
||||
@@ -697,6 +715,13 @@ def test_selective_scan_varlen(
|
||||
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
|
||||
# tests correctness in case subset of the sequences are padded
|
||||
@pytest.mark.parametrize("with_padding", [True, False])
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cpu(),
|
||||
reason=(
|
||||
"CPU kernel for selective_state_update only supports "
|
||||
"Mamba 2 (scalar A/dt), not Mamba 1."
|
||||
),
|
||||
)
|
||||
def test_selective_state_update_with_batch_indices(
|
||||
with_padding, dim, dstate, has_z, itype
|
||||
):
|
||||
@@ -789,6 +814,13 @@ def test_selective_state_update_with_batch_indices(
|
||||
@pytest.mark.parametrize("ngroups", [1, 4])
|
||||
@pytest.mark.parametrize("dstate", [16, 64])
|
||||
@pytest.mark.parametrize("dim", [2048, 4096])
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cpu(),
|
||||
reason=(
|
||||
"CPU kernel for selective_state_update only supports "
|
||||
"Mamba 2 (scalar A/dt), not Mamba 1."
|
||||
),
|
||||
)
|
||||
def test_selective_state_update_with_heads_with_batch_indices(
|
||||
dim, dstate, ngroups, has_z, tie_hdim, itype
|
||||
):
|
||||
@@ -862,6 +894,13 @@ def test_selective_state_update_with_heads_with_batch_indices(
|
||||
@pytest.mark.parametrize("dstate", [16, 64])
|
||||
@pytest.mark.parametrize("dim", [2048, 4096])
|
||||
@pytest.mark.parametrize("max_seq_len", [2, 4])
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cpu(),
|
||||
reason=(
|
||||
"CPU kernel for selective_state_update only supports "
|
||||
"Mamba 2 (scalar A/dt), not Mamba 1."
|
||||
),
|
||||
)
|
||||
def test_selective_state_update_with_num_accepted_tokens(
|
||||
dim, dstate, has_z, itype, max_seq_len
|
||||
):
|
||||
@@ -988,6 +1027,13 @@ def test_selective_state_update_with_num_accepted_tokens(
|
||||
@pytest.mark.parametrize("dstate", [16, 64])
|
||||
@pytest.mark.parametrize("dim", [2048, 4096])
|
||||
@pytest.mark.parametrize("max_seq_len", [2, 4])
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cpu(),
|
||||
reason=(
|
||||
"CPU kernel for selective_state_update only supports "
|
||||
"Mamba 2 (scalar A/dt), not Mamba 1."
|
||||
),
|
||||
)
|
||||
def test_selective_state_update_varlen_with_num_accepted(
|
||||
dim, dstate, has_z, itype, max_seq_len
|
||||
):
|
||||
|
||||
@@ -0,0 +1,880 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end tests for routed-expert capture on the monolithic MoE path.
|
||||
|
||||
These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the
|
||||
expert IDs picked by FlashInfer's fused router-and-experts kernels (the
|
||||
"monolithic" path). When ``set_capture_fn`` is installed on
|
||||
a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call
|
||||
should:
|
||||
|
||||
* allocate an int16 ``(num_tokens, top_k)`` buffer,
|
||||
* pass it to FlashInfer as ``routing_replay_out``,
|
||||
* invoke the callback after the kernel returns.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEParallelConfig,
|
||||
FusedMoEQuantConfig,
|
||||
RoutingMethodType,
|
||||
fp8_w8a8_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
|
||||
TrtLlmBf16ExpertsMonolithic,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
|
||||
TrtLlmFp8ExpertsMonolithic,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
|
||||
TrtLlmNvFp4ExpertsMonolithic,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
try:
|
||||
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
|
||||
except ImportError:
|
||||
pytest.skip("flashinfer not available", allow_module_level=True)
|
||||
|
||||
if not has_flashinfer_trtllm_fused_moe() or not current_platform.is_cuda():
|
||||
pytest.skip(
|
||||
"Requires FlashInfer TRT-LLM fused MoE on CUDA",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip(
|
||||
"TRT-LLM fused MoE kernels require SM100+",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def _shuffle_bf16_weights_block_major_k(
|
||||
w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128
|
||||
) -> torch.Tensor:
|
||||
"""Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by
|
||||
``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert
|
||||
row shuffle.
|
||||
"""
|
||||
from flashinfer import shuffle_matrix_a
|
||||
from flashinfer.fused_moe import convert_to_block_layout
|
||||
|
||||
num_experts = w.shape[0]
|
||||
shuffled = []
|
||||
for i in range(num_experts):
|
||||
t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m)
|
||||
shuffled.append(convert_to_block_layout(t, block_k))
|
||||
return torch.stack(shuffled).view(torch.bfloat16)
|
||||
|
||||
|
||||
def _make_bf16_monolithic_experts(
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
routing_method: RoutingMethodType,
|
||||
device: torch.device,
|
||||
) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]:
|
||||
"""Construct the monolithic BF16 experts plus the BlockMajorK weights
|
||||
expected by ``trtllm_bf16_moe``.
|
||||
"""
|
||||
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
|
||||
moe_config = FusedMoEConfig(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=top_k,
|
||||
hidden_dim=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
moe_parallel_config=parallel_cfg,
|
||||
in_dtype=torch.bfloat16,
|
||||
activation=MoEActivation.SILU,
|
||||
device=device,
|
||||
routing_method=routing_method,
|
||||
max_num_tokens=max(8, 1),
|
||||
)
|
||||
quant_config = FusedMoEQuantConfig.make(
|
||||
quant_dtype=None,
|
||||
per_act_token_quant=False,
|
||||
per_out_ch_quant=False,
|
||||
block_shape=None,
|
||||
)
|
||||
|
||||
experts = TrtLlmBf16ExpertsMonolithic(
|
||||
moe_config=moe_config, quant_config=quant_config
|
||||
)
|
||||
|
||||
gemm1 = (
|
||||
torch.randn(
|
||||
num_experts,
|
||||
2 * intermediate_size,
|
||||
hidden_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
* 0.1
|
||||
)
|
||||
gemm2 = (
|
||||
torch.randn(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
* 0.1
|
||||
)
|
||||
w13 = _shuffle_bf16_weights_block_major_k(gemm1)
|
||||
w2 = _shuffle_bf16_weights_block_major_k(gemm2)
|
||||
return experts, w13, w2
|
||||
|
||||
|
||||
def _run_bf16_monolithic(
|
||||
experts: TrtLlmBf16ExpertsMonolithic,
|
||||
hidden_states: torch.Tensor,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
num_experts: int,
|
||||
*,
|
||||
n_group: int | None = None,
|
||||
topk_group: int | None = None,
|
||||
routed_scaling_factor: float | None = None,
|
||||
e_score_correction_bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
return experts.apply(
|
||||
hidden_states=hidden_states,
|
||||
w1=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
activation=MoEActivation.SILU,
|
||||
global_num_experts=num_experts,
|
||||
expert_map=None,
|
||||
a1q_scale=None,
|
||||
apply_router_weight_on_input=False,
|
||||
num_expert_group=n_group,
|
||||
topk_group=topk_group,
|
||||
e_score_correction_bias=e_score_correction_bias,
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
)
|
||||
|
||||
|
||||
_DSV3_NUM_EXPERTS = 32
|
||||
_DSV3_N_GROUP = 4
|
||||
_DSV3_TOPK_GROUP = 2
|
||||
|
||||
|
||||
def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor:
|
||||
return torch.randn(num_experts, device=device, dtype=torch.bfloat16)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
|
||||
@pytest.mark.parametrize("top_k", [2, 4])
|
||||
def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts(
|
||||
num_tokens: int,
|
||||
top_k: int,
|
||||
) -> None:
|
||||
"""The capture callback should receive the int16 routed-expert IDs the
|
||||
kernel actually used, the values should be valid expert indices, and
|
||||
each token should pick ``top_k`` distinct experts."""
|
||||
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
|
||||
pytest.skip(
|
||||
f"DSV3 requires top_k <= n_group * topk_group "
|
||||
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
|
||||
)
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
num_experts = _DSV3_NUM_EXPERTS
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
|
||||
experts, w13, w2 = _make_bf16_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
device=device,
|
||||
)
|
||||
|
||||
captured: list[torch.Tensor] = []
|
||||
|
||||
def capture_fn(replay_out: torch.Tensor) -> None:
|
||||
captured.append(replay_out.clone())
|
||||
|
||||
assert experts.supports_routing_replay_capture()
|
||||
experts.set_capture_fn(capture_fn)
|
||||
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
routing_bias = _make_dsv3_routing_bias(num_experts, device)
|
||||
|
||||
_ = _run_bf16_monolithic(
|
||||
experts,
|
||||
hidden_states=hidden_states,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
n_group=_DSV3_N_GROUP,
|
||||
topk_group=_DSV3_TOPK_GROUP,
|
||||
routed_scaling_factor=1.0,
|
||||
e_score_correction_bias=routing_bias,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
replay = captured[0]
|
||||
assert replay.dtype == torch.int16
|
||||
assert replay.shape == (num_tokens, top_k)
|
||||
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
|
||||
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
|
||||
|
||||
for t in range(num_tokens):
|
||||
unique = replay[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts, "
|
||||
f"got {unique.numel()} ({replay[t].tolist()})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
|
||||
@pytest.mark.parametrize(
|
||||
"routing_method",
|
||||
[
|
||||
RoutingMethodType.Renormalize,
|
||||
RoutingMethodType.RenormalizeNaive,
|
||||
],
|
||||
)
|
||||
def test_trtllm_bf16_monolithic_routing_replay_non_dsv3(
|
||||
num_tokens: int,
|
||||
routing_method: RoutingMethodType,
|
||||
) -> None:
|
||||
"""Routing replay works for non-DeepSeekV3 routing methods too.
|
||||
FlashInfer's ``routing_replay_out`` is routing-method-agnostic."""
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
num_experts = 8
|
||||
top_k = 2
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
|
||||
experts, w13, w2 = _make_bf16_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
routing_method=routing_method,
|
||||
device=device,
|
||||
)
|
||||
|
||||
captured: list[torch.Tensor] = []
|
||||
experts.set_capture_fn(lambda r: captured.append(r.clone()))
|
||||
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
|
||||
_ = _run_bf16_monolithic(
|
||||
experts,
|
||||
hidden_states=hidden_states,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
replay = captured[0]
|
||||
assert replay.dtype == torch.int16
|
||||
assert replay.shape == (num_tokens, top_k)
|
||||
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
|
||||
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
|
||||
for t in range(num_tokens):
|
||||
unique = replay[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts, "
|
||||
f"got {unique.numel()} ({replay[t].tolist()})"
|
||||
)
|
||||
|
||||
|
||||
def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None:
|
||||
"""With no callback installed the kernel should not see a
|
||||
``routing_replay_out`` tensor — verify the helper short-circuits."""
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
experts, _, _ = _make_bf16_monolithic_experts(
|
||||
num_experts=_DSV3_NUM_EXPERTS,
|
||||
top_k=2,
|
||||
hidden_size=1024,
|
||||
intermediate_size=1024,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
device=device,
|
||||
)
|
||||
# No callback installed.
|
||||
buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device)
|
||||
assert buf is None
|
||||
|
||||
# Dispatch is also a no-op.
|
||||
experts._maybe_dispatch_routing_replay(buf, num_tokens=4)
|
||||
|
||||
|
||||
def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None:
|
||||
"""FlashInfer's ``routing_replay_out`` is supported by all routing
|
||||
methods, so ``supports_routing_replay_capture`` should be True
|
||||
regardless of routing method."""
|
||||
device = torch.device("cuda:0")
|
||||
for routing_method in (
|
||||
RoutingMethodType.DeepSeekV3,
|
||||
RoutingMethodType.Renormalize,
|
||||
RoutingMethodType.RenormalizeNaive,
|
||||
):
|
||||
experts, _, _ = _make_bf16_monolithic_experts(
|
||||
num_experts=_DSV3_NUM_EXPERTS,
|
||||
top_k=2,
|
||||
hidden_size=1024,
|
||||
intermediate_size=1024,
|
||||
routing_method=routing_method,
|
||||
device=device,
|
||||
)
|
||||
assert experts.supports_routing_replay_capture() is True, (
|
||||
f"{routing_method!r} should support routing replay capture"
|
||||
)
|
||||
|
||||
|
||||
def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None:
|
||||
"""When capture is installed, the allocated buffer is int16 and shaped
|
||||
``(num_tokens, experts_per_token)``."""
|
||||
device = torch.device("cuda:0")
|
||||
experts, _, _ = _make_bf16_monolithic_experts(
|
||||
num_experts=_DSV3_NUM_EXPERTS,
|
||||
top_k=4,
|
||||
hidden_size=1024,
|
||||
intermediate_size=1024,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
device=device,
|
||||
)
|
||||
experts.set_capture_fn(lambda r: None)
|
||||
buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device)
|
||||
assert buf is not None
|
||||
assert buf.dtype == torch.int16
|
||||
assert buf.shape[0] >= 11
|
||||
assert buf.shape[1] == 4
|
||||
assert buf.device.type == "cuda"
|
||||
|
||||
|
||||
def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
|
||||
"""End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback
|
||||
on the monolithic experts and verify the captured rows land in the
|
||||
capturer's device buffer at the correct layer slot.
|
||||
|
||||
Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
|
||||
for the monolithic path: a single closure is installed on the monolithic
|
||||
``fused_experts`` (in addition to ``router.set_capture_fn`` on the
|
||||
non-monolithic path) and the capturer routes per-layer based on the
|
||||
closed-over ``layer_id``.
|
||||
"""
|
||||
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
|
||||
RoutedExpertsCapturer,
|
||||
)
|
||||
|
||||
torch.manual_seed(7)
|
||||
device = torch.device("cuda:0")
|
||||
num_tokens = 4
|
||||
top_k = 2
|
||||
num_experts = _DSV3_NUM_EXPERTS
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
|
||||
experts, w13, w2 = _make_bf16_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
device=device,
|
||||
)
|
||||
|
||||
num_layers = 3
|
||||
layer_id = 1
|
||||
capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
|
||||
capturer.dp_rank = 0
|
||||
capturer.tp_size = 1
|
||||
capturer.device_buffer = torch.full(
|
||||
(num_tokens + 4, num_layers, top_k),
|
||||
-1,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def capture_fn(replay_out: torch.Tensor) -> None:
|
||||
capturer.capture(layer_id, replay_out)
|
||||
|
||||
experts.set_capture_fn(capture_fn)
|
||||
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
routing_bias = _make_dsv3_routing_bias(num_experts, device)
|
||||
|
||||
# Patch get_forward_context to return a dp_metadata=None context so the
|
||||
# capturer takes the single-DP branch.
|
||||
import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec
|
||||
|
||||
with patch.object(
|
||||
rec,
|
||||
"get_forward_context",
|
||||
return_value=SimpleNamespace(dp_metadata=None),
|
||||
):
|
||||
_ = _run_bf16_monolithic(
|
||||
experts,
|
||||
hidden_states=hidden_states,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
n_group=_DSV3_N_GROUP,
|
||||
topk_group=_DSV3_TOPK_GROUP,
|
||||
routed_scaling_factor=1.0,
|
||||
e_score_correction_bias=routing_bias,
|
||||
)
|
||||
|
||||
captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu()
|
||||
# Valid expert IDs at this layer.
|
||||
assert (captured >= 0).all()
|
||||
assert (captured < num_experts).all()
|
||||
for t in range(num_tokens):
|
||||
unique = captured[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts at layer "
|
||||
f"{layer_id}, got {unique.numel()}"
|
||||
)
|
||||
|
||||
# Other layers / trailing token rows untouched.
|
||||
for other_layer in range(num_layers):
|
||||
if other_layer == layer_id:
|
||||
continue
|
||||
assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), (
|
||||
f"layer {other_layer} should be untouched, got writes"
|
||||
)
|
||||
assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), (
|
||||
"tail rows beyond num_tokens should remain sentinel"
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic``
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_fp8_block_scale_monolithic_experts(
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
device: torch.device,
|
||||
) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]:
|
||||
"""Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale
|
||||
code path with DSV3 routing.
|
||||
|
||||
Weights are shuffled into the BlockMajorK layout the kernel expects
|
||||
(same helper the vLLM weight loader uses for DeepSeek-FP8 models).
|
||||
"""
|
||||
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
|
||||
_shuffle_deepseek_fp8_moe_weights,
|
||||
)
|
||||
|
||||
block_k = 128
|
||||
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
|
||||
moe_config = FusedMoEConfig(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=top_k,
|
||||
hidden_dim=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
moe_parallel_config=parallel_cfg,
|
||||
in_dtype=torch.bfloat16,
|
||||
activation=MoEActivation.SILU,
|
||||
device=device,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
max_num_tokens=max(8, 1),
|
||||
)
|
||||
|
||||
# Random fp8 weights + ones-block scales (the kernel decoder only cares
|
||||
# that the per-block scales are present and finite for routing/replay).
|
||||
gemm1 = torch.randn(
|
||||
num_experts, 2 * intermediate_size, hidden_size, device=device
|
||||
).to(torch.float8_e4m3fn)
|
||||
gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2)
|
||||
|
||||
w1_scale = torch.ones(
|
||||
num_experts,
|
||||
2 * intermediate_size // block_k,
|
||||
hidden_size // block_k,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
w2_scale = torch.ones(
|
||||
num_experts,
|
||||
hidden_size // block_k,
|
||||
intermediate_size // block_k,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
quant_config = fp8_w8a8_moe_quant_config(
|
||||
w1_scale=w1_scale,
|
||||
w2_scale=w2_scale,
|
||||
block_shape=[block_k, block_k],
|
||||
per_act_token_quant=False,
|
||||
)
|
||||
|
||||
experts = TrtLlmFp8ExpertsMonolithic(
|
||||
moe_config=moe_config, quant_config=quant_config
|
||||
)
|
||||
return experts, w13_shuffled, w2_shuffled
|
||||
|
||||
|
||||
def _run_fp8_block_scale_monolithic(
|
||||
experts: TrtLlmFp8ExpertsMonolithic,
|
||||
hidden_states_fp8: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
w13: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
num_experts: int,
|
||||
routing_bias: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return experts.apply(
|
||||
hidden_states=hidden_states_fp8,
|
||||
w1=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
activation=MoEActivation.SILU,
|
||||
global_num_experts=num_experts,
|
||||
expert_map=None,
|
||||
# The block-scale apply path reads ``a1q_scale`` and transposes it
|
||||
# to ``(hidden_size/128, num_tokens)`` for the kernel call.
|
||||
a1q_scale=hidden_states_scale,
|
||||
apply_router_weight_on_input=False,
|
||||
num_expert_group=_DSV3_N_GROUP,
|
||||
topk_group=_DSV3_TOPK_GROUP,
|
||||
e_score_correction_bias=routing_bias,
|
||||
routed_scaling_factor=1.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
|
||||
@pytest.mark.parametrize("top_k", [2, 4])
|
||||
def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts(
|
||||
num_tokens: int,
|
||||
top_k: int,
|
||||
) -> None:
|
||||
"""End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale
|
||||
path, DSV3 routing) captures valid expert IDs."""
|
||||
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
|
||||
pytest.skip(
|
||||
f"DSV3 requires top_k <= n_group * topk_group "
|
||||
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
|
||||
)
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
num_experts = _DSV3_NUM_EXPERTS
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
|
||||
experts, w13, w2 = _make_fp8_block_scale_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
device=device,
|
||||
)
|
||||
assert experts.supports_routing_replay_capture()
|
||||
|
||||
captured: list[torch.Tensor] = []
|
||||
experts.set_capture_fn(lambda r: captured.append(r.clone()))
|
||||
|
||||
# Per-token / per-block hidden scales (ones is fine for the routing
|
||||
# path; the GEMM output isn't being asserted on).
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
).to(torch.float8_e4m3fn)
|
||||
hidden_states_scale = torch.ones(
|
||||
num_tokens, hidden_size // 128, device=device, dtype=torch.float32
|
||||
)
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
routing_bias = _make_dsv3_routing_bias(num_experts, device)
|
||||
|
||||
_ = _run_fp8_block_scale_monolithic(
|
||||
experts,
|
||||
hidden_states_fp8=hidden_states,
|
||||
hidden_states_scale=hidden_states_scale,
|
||||
w13=w13,
|
||||
w2=w2,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
routing_bias=routing_bias,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
replay = captured[0]
|
||||
assert replay.dtype == torch.int16
|
||||
assert replay.shape == (num_tokens, top_k)
|
||||
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
|
||||
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
|
||||
for t in range(num_tokens):
|
||||
unique = replay[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts, "
|
||||
f"got {unique.numel()} ({replay[t].tolist()})"
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic``
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_nvfp4_monolithic_experts(
|
||||
num_experts: int,
|
||||
top_k: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
device: torch.device,
|
||||
) -> tuple[
|
||||
TrtLlmNvFp4ExpertsMonolithic,
|
||||
torch.Tensor, # w13 (packed nvfp4 uint8)
|
||||
torch.Tensor, # w13 block-scale (fp8)
|
||||
torch.Tensor, # w2 (packed nvfp4 uint8)
|
||||
torch.Tensor, # w2 block-scale (fp8)
|
||||
torch.Tensor, # input global scale (per-tensor float32)
|
||||
]:
|
||||
"""Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights
|
||||
and DSV3 routing.
|
||||
|
||||
NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor
|
||||
"global" scale. We follow the layout in ``test_ocp_mx_moe.py`` /
|
||||
``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``:
|
||||
* weights: uint8 (packed fp4) ``(E, M, K//2)``
|
||||
* weight scales: fp8 ``(E, M, K//16)``
|
||||
* hidden states: uint8 (packed fp4) ``(N, K//2)``
|
||||
* hidden state scales: fp8 ``(N, K//16)``
|
||||
"""
|
||||
from flashinfer import fp4_quantize
|
||||
|
||||
block_size = 16
|
||||
parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
|
||||
moe_config = FusedMoEConfig(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=top_k,
|
||||
hidden_dim=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
moe_parallel_config=parallel_cfg,
|
||||
in_dtype=torch.bfloat16,
|
||||
activation=MoEActivation.SILU,
|
||||
device=device,
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
max_num_tokens=max(8, 1),
|
||||
)
|
||||
|
||||
gemm1 = torch.randn(
|
||||
num_experts,
|
||||
2 * intermediate_size,
|
||||
hidden_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
gemm2 = torch.randn(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
# Per-tensor weight scaling factor (used to build ``g1_alphas`` /
|
||||
# ``g2_alphas`` below).
|
||||
w_global_scale = torch.tensor(1.0, device=device)
|
||||
# Per-tensor input scaling factor.
|
||||
a_global_scale = torch.tensor(1.0, device=device)
|
||||
|
||||
w13_q, w13_scale = fp4_quantize(
|
||||
gemm1,
|
||||
w_global_scale,
|
||||
block_size,
|
||||
sf_use_ue8m0=False,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
|
||||
num_experts, 2 * intermediate_size, hidden_size // block_size
|
||||
)
|
||||
w2_q, w2_scale = fp4_quantize(
|
||||
gemm2,
|
||||
w_global_scale,
|
||||
block_size,
|
||||
sf_use_ue8m0=False,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
|
||||
num_experts, hidden_size, intermediate_size // block_size
|
||||
)
|
||||
|
||||
# NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2,
|
||||
# g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these.
|
||||
g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
|
||||
a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
|
||||
|
||||
quant_config = FusedMoEQuantConfig.make(
|
||||
quant_dtype="nvfp4",
|
||||
per_act_token_quant=False,
|
||||
per_out_ch_quant=False,
|
||||
block_shape=None,
|
||||
w1_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
g1_alphas=g_alphas,
|
||||
g2_alphas=g_alphas,
|
||||
a1_gscale=a_global_scale,
|
||||
a2_gscale=a2_gscale,
|
||||
)
|
||||
|
||||
experts = TrtLlmNvFp4ExpertsMonolithic(
|
||||
moe_config=moe_config, quant_config=quant_config
|
||||
)
|
||||
return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale
|
||||
|
||||
|
||||
def _run_nvfp4_monolithic(
|
||||
experts: TrtLlmNvFp4ExpertsMonolithic,
|
||||
hidden_states_q: torch.Tensor,
|
||||
hidden_states_scale: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
num_experts: int,
|
||||
routing_bias: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""The monolithic NVFP4 apply expects packed fp4 hidden states + the
|
||||
matching fp8 per-block scale stored in the ``a1q_scale`` slot."""
|
||||
# Stash the weight tensors on the experts in the locations the apply()
|
||||
# implementation reads from (it pulls them from quant_config / scales
|
||||
# already; w1/w2 come in as args).
|
||||
return experts.apply(
|
||||
hidden_states=hidden_states_q,
|
||||
w1=experts._w13_packed,
|
||||
w2=experts._w2_packed,
|
||||
router_logits=router_logits,
|
||||
activation=MoEActivation.SILU,
|
||||
global_num_experts=num_experts,
|
||||
expert_map=None,
|
||||
a1q_scale=hidden_states_scale,
|
||||
apply_router_weight_on_input=False,
|
||||
num_expert_group=_DSV3_N_GROUP,
|
||||
topk_group=_DSV3_TOPK_GROUP,
|
||||
e_score_correction_bias=routing_bias,
|
||||
routed_scaling_factor=1.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [2, 7, 16])
|
||||
@pytest.mark.parametrize("top_k", [2, 4])
|
||||
def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts(
|
||||
num_tokens: int,
|
||||
top_k: int,
|
||||
) -> None:
|
||||
"""End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs
|
||||
on the DSV3 routing path."""
|
||||
if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
|
||||
pytest.skip(
|
||||
f"DSV3 requires top_k <= n_group * topk_group "
|
||||
f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
|
||||
)
|
||||
from flashinfer import fp4_quantize
|
||||
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda:0")
|
||||
|
||||
num_experts = _DSV3_NUM_EXPERTS
|
||||
hidden_size = 1024
|
||||
intermediate_size = 1024
|
||||
block_size = 16
|
||||
|
||||
experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts(
|
||||
num_experts=num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
device=device,
|
||||
)
|
||||
# The apply() reads w1/w2 from its args, but we keep them on the experts
|
||||
# for convenience of the helper.
|
||||
experts._w13_packed = w13_q
|
||||
experts._w2_packed = w2_q
|
||||
|
||||
assert experts.supports_routing_replay_capture()
|
||||
captured: list[torch.Tensor] = []
|
||||
experts.set_capture_fn(lambda r: captured.append(r.clone()))
|
||||
|
||||
hidden_states = (
|
||||
torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
hidden_states_q, hidden_states_scale = fp4_quantize(
|
||||
hidden_states,
|
||||
a_gs,
|
||||
block_size,
|
||||
sf_use_ue8m0=False,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
# The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave
|
||||
# ``hidden_states_scale`` in its native (uint8 packed) form.
|
||||
router_logits = torch.rand(
|
||||
num_tokens, num_experts, device=device, dtype=torch.float32
|
||||
)
|
||||
routing_bias = _make_dsv3_routing_bias(num_experts, device)
|
||||
|
||||
_ = _run_nvfp4_monolithic(
|
||||
experts,
|
||||
hidden_states_q=hidden_states_q,
|
||||
hidden_states_scale=hidden_states_scale,
|
||||
router_logits=router_logits,
|
||||
num_experts=num_experts,
|
||||
routing_bias=routing_bias,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
replay = captured[0]
|
||||
assert replay.dtype == torch.int16
|
||||
assert replay.shape == (num_tokens, top_k)
|
||||
assert (replay >= 0).all(), f"got out-of-range values: {replay}"
|
||||
assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
|
||||
for t in range(num_tokens):
|
||||
unique = replay[t].unique()
|
||||
assert unique.numel() == top_k, (
|
||||
f"token {t}: expected {top_k} distinct experts, "
|
||||
f"got {unique.numel()} ({replay[t].tolist()})"
|
||||
)
|
||||
@@ -0,0 +1,31 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
|
||||
def test_cutlass_group_gemm_python_guard_allows_thor(monkeypatch):
|
||||
seen_capabilities: list[int] = []
|
||||
|
||||
def fake_cutlass_group_gemm_supported(capability: int) -> bool:
|
||||
seen_capabilities.append(capability)
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(
|
||||
torch.ops,
|
||||
"_C",
|
||||
SimpleNamespace(cutlass_group_gemm_supported=fake_cutlass_group_gemm_supported),
|
||||
)
|
||||
|
||||
# CUDA 12 reports Thor as SM101 and CUDA 13 reports it as SM110. Both
|
||||
# should reach the C++ query for the SM10x/SM11x CUTLASS MoE kernel.
|
||||
assert ops.cutlass_group_gemm_supported(101)
|
||||
assert ops.cutlass_group_gemm_supported(110)
|
||||
|
||||
# SM120 uses separate kernels and must not be advertised by this path.
|
||||
assert not ops.cutlass_group_gemm_supported(120)
|
||||
assert seen_capabilities == [101, 110]
|
||||
@@ -0,0 +1,306 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import ast
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.model_executor.warmup.jit_warmup import (
|
||||
VllmJitKernel,
|
||||
WarmupIntRange,
|
||||
get_ast_full_name,
|
||||
zip_inputs,
|
||||
)
|
||||
|
||||
|
||||
def _next_power_of_2(value: int) -> int:
|
||||
return 1 << max(0, value - 1).bit_length()
|
||||
|
||||
|
||||
def _round_up(value: int, *, multiple: int) -> int:
|
||||
return ((value + multiple - 1) // multiple) * multiple
|
||||
|
||||
|
||||
def _config(
|
||||
*,
|
||||
bias: int = 0,
|
||||
disabled: bool = False,
|
||||
name: str = "base",
|
||||
vectorized: bool = False,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
bias=bias,
|
||||
disabled=disabled,
|
||||
name=name,
|
||||
vectorized=vectorized,
|
||||
)
|
||||
|
||||
|
||||
class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
block_size: int
|
||||
work: int
|
||||
vector_width: int
|
||||
descriptor: tuple[object, ...]
|
||||
enabled: bool
|
||||
|
||||
def dispatch( # type: ignore[override]
|
||||
self,
|
||||
*,
|
||||
tokens: int,
|
||||
cfg: Any,
|
||||
lanes: int = 1,
|
||||
mode: str = "default",
|
||||
debug: int = 0,
|
||||
) -> CompileKey:
|
||||
block_size = _next_power_of_2(tokens)
|
||||
work: int = block_size * lanes + cfg.bias
|
||||
return self.CompileKey(
|
||||
block_size=block_size,
|
||||
work=work,
|
||||
vector_width=4 if cfg.vectorized and block_size >= 4 else 1,
|
||||
descriptor=(
|
||||
cfg.name,
|
||||
mode,
|
||||
-block_size,
|
||||
block_size % 3,
|
||||
block_size**2,
|
||||
),
|
||||
enabled=not cfg.disabled,
|
||||
)
|
||||
|
||||
def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]:
|
||||
return self._trace_dispatch(self.dispatch)(
|
||||
tokens=WarmupIntRange(1, max_tokens + 1),
|
||||
cfg=cfg,
|
||||
# This argument is intentionally unused by dispatch expressions.
|
||||
debug=WarmupIntRange(0, 100),
|
||||
)
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class RecordingToyKernel(ToyKernel):
|
||||
def __init__(self) -> None:
|
||||
self.compiled: list[ToyKernel.CompileKey] = []
|
||||
super().__init__()
|
||||
|
||||
def compile(self, compile_key: ToyKernel.CompileKey) -> None:
|
||||
self.compiled.append(compile_key)
|
||||
|
||||
|
||||
def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None:
|
||||
cfg = _config()
|
||||
|
||||
assert ToyKernel().get_warmup_keys(5, cfg) == [
|
||||
ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
|
||||
ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
|
||||
ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
|
||||
ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
|
||||
]
|
||||
|
||||
|
||||
def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
|
||||
cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
|
||||
|
||||
assert ToyKernel().compile_key(
|
||||
{
|
||||
"tokens": 4,
|
||||
"cfg": cfg,
|
||||
"lanes": 2,
|
||||
}
|
||||
) == ToyKernel.CompileKey(
|
||||
block_size=4,
|
||||
work=11,
|
||||
vector_width=4,
|
||||
descriptor=("cfg", "default", -4, 1, 16),
|
||||
enabled=False,
|
||||
)
|
||||
|
||||
|
||||
def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
|
||||
cfg = _config(vectorized=True)
|
||||
|
||||
keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
|
||||
zip_inputs(
|
||||
dict(tokens=1, mode="small"),
|
||||
dict(tokens=4, mode="wide"),
|
||||
),
|
||||
cfg=cfg,
|
||||
lanes=(1, 2),
|
||||
)
|
||||
|
||||
assert keys == [
|
||||
ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True),
|
||||
ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True),
|
||||
ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True),
|
||||
ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True),
|
||||
]
|
||||
|
||||
|
||||
def test_zip_inputs_validates_input_rows() -> None:
|
||||
with pytest.raises(ValueError, match="requires at least one"):
|
||||
zip_inputs()
|
||||
with pytest.raises(ValueError, match="rows must be mappings"):
|
||||
zip_inputs(cast(Any, ("tokens", 1)))
|
||||
with pytest.raises(ValueError, match="at least one dispatch input name"):
|
||||
zip_inputs({})
|
||||
with pytest.raises(ValueError, match="dispatch input names must be strings"):
|
||||
zip_inputs(cast(Any, {1: 2}))
|
||||
with pytest.raises(ValueError, match="same dispatch input names"):
|
||||
zip_inputs({"tokens": 1}, {"mode": "small"})
|
||||
|
||||
|
||||
def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None:
|
||||
kernel = ToyKernel()
|
||||
|
||||
with pytest.raises(TypeError, match="zip_inputs"):
|
||||
kernel._trace_dispatch(kernel.dispatch)(
|
||||
cast(Any, {"tokens": 1}),
|
||||
cfg=_config(),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="specified more than once"):
|
||||
kernel._trace_dispatch(kernel.dispatch)(
|
||||
zip_inputs(dict(tokens=1, mode="small")),
|
||||
tokens=2,
|
||||
cfg=_config(),
|
||||
)
|
||||
|
||||
|
||||
def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
|
||||
class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: int
|
||||
|
||||
def dispatch( # type: ignore[override]
|
||||
self,
|
||||
*,
|
||||
tokens: int,
|
||||
block_size: int,
|
||||
) -> CompileKey:
|
||||
return self.CompileKey(value=_round_up(tokens, multiple=block_size))
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: int
|
||||
|
||||
def dispatch( # type: ignore[override]
|
||||
self,
|
||||
*,
|
||||
tokens: int,
|
||||
block_size: int,
|
||||
) -> CompileKey:
|
||||
return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size}))
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
assert HelperKernel().compile_key(
|
||||
{
|
||||
"tokens": 5,
|
||||
"block_size": 4,
|
||||
}
|
||||
) == HelperKernel.CompileKey(value=8)
|
||||
with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"):
|
||||
StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
|
||||
|
||||
|
||||
def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
|
||||
class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: int
|
||||
|
||||
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
|
||||
if value > 0:
|
||||
value = 1
|
||||
return self.CompileKey(value=value)
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: int
|
||||
|
||||
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
|
||||
return self.CompileKey(**{"value": value})
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="local assignments"):
|
||||
BranchKernel()
|
||||
with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"):
|
||||
KwargsReturnKernel()
|
||||
|
||||
|
||||
def test_dispatch_reports_unsupported_expression_with_context() -> None:
|
||||
class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
|
||||
@dataclass(frozen=True)
|
||||
class CompileKey:
|
||||
value: object
|
||||
|
||||
def dispatch(self, *, value: int) -> CompileKey: # type: ignore[override]
|
||||
return self.CompileKey(value={value})
|
||||
|
||||
def get_warmup_keys(self) -> list[CompileKey]:
|
||||
return []
|
||||
|
||||
def compile(self, compile_key: CompileKey) -> None:
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
UnsupportedKernel().compile_key({"value": 1})
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "Unsupported dispatch expression" in message
|
||||
assert "{value}" in message
|
||||
assert "Supported dispatch expressions" in message
|
||||
|
||||
|
||||
def test_warmup_compiles_all_returned_keys_in_order() -> None:
|
||||
kernel = RecordingToyKernel()
|
||||
cfg = _config()
|
||||
|
||||
kernel.warmup(3, cfg)
|
||||
|
||||
assert kernel.compiled == [
|
||||
ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
|
||||
ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
|
||||
ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
|
||||
]
|
||||
|
||||
|
||||
def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
|
||||
dotted_expr = ast.parse("foo.bar.baz").body[0]
|
||||
call_expr = ast.parse("foo()").body[0]
|
||||
assert isinstance(dotted_expr, ast.Expr)
|
||||
assert isinstance(call_expr, ast.Expr)
|
||||
|
||||
assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz"
|
||||
assert get_ast_full_name(call_expr.value) is None
|
||||
@@ -66,6 +66,12 @@ def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
|
||||
)
|
||||
|
||||
|
||||
def _make_modular_routed_experts():
|
||||
return types.SimpleNamespace(
|
||||
quant_method=types.SimpleNamespace(is_monolithic=False),
|
||||
)
|
||||
|
||||
|
||||
def test_base_router_capture_pre_eplb_mapping():
|
||||
router = _make_router()
|
||||
captured = []
|
||||
@@ -122,6 +128,8 @@ def test_gpu_model_runner_binds_router_capture(monkeypatch):
|
||||
def __init__(self):
|
||||
self.layer_id = 7
|
||||
self.router = _make_router()
|
||||
self.routed_experts = _make_modular_routed_experts()
|
||||
self._quant_method = self.routed_experts.quant_method
|
||||
|
||||
class DummyCapturer:
|
||||
def __init__(self):
|
||||
@@ -160,6 +168,8 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
|
||||
def __init__(self):
|
||||
self.layer_id = 11
|
||||
self.router = _make_router()
|
||||
self.routed_experts = _make_modular_routed_experts()
|
||||
self._quant_method = self.routed_experts.quant_method
|
||||
|
||||
class DummyCapturer:
|
||||
def __init__(self):
|
||||
@@ -197,6 +207,8 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
|
||||
def __init__(self, layer_id):
|
||||
self.layer_id = layer_id
|
||||
self.router = _make_router()
|
||||
self.routed_experts = _make_modular_routed_experts()
|
||||
self._quant_method = self.routed_experts.quant_method
|
||||
|
||||
target_module = DummyFusedMoE(layer_id=7)
|
||||
draft_module = DummyFusedMoE(layer_id=0)
|
||||
@@ -222,6 +234,49 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
|
||||
assert draft_module.router.capture_fn is None
|
||||
|
||||
|
||||
def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
|
||||
from vllm.v1.worker import gpu_model_runner as gmr
|
||||
|
||||
class DummyFusedMoE:
|
||||
def __init__(self):
|
||||
self.layer_id = 3
|
||||
self.router = _make_router()
|
||||
# Use a concrete monolithic expert and override its capability
|
||||
# instead of instantiating the abstract base class directly.
|
||||
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
|
||||
CPUExpertsFp8,
|
||||
)
|
||||
|
||||
fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8)
|
||||
self.routed_experts = types.SimpleNamespace(
|
||||
quant_method=types.SimpleNamespace(
|
||||
is_monolithic=True,
|
||||
moe_kernel=types.SimpleNamespace(
|
||||
impl=types.SimpleNamespace(fused_experts=fused_experts)
|
||||
),
|
||||
)
|
||||
)
|
||||
self._quant_method = self.routed_experts.quant_method
|
||||
self._quant_method.moe_kernel.impl.fused_experts = fused_experts
|
||||
fused_experts.supports_routing_replay_capture = lambda: False
|
||||
|
||||
class DummyCapturer:
|
||||
def capture(self, layer_id, topk_ids):
|
||||
pass
|
||||
|
||||
dummy_module = DummyFusedMoE()
|
||||
import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
|
||||
|
||||
monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
|
||||
|
||||
dummy_self = types.SimpleNamespace(
|
||||
model=types.SimpleNamespace(modules=lambda: [dummy_module])
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="monolithic MoE kernel"):
|
||||
gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
|
||||
|
||||
|
||||
def test_routed_experts_capturer_single_dp_no_metadata():
|
||||
"""dp_metadata is None: capture writes the full topk_ids rows."""
|
||||
capturer = _capturer_with_buffer(dp_rank=0)
|
||||
|
||||
@@ -9,6 +9,7 @@ import torch.nn.functional as F
|
||||
from transformers import AutoModel
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.mem_constants import MiB_bytes
|
||||
|
||||
from ....conftest import HfRunner
|
||||
from ....utils import VLLM_PATH
|
||||
@@ -106,7 +107,12 @@ def test_prm_models(
|
||||
if current_platform.is_cpu():
|
||||
pytest.skip("CPU only supports V1")
|
||||
|
||||
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=1024,
|
||||
dtype=dtype,
|
||||
kv_cache_memory_bytes=64 * MiB_bytes,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(math_step_prompts)
|
||||
|
||||
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
|
||||
@@ -145,7 +151,12 @@ def test_prm_models_with_golden_outputs(
|
||||
if not FIXTURE_REWARD_RESULT.get(model):
|
||||
pytest.skip(f"No available golden outputs for {model}.")
|
||||
|
||||
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
|
||||
with vllm_runner(
|
||||
model,
|
||||
max_model_len=1024,
|
||||
dtype=dtype,
|
||||
kv_cache_memory_bytes=64 * MiB_bytes,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_classify(math_step_prompts)
|
||||
|
||||
golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])
|
||||
|
||||
@@ -74,6 +74,44 @@ def test_cosmos3_new_checkpoint_weights_mapper():
|
||||
)
|
||||
|
||||
|
||||
def test_cosmos3_modelopt_quantizer_weights_mapper():
|
||||
"""ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers
|
||||
(``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable
|
||||
``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former
|
||||
(it has no parameter for them) while keeping the latter."""
|
||||
from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration
|
||||
|
||||
mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper
|
||||
|
||||
# Native ModelOpt quantizer buffers are dropped.
|
||||
assert (
|
||||
mapper.apply_list(
|
||||
[
|
||||
"layers.0.self_attn.to_q.input_quantizer._amax",
|
||||
"layers.0.self_attn.to_q.weight_quantizer._amax",
|
||||
"layers.0.self_attn.to_q.weight_quantizer._scale",
|
||||
"layers.0.mlp.down_proj.output_quantizer._amax",
|
||||
]
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
# The FP8 scale sidecars vLLM actually consumes are kept and remapped.
|
||||
assert mapper.apply_list(
|
||||
[
|
||||
"layers.0.self_attn.to_q.weight",
|
||||
"layers.0.self_attn.to_q.weight_scale",
|
||||
"layers.0.self_attn.to_q.input_scale",
|
||||
"layers.0.mlp.down_proj.input_scale",
|
||||
]
|
||||
) == [
|
||||
"language_model.model.layers.0.self_attn.q_proj.weight",
|
||||
"language_model.model.layers.0.self_attn.q_proj.weight_scale",
|
||||
"language_model.model.layers.0.self_attn.q_proj.input_scale",
|
||||
"language_model.model.layers.0.mlp.down_proj.input_scale",
|
||||
]
|
||||
|
||||
|
||||
def test_cosmos3_edge_checkpoint_weights_mapper():
|
||||
from vllm.model_executor.models.cosmos3_edge import (
|
||||
Cosmos3EdgeForConditionalGeneration,
|
||||
|
||||
@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.
|
||||
import pytest
|
||||
|
||||
from tests.quantization.utils import is_quant_method_supported
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_kv_cache_dtype
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
|
||||
from ..utils import check_logprobs_close
|
||||
|
||||
|
||||
@@ -70,13 +70,7 @@ def test_models(
|
||||
if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():
|
||||
pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")
|
||||
|
||||
if not (
|
||||
current_platform.is_xpu()
|
||||
or (
|
||||
get_flash_attn_version() == 3
|
||||
and current_platform.is_device_capability_family(90)
|
||||
)
|
||||
):
|
||||
if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype):
|
||||
pytest.skip(
|
||||
f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
|
||||
)
|
||||
|
||||
@@ -6,6 +6,7 @@ import mimetypes
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from io import BytesIO
|
||||
from tempfile import NamedTemporaryFile, TemporaryDirectory
|
||||
|
||||
import aiohttp
|
||||
@@ -111,6 +112,34 @@ async def test_fetch_image_base64(
|
||||
assert _image_equals(data_image_sync, data_image_async)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_image_keep_original_mode():
|
||||
"""media_io_kwargs can disable the default RGB conversion."""
|
||||
# RGBA image: opaque black pixel on a fully transparent background
|
||||
rgba_image = Image.new("RGBA", (4, 4), (0, 0, 0, 0))
|
||||
rgba_image.putpixel((2, 2), (0, 0, 0, 255))
|
||||
buffer = BytesIO()
|
||||
rgba_image.save(buffer, "PNG")
|
||||
data_url = (
|
||||
f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
|
||||
)
|
||||
|
||||
# Default behavior: RGBA is composited onto a white background
|
||||
default_image = MediaConnector().fetch_image(data_url)
|
||||
assert default_image.mode == "RGB"
|
||||
assert default_image.getpixel((0, 0)) == (255, 255, 255)
|
||||
assert default_image.getpixel((2, 2)) == (0, 0, 0)
|
||||
|
||||
# image_mode=None via media_io_kwargs: original mode is preserved
|
||||
connector = MediaConnector(media_io_kwargs={"image": {"image_mode": None}})
|
||||
image_sync = connector.fetch_image(data_url)
|
||||
image_async = await connector.fetch_image_async(data_url)
|
||||
for image in (image_sync, image_async):
|
||||
assert image.mode == "RGBA"
|
||||
assert image.getpixel((0, 0)) == (0, 0, 0, 0)
|
||||
assert image.getpixel((2, 2)) == (0, 0, 0, 255)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
|
||||
async def test_fetch_image_local_files(image_url: str):
|
||||
|
||||
@@ -80,6 +80,29 @@ def test_image_media_io_rgba_custom_background(tmp_path):
|
||||
assert green_numpy[0][0][2] == 0 # B
|
||||
|
||||
|
||||
def test_image_media_io_no_mode_conversion(tmp_path):
|
||||
"""image_mode=None skips conversion and preserves the original mode."""
|
||||
# RGBA image: opaque black pixel on a fully transparent background
|
||||
rgba_image = Image.new("RGBA", (10, 10), (0, 0, 0, 0))
|
||||
rgba_image.putpixel((5, 5), (0, 0, 0, 255))
|
||||
test_image_path = tmp_path / "test_rgba.png"
|
||||
rgba_image.save(test_image_path)
|
||||
|
||||
# Default behavior: RGBA is composited onto a white background
|
||||
image_io_default = ImageMediaIO()
|
||||
converted_default = image_io_default.load_file(test_image_path)
|
||||
assert converted_default.media.mode == "RGB"
|
||||
assert converted_default.media.getpixel((0, 0)) == (255, 255, 255)
|
||||
assert converted_default.media.getpixel((5, 5)) == (0, 0, 0)
|
||||
|
||||
# image_mode=None: original mode and alpha channel are preserved
|
||||
image_io_keep = ImageMediaIO(image_mode=None)
|
||||
converted_keep = image_io_keep.load_file(test_image_path)
|
||||
assert converted_keep.media.mode == "RGBA"
|
||||
assert converted_keep.media.getpixel((0, 0)) == (0, 0, 0, 0)
|
||||
assert converted_keep.media.getpixel((5, 5)) == (0, 0, 0, 255)
|
||||
|
||||
|
||||
def test_image_media_io_rgba_background_color_validation():
|
||||
"""Test that invalid rgba_background_color values are properly rejected."""
|
||||
|
||||
|
||||
@@ -48,12 +48,12 @@ def _check_dense_embedding(data, index=0):
|
||||
def _check_sparse_embedding(data, check_tokens=False):
|
||||
expected_weights = [
|
||||
{"token_id": 32, "weight": 0.0552978515625, "token": "?"},
|
||||
{"token_id": 70, "weight": 0.09808349609375, "token": " the"},
|
||||
{"token_id": 83, "weight": 0.08154296875, "token": " is"},
|
||||
{"token_id": 111, "weight": 0.11810302734375, "token": " of"},
|
||||
{"token_id": 4865, "weight": 0.1171875, "token": " What"},
|
||||
{"token_id": 9942, "weight": 0.292236328125, "token": " France"},
|
||||
{"token_id": 10323, "weight": 0.2802734375, "token": " capital"},
|
||||
{"token_id": 70, "weight": 0.09808349609375, "token": "the"},
|
||||
{"token_id": 83, "weight": 0.08154296875, "token": "is"},
|
||||
{"token_id": 111, "weight": 0.11810302734375, "token": "of"},
|
||||
{"token_id": 4865, "weight": 0.1171875, "token": "What"},
|
||||
{"token_id": 9942, "weight": 0.292236328125, "token": "France"},
|
||||
{"token_id": 10323, "weight": 0.2802734375, "token": "capital"},
|
||||
]
|
||||
expected_embed = {x["token_id"]: x for x in expected_weights}
|
||||
|
||||
|
||||
@@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
|
||||
assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None
|
||||
|
||||
|
||||
def test_modelopt_mixed_precision_composes_gemma4_mappers():
|
||||
from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
|
||||
from vllm.model_executor.models.gemma4_mm import (
|
||||
Gemma4ForConditionalGeneration,
|
||||
)
|
||||
|
||||
config = _mixed_precision_config(
|
||||
{
|
||||
"model.language_model.layers.0.experts": {
|
||||
"quant_algo": "NVFP4",
|
||||
"group_size": 16,
|
||||
},
|
||||
"model.language_model.layers.1.moe.experts.gate_up_proj": {
|
||||
"quant_algo": "NVFP4",
|
||||
"group_size": 16,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
config.apply_vllm_mapper(
|
||||
Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
|
||||
)
|
||||
config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
|
||||
|
||||
expected_prefix = "language_model.model.layers.0.moe.experts"
|
||||
assert set(config.quantized_layers) == {
|
||||
expected_prefix,
|
||||
"language_model.model.layers.1.moe.gate_up_proj",
|
||||
}
|
||||
assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
|
||||
|
||||
|
||||
def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
|
||||
from vllm.model_executor.layers.linear import LinearBase
|
||||
|
||||
|
||||
@@ -67,6 +67,23 @@ def test_memory_profiling():
|
||||
non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa
|
||||
assert abs(non_torch_ratio - 1) <= 0.05
|
||||
assert result.torch_peak_increase == 1024 * 1024 * 1024
|
||||
|
||||
expected_total_consumed = (256 + 512) * 1024 * 1024
|
||||
total_consumed_ratio = result.total_consumed / expected_total_consumed
|
||||
assert abs(total_consumed_ratio - 1) <= 0.05, (
|
||||
f"total_consumed={result.total_consumed}, "
|
||||
f"expected={expected_total_consumed}, "
|
||||
f"ratio={total_consumed_ratio}"
|
||||
)
|
||||
|
||||
expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
|
||||
non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
|
||||
assert abs(non_kv_ratio - 1) <= 0.05, (
|
||||
f"non_kv_cache_memory={result.non_kv_cache_memory}, "
|
||||
f"expected={expected_non_kv}, "
|
||||
f"ratio={non_kv_ratio}"
|
||||
)
|
||||
|
||||
del weights
|
||||
lib.cudaFree(handle1)
|
||||
lib.cudaFree(handle2)
|
||||
|
||||
@@ -21,6 +21,7 @@ from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.utils.torch_utils import (
|
||||
STR_DTYPE_TO_TORCH_DTYPE,
|
||||
is_quantized_kv_cache,
|
||||
is_torch_equal_or_newer,
|
||||
set_random_seed,
|
||||
)
|
||||
@@ -45,6 +46,11 @@ BACKENDS_TO_TEST = [
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
|
||||
FP8_KV_CACHE_DTYPES = {
|
||||
"fp8": torch.float8_e4m3fn,
|
||||
"fp8_e4m3": torch.float8_e4m3fn,
|
||||
}
|
||||
|
||||
# Remove flashinfer from the list if it's not available
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
@@ -110,6 +116,7 @@ def create_and_prepopulate_kv_cache(
|
||||
num_blocks: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
randomize_blocks: bool = True,
|
||||
kv_cache_dtype: str = "auto",
|
||||
) -> torch.Tensor:
|
||||
"""Create and prepopulate a KV cache with context data.
|
||||
|
||||
@@ -140,8 +147,18 @@ def create_and_prepopulate_kv_cache(
|
||||
block_table = common_attn_metadata.block_table_tensor
|
||||
slot_mapping = common_attn_metadata.slot_mapping
|
||||
|
||||
# For an fp8 kv cache, store the cache in the fp8 dtype so that assigning
|
||||
# the higher-precision context tensors quantizes them, mirroring runtime.
|
||||
fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
|
||||
storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype
|
||||
|
||||
kv_cache = torch.zeros(
|
||||
num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device
|
||||
num_blocks,
|
||||
block_size,
|
||||
num_kv_heads,
|
||||
2 * head_size,
|
||||
dtype=storage_dtype,
|
||||
device=device,
|
||||
)
|
||||
kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
|
||||
|
||||
@@ -195,7 +212,12 @@ def create_and_prepopulate_kv_cache(
|
||||
] * block_size + token_inter_block_offsets.to(device)
|
||||
|
||||
# Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
|
||||
return kv_cache.transpose(1, 2).contiguous()
|
||||
kv_cache = kv_cache.transpose(1, 2).contiguous()
|
||||
|
||||
if fp8_kv_cache:
|
||||
kv_cache = kv_cache.view(torch.uint8)
|
||||
|
||||
return kv_cache
|
||||
|
||||
|
||||
class MockAttentionLayer:
|
||||
@@ -224,6 +246,7 @@ def run_attention_backend(
|
||||
kv_cache: torch.Tensor,
|
||||
attn_type: AttentionType = AttentionType.DECODER,
|
||||
sliding_window: int | None = None,
|
||||
kv_cache_dtype: str = "auto",
|
||||
) -> torch.Tensor:
|
||||
"""Run attention computation using the specified backend's AttentionImpl."""
|
||||
|
||||
@@ -291,13 +314,16 @@ def run_attention_backend(
|
||||
alibi_slopes=None,
|
||||
sliding_window=sliding_window,
|
||||
attn_type=attn_type,
|
||||
kv_cache_dtype="auto",
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
# Create mock layer and output buffer
|
||||
mock_layer = MockAttentionLayer(device)
|
||||
output = torch.empty_like(query)
|
||||
|
||||
if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:
|
||||
query = query.to(current_platform.fp8_dtype())
|
||||
|
||||
# Run forward pass
|
||||
# NOTE: The query, key, and value are already shaped correctly
|
||||
# in the calling test function.
|
||||
@@ -324,6 +350,7 @@ def _test_backend_correctness(
|
||||
atol: float = 1e-2,
|
||||
rtol: float = 1e-2,
|
||||
tensor_parallel_size: int = 1,
|
||||
kv_cache_dtype: str = "auto",
|
||||
):
|
||||
"""
|
||||
Test that all backends produce similar outputs to a reference implementation
|
||||
@@ -372,6 +399,7 @@ def _test_backend_correctness(
|
||||
num_gpu_blocks=8192,
|
||||
hf_config_override=hf_config_override,
|
||||
)
|
||||
vllm_config.cache_config.cache_dtype = kv_cache_dtype
|
||||
device = torch.device(f"{DEVICE_TYPE}:0")
|
||||
|
||||
kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type)
|
||||
@@ -392,6 +420,13 @@ def _test_backend_correctness(
|
||||
block_size = vllm_config.cache_config.block_size
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
|
||||
fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
|
||||
if fp8_kv_cache:
|
||||
query_fp8_dtype = current_platform.fp8_dtype()
|
||||
kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype]
|
||||
atol = max(atol, 6e-2)
|
||||
rtol = max(rtol, 1e-1)
|
||||
|
||||
# 2. Generate data and compute SDPA reference output
|
||||
all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
|
||||
all_sdpa_outputs = []
|
||||
@@ -407,10 +442,17 @@ def _test_backend_correctness(
|
||||
k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
|
||||
|
||||
if fp8_kv_cache:
|
||||
q_ref = q.to(query_fp8_dtype).to(dtype)
|
||||
k_ref = k_full.to(kv_fp8_dtype).to(dtype)
|
||||
v_ref = v_full.to(kv_fp8_dtype).to(dtype)
|
||||
else:
|
||||
q_ref, k_ref, v_ref = q, k_full, v_full
|
||||
|
||||
# SDPA expects (N, H, L, D), so unsqueeze batch and permute
|
||||
q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
|
||||
q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2)
|
||||
k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2)
|
||||
v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2)
|
||||
|
||||
if num_q_heads != num_kv_heads:
|
||||
assert num_q_heads % num_kv_heads == 0, (
|
||||
@@ -471,6 +513,7 @@ def _test_backend_correctness(
|
||||
num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
randomize_blocks=True,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
# 4. Run vLLM backends and compare
|
||||
@@ -488,6 +531,12 @@ def _test_backend_correctness(
|
||||
else:
|
||||
backend_cls = None
|
||||
|
||||
if is_quantized_kv_cache(kv_cache_dtype) and (
|
||||
backend_cls is None
|
||||
or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
|
||||
):
|
||||
continue
|
||||
|
||||
if backend_name == AttentionBackendEnum.FLASHINFER:
|
||||
set_kv_cache_layout("HND")
|
||||
reset_kv_cache_layout = True
|
||||
@@ -521,6 +570,7 @@ def _test_backend_correctness(
|
||||
kv_cache_for_backend,
|
||||
sliding_window=sliding_window,
|
||||
attn_type=attn_type,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
finally:
|
||||
if reset_kv_cache_layout:
|
||||
@@ -570,8 +620,13 @@ def _test_backend_correctness(
|
||||
)
|
||||
@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
|
||||
def test_causal_backend_correctness(
|
||||
default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
|
||||
default_vllm_config,
|
||||
batch_spec_name: str,
|
||||
model: str,
|
||||
tensor_parallel_size: int,
|
||||
kv_cache_dtype: str,
|
||||
):
|
||||
"""Test backend's correctness with causal attention."""
|
||||
|
||||
@@ -612,6 +667,7 @@ def test_causal_backend_correctness(
|
||||
SMALL_BLOCK_BACKENDS,
|
||||
causal_mask_mod,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
# Fast FlexAttention needs to run with block_size=128
|
||||
@@ -623,6 +679,7 @@ def test_causal_backend_correctness(
|
||||
causal_mask_mod,
|
||||
block_size=128,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -119,6 +119,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
|
||||
layer.kv_b_proj.quant_method = None
|
||||
layer.is_aiter_triton_fp4_bmm_enabled = False
|
||||
layer.is_aiter_triton_fp8_bmm_enabled = False
|
||||
layer.dcp_q_replicate = False
|
||||
layer.quant_config = None
|
||||
layer.layer_name = "test"
|
||||
|
||||
|
||||
@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (
|
||||
EncoderOnlyAttentionSpec,
|
||||
FullAttentionSpec,
|
||||
MambaSpec,
|
||||
get_kv_quant_mode,
|
||||
)
|
||||
|
||||
|
||||
@@ -178,6 +179,7 @@ def create_standard_kv_cache_spec(
|
||||
head_size=vllm_config.model_config.get_head_size(),
|
||||
dtype=vllm_config.model_config.dtype,
|
||||
sliding_window=vllm_config.model_config.get_sliding_window(),
|
||||
kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -9,7 +9,6 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler
|
||||
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import RequestStatus
|
||||
from vllm.v1.structured_output import StructuredOutputGrammar
|
||||
from vllm.v1.utils import ConstantList
|
||||
|
||||
from .utils import create_requests, create_scheduler
|
||||
@@ -263,7 +262,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
scheduler = object.__new__(AsyncScheduler)
|
||||
request = create_requests(num_requests=1, num_tokens=1)[0]
|
||||
request.structured_output_request = Mock()
|
||||
request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
|
||||
request.structured_output_request.grammar = Mock()
|
||||
request.structured_output_request.grammar.accept_tokens.return_value = False
|
||||
request.status = RequestStatus.RUNNING
|
||||
request.num_computed_tokens = request.num_tokens
|
||||
@@ -285,7 +284,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
scheduler.kv_event_publisher = Mock()
|
||||
scheduler.finished_req_ids = set()
|
||||
scheduler.finished_req_ids_dict = None
|
||||
scheduler.grammar_compile_error_reqs = set()
|
||||
scheduler.vllm_config = Mock()
|
||||
scheduler.vllm_config.model_config.enable_return_routed_experts = False
|
||||
scheduler.enable_return_routed_experts = False
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import copy
|
||||
import hashlib
|
||||
import importlib
|
||||
from collections.abc import Callable
|
||||
@@ -1477,6 +1478,128 @@ def test_get_max_concurrency_for_kv_cache_config():
|
||||
assert num_tokens == max_concurrency_hybrid_model * max_model_len
|
||||
assert max_concurrency == max_concurrency_hybrid_model
|
||||
|
||||
# Unequal group sizes in the standard layout: each group's pages cost
|
||||
# whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the
|
||||
# same as the equal-hybrid case above, regardless of the second group
|
||||
# holding only 2 layers.
|
||||
kv_cache_config_unequal_groups = KVCacheConfig(
|
||||
num_blocks=1153 * 3,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
|
||||
KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec),
|
||||
],
|
||||
)
|
||||
assert (
|
||||
get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_unequal_groups
|
||||
)
|
||||
== 3
|
||||
)
|
||||
|
||||
# UniformTypeKVCacheSpecs group (worker config shape): the aggregated
|
||||
# spec's memory/page ratio equals a single layer's page count, so the
|
||||
# group needs 1024 blocks and the request 1153 in total. The previous
|
||||
# formula normalized both groups' memory by the first group's page size,
|
||||
# reporting 3459/1057 = 3.27 here instead of 3 — and a different value
|
||||
# again for the scheduler-config shape below.
|
||||
uniform_full_spec = UniformTypeKVCacheSpecs(
|
||||
block_size=full_attention_spec.block_size,
|
||||
kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)},
|
||||
)
|
||||
kv_cache_config_uniform_group = KVCacheConfig(
|
||||
num_blocks=1153 * 3,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec),
|
||||
KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec),
|
||||
],
|
||||
)
|
||||
assert (
|
||||
get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_uniform_group
|
||||
)
|
||||
== 3
|
||||
)
|
||||
|
||||
# Scheduler-config shape: generate_scheduler_kv_cache_config replaces the
|
||||
# uniform-type group's spec with a representative per-layer spec.
|
||||
# Capacity must not change between the two shapes (the engine computes
|
||||
# on the scheduler config, the worker loop on the worker config).
|
||||
kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config(
|
||||
[copy.deepcopy(kv_cache_config_uniform_group)]
|
||||
)
|
||||
assert get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_scheduler_shape
|
||||
) == get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_uniform_group
|
||||
)
|
||||
|
||||
|
||||
def test_get_max_concurrency_packed_kv_cache_config():
|
||||
from vllm.v1.core.kv_cache_utils import (
|
||||
_get_kv_cache_config_packed,
|
||||
_use_packed_kv_cache_config,
|
||||
)
|
||||
|
||||
model_config = ModelConfig(
|
||||
"Qwen/Qwen1.5-7B",
|
||||
runner="generate",
|
||||
dtype="float16",
|
||||
max_model_len=16384,
|
||||
)
|
||||
scheduler_config = SchedulerConfig(
|
||||
max_num_batched_tokens=1024,
|
||||
enable_chunked_prefill=True,
|
||||
max_model_len=model_config.max_model_len,
|
||||
is_encoder_decoder=model_config.is_encoder_decoder,
|
||||
async_scheduling=False,
|
||||
)
|
||||
vllm_config = VllmConfig(
|
||||
model_config=model_config,
|
||||
scheduler_config=scheduler_config,
|
||||
)
|
||||
|
||||
# All-UniformTypeKVCacheSpecs groups select the packed layout.
|
||||
mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
|
||||
swa_specs = {
|
||||
f"layer_{i}": SlidingWindowMLASpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=576,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
)
|
||||
for i in range(4, 6)
|
||||
}
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(
|
||||
list(mla_specs),
|
||||
UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
|
||||
),
|
||||
KVCacheGroupSpec(
|
||||
list(swa_specs),
|
||||
UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
|
||||
),
|
||||
]
|
||||
assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
|
||||
num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
|
||||
vllm_config, kv_cache_groups, 2 * GiB_bytes
|
||||
)
|
||||
assert num_blocks > 0
|
||||
kv_cache_config_packed = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=kv_cache_tensors,
|
||||
kv_cache_groups=kv_cache_groups,
|
||||
)
|
||||
# Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
|
||||
# the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
|
||||
# previous formula normalized by the first group's page size and gave
|
||||
# 1061 blocks per request instead of 1097.
|
||||
assert get_max_concurrency_for_kv_cache_config(
|
||||
vllm_config, kv_cache_config_packed
|
||||
) == num_blocks / (1024 + 73)
|
||||
|
||||
|
||||
def test_allocate_with_lookahead():
|
||||
"""Verify that lookahead tokens correctly affect block allocation"""
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import dataclasses
|
||||
from concurrent.futures import Future
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
@@ -37,7 +36,7 @@ from vllm.v1.kv_cache_interface import (
|
||||
)
|
||||
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
from vllm.v1.structured_output import StructuredOutputGrammar, StructuredOutputManager
|
||||
from vllm.v1.structured_output import StructuredOutputManager
|
||||
|
||||
from .utils import EOS_TOKEN_ID, create_requests, create_scheduler, mock_kv
|
||||
|
||||
@@ -3007,58 +3006,6 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
|
||||
assert len(scheduler.skipped_waiting) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_grammar", [True, False])
|
||||
def test_grammar_compile_error_finishes_only_request(async_grammar: bool):
|
||||
scheduler = create_scheduler()
|
||||
manager = scheduler.structured_output_manager
|
||||
manager.backend = Mock()
|
||||
manager.backend.compile_grammar.side_effect = RuntimeError(
|
||||
"forced FSM compilation error"
|
||||
)
|
||||
manager._use_async_grammar_compilation = async_grammar
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
max_tokens=16,
|
||||
structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),
|
||||
)
|
||||
sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
|
||||
request = Request(
|
||||
request_id="grammar-error",
|
||||
prompt_token_ids=[0, 1],
|
||||
sampling_params=sampling_params,
|
||||
pooling_params=None,
|
||||
)
|
||||
|
||||
manager.grammar_init(request)
|
||||
assert request.structured_output_request is not None
|
||||
grammar_future = request.structured_output_request._grammar
|
||||
assert isinstance(grammar_future, Future)
|
||||
assert isinstance(grammar_future.exception(timeout=5), RuntimeError)
|
||||
|
||||
scheduler.add_request(request)
|
||||
scheduler_output = scheduler.schedule()
|
||||
assert not scheduler_output.num_scheduled_tokens
|
||||
|
||||
engine_core_outputs = scheduler.update_from_output(
|
||||
scheduler_output,
|
||||
ModelRunnerOutput(req_ids=[], req_id_to_index={}),
|
||||
)
|
||||
|
||||
assert request.status == RequestStatus.FINISHED_ERROR
|
||||
assert request.request_id not in scheduler.requests
|
||||
output = engine_core_outputs[0].outputs[0]
|
||||
assert output.request_id == request.request_id
|
||||
assert output.finish_reason == FinishReason.ERROR
|
||||
assert output.stop_reason is None
|
||||
|
||||
healthy_request = create_requests(num_requests=1, req_ids=["healthy-request"])[0]
|
||||
scheduler.add_request(healthy_request)
|
||||
next_output = scheduler.schedule()
|
||||
assert [req.req_id for req in next_output.scheduled_new_reqs] == [
|
||||
healthy_request.request_id
|
||||
]
|
||||
|
||||
|
||||
def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
scheduler = object.__new__(Scheduler)
|
||||
sampling_params = SamplingParams(ignore_eos=True, max_tokens=4)
|
||||
@@ -3072,7 +3019,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
pooling_params=None,
|
||||
)
|
||||
request.structured_output_request = Mock()
|
||||
request.structured_output_request.grammar = Mock(spec=StructuredOutputGrammar)
|
||||
request.structured_output_request.grammar = Mock()
|
||||
request.structured_output_request.grammar.accept_tokens.return_value = False
|
||||
request.status = RequestStatus.RUNNING
|
||||
request.num_computed_tokens = request.num_tokens
|
||||
@@ -3093,7 +3040,6 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
|
||||
scheduler.kv_event_publisher = Mock()
|
||||
scheduler.finished_req_ids = set()
|
||||
scheduler.finished_req_ids_dict = None
|
||||
scheduler.grammar_compile_error_reqs = set()
|
||||
scheduler.vllm_config = Mock()
|
||||
scheduler.vllm_config.model_config.enable_return_routed_experts = False
|
||||
scheduler.enable_return_routed_experts = False
|
||||
|
||||
@@ -7,7 +7,13 @@ import torch
|
||||
|
||||
from tests.v1.kv_connector.unit.utils import create_vllm_config
|
||||
from vllm.config import KVEventsConfig, KVTransferConfig
|
||||
from vllm.distributed.kv_events import MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored
|
||||
from vllm.distributed.kv_events import (
|
||||
MEDIUM_CPU,
|
||||
MEDIUM_FS,
|
||||
MEDIUM_OBJ,
|
||||
BlockRemoved,
|
||||
BlockStored,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
|
||||
build_offloading_config,
|
||||
)
|
||||
@@ -62,8 +68,9 @@ def _wire_hash(block_hash: BlockHash):
|
||||
return maybe_convert_block_hash(block_hash)
|
||||
|
||||
|
||||
def _request(*, block_hashes: list[BlockHash], token_count: int):
|
||||
def _request(*, block_hashes: list[BlockHash], token_count: int, req_id: str = "req"):
|
||||
req = MagicMock()
|
||||
req.request_id = req_id
|
||||
req.block_hashes = block_hashes
|
||||
req.all_token_ids = list(range(1, token_count + 1))
|
||||
req.lora_request = None
|
||||
@@ -104,10 +111,32 @@ def _record_chunks(
|
||||
return keys
|
||||
|
||||
|
||||
def _record_lookup_chunks(
|
||||
tracker: OffloadingEventsTracker,
|
||||
req,
|
||||
group_config: GroupOffloadConfig,
|
||||
num_chunks: int,
|
||||
) -> list[OffloadKey]:
|
||||
keys: list[OffloadKey] = []
|
||||
hbf = group_config.hashes_per_chunk
|
||||
for chunk_idx in range(num_chunks):
|
||||
tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
|
||||
assert tail_hash is not None
|
||||
key = make_offload_key(tail_hash, group_config.group_idx)
|
||||
tracker.record_lookup(
|
||||
req,
|
||||
group_config,
|
||||
chunk_idx,
|
||||
key,
|
||||
)
|
||||
keys.append(key)
|
||||
return keys
|
||||
|
||||
|
||||
def _stored_event(
|
||||
keys: list[OffloadKey],
|
||||
locality: Locality | None = None,
|
||||
medium: str = _CPU_MEDIUM,
|
||||
locality: Locality | None = None,
|
||||
) -> OffloadingEvent:
|
||||
return OffloadingEvent(
|
||||
keys=keys,
|
||||
@@ -119,8 +148,8 @@ def _stored_event(
|
||||
|
||||
def _removed_event(
|
||||
keys: list[OffloadKey],
|
||||
locality: Locality | None = None,
|
||||
medium: str = _CPU_MEDIUM,
|
||||
locality: Locality | None = None,
|
||||
) -> OffloadingEvent:
|
||||
return OffloadingEvent(
|
||||
keys=keys,
|
||||
@@ -130,6 +159,21 @@ def _removed_event(
|
||||
)
|
||||
|
||||
|
||||
def _lookup_chunk() -> tuple[
|
||||
OffloadingEventsTracker, MagicMock, GroupOffloadConfig, OffloadKey
|
||||
]:
|
||||
tracker = _tracker()
|
||||
req = _request(block_hashes=[_hash(0)], token_count=4)
|
||||
group_config = _group_config()
|
||||
key = _record_lookup_chunks(
|
||||
tracker,
|
||||
req,
|
||||
group_config,
|
||||
num_chunks=1,
|
||||
)[0]
|
||||
return tracker, req, group_config, key
|
||||
|
||||
|
||||
def test_take_events_forwards_locality_to_rich_store():
|
||||
tracker = _tracker()
|
||||
req = _request(block_hashes=[_hash(0)], token_count=4)
|
||||
@@ -220,18 +264,37 @@ def test_take_events_publishes_routable_block_stored():
|
||||
assert len(tracker._pending_event_metadata) == 6
|
||||
|
||||
|
||||
def test_take_events_factor_gt_1_chunk_store_and_remove():
|
||||
def test_promotion_emits_full_cpu_stored_event():
|
||||
tracker, _, _, key = _lookup_chunk()
|
||||
|
||||
[event] = tracker.take_events([_stored_event([key])])
|
||||
|
||||
assert isinstance(event, BlockStored)
|
||||
assert event.medium == MEDIUM_CPU
|
||||
assert event.block_hashes == [_wire_hash(_hash(0))]
|
||||
assert event.parent_block_hash is None
|
||||
assert event.token_ids == [1, 2, 3, 4]
|
||||
assert event.block_size == 4
|
||||
assert event.lora_id is None
|
||||
assert event.lora_name is None
|
||||
assert event.extra_keys is None
|
||||
assert event.group_idx == 0
|
||||
assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
|
||||
assert event.kv_cache_spec_sliding_window is None
|
||||
|
||||
|
||||
def test_lookup_promotion_factor_gt_1_store_and_remove():
|
||||
block_size = 4
|
||||
blocks_per_chunk = 3
|
||||
blocks_per_chunk = 2
|
||||
tracker = _tracker()
|
||||
group_config = _group_config(
|
||||
block_size=block_size, blocks_per_chunk=blocks_per_chunk
|
||||
)
|
||||
req = _request(
|
||||
block_hashes=[_hash(i) for i in range(6)],
|
||||
block_hashes=[_hash(i) for i in range(4)],
|
||||
token_count=block_size * blocks_per_chunk * 2,
|
||||
)
|
||||
keys = _record_chunks(tracker, req, group_config, num_chunks=2)
|
||||
keys = _record_lookup_chunks(tracker, req, group_config, num_chunks=2)
|
||||
|
||||
stored = list(tracker.take_events([_stored_event(keys)]))
|
||||
assert len(stored) == 2
|
||||
@@ -293,6 +356,7 @@ def test_take_events_opt_out_keeps_placeholders():
|
||||
group_config = _group_config()
|
||||
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
|
||||
keys = _record_chunks(tracker, req, group_config, num_chunks=3)
|
||||
_record_lookup_chunks(tracker, req, group_config, num_chunks=3)
|
||||
|
||||
assert not tracker.self_describing_enabled
|
||||
assert not tracker._pending_event_metadata
|
||||
@@ -315,11 +379,21 @@ def test_take_events_opt_out_keeps_placeholders():
|
||||
assert len(events[3].block_hashes) == 3
|
||||
|
||||
|
||||
def test_record_store_skips_sliding_window_group():
|
||||
@pytest.mark.parametrize(
|
||||
"sliding_window_size_in_chunks",
|
||||
[1, 2],
|
||||
ids=["ssm", "sliding-window"],
|
||||
)
|
||||
def test_event_metadata_skips_non_full_attention_group(
|
||||
sliding_window_size_in_chunks: int,
|
||||
):
|
||||
tracker = _tracker()
|
||||
group_config = _group_config(sliding_window_size_in_chunks=2)
|
||||
group_config = _group_config(
|
||||
sliding_window_size_in_chunks=sliding_window_size_in_chunks
|
||||
)
|
||||
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
|
||||
keys = _record_chunks(tracker, req, group_config, num_chunks=3)
|
||||
_record_lookup_chunks(tracker, req, group_config, num_chunks=3)
|
||||
|
||||
assert not tracker._pending_event_metadata
|
||||
|
||||
@@ -329,6 +403,57 @@ def test_record_store_skips_sliding_window_group():
|
||||
assert events[0].block_size == 0
|
||||
|
||||
|
||||
def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
|
||||
tracker = _tracker()
|
||||
block_hashes = [_hash(0), _hash(1)]
|
||||
req = _request(block_hashes=block_hashes, token_count=8)
|
||||
group_config = _group_config(blocks_per_chunk=2)
|
||||
key = _record_chunks(tracker, req, group_config, num_chunks=1)[0]
|
||||
confirmed_meta = tracker._pending_event_metadata[key]
|
||||
lookup_req = _request(
|
||||
block_hashes=block_hashes,
|
||||
token_count=8,
|
||||
req_id="new-request",
|
||||
)
|
||||
|
||||
tracker.record_lookup(
|
||||
lookup_req,
|
||||
group_config,
|
||||
0,
|
||||
key,
|
||||
)
|
||||
assert tracker._pending_event_metadata[key] is confirmed_meta
|
||||
|
||||
removed = list(tracker.take_events([_removed_event([key])]))
|
||||
assert len(removed) == 1
|
||||
assert removed[0].block_hashes == [
|
||||
_wire_hash(_hash(0)),
|
||||
_wire_hash(_hash(1)),
|
||||
]
|
||||
|
||||
stored = list(tracker.take_events([_stored_event([key])]))
|
||||
assert len(stored) == 1
|
||||
assert stored[0].block_size == 0
|
||||
assert stored[0].token_ids == []
|
||||
|
||||
tracker.record_lookup(lookup_req, group_config, 0, key)
|
||||
removed = list(tracker.take_events([_removed_event([key])]))
|
||||
assert removed[0].block_hashes == [
|
||||
_wire_hash(_hash(0)),
|
||||
_wire_hash(_hash(1)),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("medium", [MEDIUM_FS, MEDIUM_OBJ])
|
||||
def test_secondary_stored_event_does_not_mutate_cpu_metadata(medium: str):
|
||||
tracker, _, _, key = _lookup_chunk()
|
||||
expected_metadata = dict(tracker._pending_event_metadata)
|
||||
|
||||
stored = list(tracker.take_events([_stored_event([key], medium)]))
|
||||
assert stored[0].token_ids == [1, 2, 3, 4]
|
||||
assert tracker._pending_event_metadata == expected_metadata
|
||||
|
||||
|
||||
def test_take_events_groups_removed_hashes_by_kv_group():
|
||||
tracker = _tracker()
|
||||
group0_config = _group_config(group_idx=0, blocks_per_chunk=2)
|
||||
@@ -378,7 +503,7 @@ def test_reset_cache_clears_side_table():
|
||||
tracker = _tracker()
|
||||
group_config = _group_config()
|
||||
req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
|
||||
_record_chunks(tracker, req, group_config, num_chunks=3)
|
||||
_record_lookup_chunks(tracker, req, group_config, num_chunks=3)
|
||||
|
||||
assert tracker._pending_event_metadata
|
||||
|
||||
@@ -387,7 +512,7 @@ def test_reset_cache_clears_side_table():
|
||||
assert not tracker._pending_event_metadata
|
||||
|
||||
|
||||
def test_tiering_rejects_self_describing_kv_events():
|
||||
def test_tiering_accepts_self_describing_kv_events():
|
||||
vllm_config = create_vllm_config(
|
||||
block_size=4,
|
||||
max_num_batched_tokens=16,
|
||||
@@ -423,5 +548,9 @@ def test_tiering_rejects_self_describing_kv_events():
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="TieringOffloadingSpec"):
|
||||
TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
|
||||
spec = TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
|
||||
tracker = OffloadingEventsTracker(spec.kv_events_config)
|
||||
|
||||
assert spec.kv_events_config.enable_kv_cache_events
|
||||
assert spec.kv_events_config.self_describing_kv_events
|
||||
assert tracker.self_describing_enabled
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user