diff --git a/.buildkite/ci_config_intel.yaml b/.buildkite/ci_config_intel.yaml index a1c0091e0f1..95fff1940fa 100644 --- a/.buildkite/ci_config_intel.yaml +++ b/.buildkite/ci_config_intel.yaml @@ -2,17 +2,16 @@ name: vllm_intel_ci job_dirs: - ".buildkite/intel_jobs" run_all_patterns: + - ".buildkite/ci_config_intel.yaml" - "docker/Dockerfile" + - "docker/Dockerfile.xpu" - "CMakeLists.txt" - "requirements/common.txt" - "requirements/xpu.txt" - - "requirements/build/cuda.txt" - - "requirements/test/cuda.txt" - "setup.py" - "csrc/" - "cmake/" run_all_exclude_patterns: - - "docker/Dockerfile." - "csrc/cpu/" - "csrc/rocm/" - "cmake/hipify.py" diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index 911b6c45e0e..dd85400f2f1 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -53,7 +53,7 @@ steps: - tests/models/language/pooling/ commands: - | - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 40m " + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 50m " pytest -x -v -s tests/models/language/generation -m cpu_model pytest -x -v -s tests/models/language/pooling -m cpu_model" @@ -68,13 +68,15 @@ steps: - vllm/v1/sample/ops/topk_topp_triton.py - vllm/v1/sample/ops/topk_topp_sampler.py - tests/v1/sample/test_topk_topp_sampler.py + - tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py commands: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m " uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel - pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp" + pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp + pytest -x -v -s tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py" - label: CPU-Quantization Model Tests depends_on: [] diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test.sh b/.buildkite/scripts/hardware_ci/run-cpu-test.sh index 27ec0068668..0f0c18b55af 100644 --- a/.buildkite/scripts/hardware_ci/run-cpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test.sh @@ -7,7 +7,8 @@ set -euox pipefail # allow to bind to different cores CORE_RANGE=${CORE_RANGE:-48-95} NUMA_NODE=${NUMA_NODE:-1} -IMAGE_NAME="cpu-test-$NUMA_NODE" +AGENT_SLOT=${AGENT_SLOT:-} +IMAGE_NAME="cpu-test-${NUMA_NODE}${AGENT_SLOT:+-${AGENT_SLOT}}" TIMEOUT_VAL=$1 TEST_COMMAND=$2 diff --git a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh index 0eadfa1f80b..e1808835fdf 100755 --- a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh +++ b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh @@ -33,6 +33,14 @@ if [[ -n "${ATTENTION_BACKEND:-}" ]]; then EXTRA_ARGS+=(--attention-backend "${ATTENTION_BACKEND}") fi +# ROCm: run eager to avoid intermittent HIP-graph decode corruption. +# See https://github.com/ROCm/clr/issues/279 +# TODO(aarushjain29): Revert after TheRock 7.14 +if command -v rocm-smi &> /dev/null || command -v amd-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then + echo "ROCm platform detected: adding --enforce-eager to avoid HIP-graph decode corruption" + EXTRA_ARGS+=(--enforce-eager) +fi + cleanup() { if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then kill "${SERVER_PID}" 2>/dev/null || true diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 7e48a125071..eeb685e9892 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -132,7 +132,7 @@ steps: - vllm/_aiter_ops.py - vllm/platforms/rocm.py commands: - - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - TARGET_TEST_SUITE=MI250 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' @@ -221,30 +221,6 @@ steps: - pytest -v -s distributed/test_shm_buffer.py - pytest -v -s distributed/test_shm_storage.py -- label: Distributed Torchrun + Shutdown Tests (2 GPUs) # TBD - timeout_in_minutes: 180 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - agent_pool: mi250_2 - num_gpus: 2 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/ - - vllm/engine/ - - vllm/executor/ - - vllm/worker/worker_base.py - - vllm/v1/engine/ - - vllm/v1/worker/ - - tests/distributed/ - - tests/v1/shutdown - - tests/v1/worker/test_worker_memory_snapshot.py - - vllm/platforms/rocm.py - commands: - - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed' - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown - - pytest -v -s v1/worker/test_worker_memory_snapshot.py - - label: Pipeline + Context Parallelism (4 GPUs) # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] @@ -617,7 +593,7 @@ steps: - pytest -v -s -m 'cpu_test' multimodal - pytest -v -s renderers - pytest -v -s tokenizers_ - - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py + - pytest -v -s reasoning - pytest -v -s tool_parsers - pytest -v -s parser - pytest -v -s transformers_utils @@ -646,7 +622,7 @@ steps: commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -v -s basic_correctness/test_mem.py - - pytest -v -s basic_correctness/test_basic_correctness.py + - VLLM_TARGET_TEST_SUITE=MI300 pytest -v -s basic_correctness/test_basic_correctness.py - pytest -v -s basic_correctness/test_cpu_offload.py - label: Distributed Model Tests (2 GPUs) # TBD @@ -668,7 +644,7 @@ steps: - tests/model_executor/model_loader/test_sharded_state_loader.py - tests/models/ commands: - - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)' - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)' - pytest models/language -v -s -m 'distributed(num_gpus=2)' @@ -887,7 +863,7 @@ steps: commands: - pytest -v -s distributed/test_custom_all_reduce.py - torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py - - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' + - TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)' - pytest -v -s -x lora/test_mixtral.py - label: Distributed Torchrun + Examples (4 GPUs) # TBD @@ -1588,10 +1564,10 @@ steps: - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset - label: Basic Models Tests (Other) # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_1 - torch_nightly: true + optional: true working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/ @@ -2240,7 +2216,7 @@ steps: - pytest -v -s tests/distributed/test_packed_tensor.py - label: Metrics, Tracing (2 GPUs) # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 20 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_2 optional: true @@ -2393,7 +2369,7 @@ steps: #------------------------------------------------------ mi300 · weight_loading -------------------------------------------------------# - label: Weight Loading Multiple GPU # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_2 num_gpus: 2 @@ -2405,7 +2381,7 @@ steps: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - label: Weight Loading Multiple GPU - Large Models # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] agent_pool: mi300_2 num_gpus: 2 @@ -3318,7 +3294,7 @@ steps: #------------------------------------------------------ mi355 · weight_loading -------------------------------------------------------# - label: Weight Loading Multiple GPU # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 num_gpus: 2 @@ -3330,7 +3306,7 @@ steps: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt - label: Weight Loading Multiple GPU - Large Models # TBD - timeout_in_minutes: 180 + timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] agent_pool: mi355_2 working_dir: "/vllm-workspace/tests" diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 57851edb0b8..f5db2e956b6 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -252,9 +252,10 @@ steps: - pytest -v -s v1/tracing mirror: amd: - device: mi300_2 + device: mi325_2 depends_on: - image-build-amd + optional: true - label: Python-only Installation key: python-only-installation @@ -359,7 +360,7 @@ steps: - pytest -v -s test_ray_env.py - pytest -v -s -m 'cpu_test' multimodal - pytest -v -s renderers - - pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py + - pytest -v -s reasoning - pytest -v -s tool_parsers - pytest -v -s tokenizers_ - pytest -v -s parser diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 4e47cbb7794..5eb799efa18 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -45,6 +45,11 @@ steps: - tests/models/test_registry.py commands: - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py + mirror: + amd: + device: mi325_1 + depends_on: + - image-build-amd - label: Basic Models Test (Other CPU) # 5min key: basic-models-test-other-cpu diff --git a/.buildkite/test_areas/weight_loading.yaml b/.buildkite/test_areas/weight_loading.yaml index 01c6bb7809b..9d7bd0bce91 100644 --- a/.buildkite/test_areas/weight_loading.yaml +++ b/.buildkite/test_areas/weight_loading.yaml @@ -13,6 +13,13 @@ steps: - tests/weight_loading commands: - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models.txt + mirror: + amd: + device: mi300_2 + depends_on: + - image-build-amd + commands: + - bash weight_loading/run_model_weight_loading_test.sh -c weight_loading/models-amd.txt # - label: Weight Loading Multiple GPU - Large Models # optional # working_dir: "/vllm-workspace/tests" diff --git a/.github/workflows/macos-smoke-test.yml b/.github/workflows/macos-smoke-test.yml index ea1c8b0feac..9068ec281b2 100644 --- a/.github/workflows/macos-smoke-test.yml +++ b/.github/workflows/macos-smoke-test.yml @@ -15,9 +15,9 @@ jobs: timeout-minutes: 30 steps: - - uses: actions/checkout@v6.0.1 + - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1 - - uses: astral-sh/setup-uv@v7 + - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0 with: enable-cache: true cache-dependency-glob: | diff --git a/CMakeLists.txt b/CMakeLists.txt index 6d130f8bda2..36b8e66f2c6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -869,9 +869,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) - cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}") else() - cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}") + cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}") endif() if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS) set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu") diff --git a/benchmarks/kernels/benchmark_fused_collective.py b/benchmarks/kernels/benchmark_fused_collective.py index 36cbd715f18..c999c16021b 100644 --- a/benchmarks/kernels/benchmark_fused_collective.py +++ b/benchmarks/kernels/benchmark_fused_collective.py @@ -80,13 +80,17 @@ _FI_MAX_SIZES = { 2: 64 * MiB, # 64MB 4: 64 * MiB, # 64MB 8: 64 * MiB, # 64MB + 16: 64 * MiB, # 64MB (multi-node) } # Global workspace tensors for FlashInfer (keyed by backend name) _FI_WORKSPACES: dict = {} -# Backends to benchmark -FLASHINFER_BACKENDS = ["trtllm", "mnnvl"] +# Backends to benchmark. trtllm is single-node only and can hang cross-node, so +# multi-node sweeps can restrict to mnnvl via FI_BACKENDS=mnnvl. +FLASHINFER_BACKENDS = [ + b for b in os.environ.get("FI_BACKENDS", "trtllm,mnnvl").split(",") if b +] def setup_flashinfer_workspace( @@ -995,7 +999,10 @@ def main(): rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) - device = torch.device(f"cuda:{rank}") + # Use LOCAL_RANK for the device so multi-node runs (global rank >= GPUs per + # node) map to a valid local GPU; falls back to global rank single-node. + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + device = torch.device(f"cuda:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 5d0876f9125..1531cc96920 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -391,16 +391,19 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int config = dict(zip(keys, config_values)) configs.append(config) - # Remove configs that are not compatible with fp8 block quantization - # BLOCK_SIZE_K must be a multiple of block_k - # BLOCK_SIZE_N must be a multiple of block_n + # Drop configs incompatible with fp8 block quantization. A tile must align + # to the quant-block scale grid, i.e. tile and block must divide one + # another. The kernel indexes scales per element (offs_bn // group_n, + # k_start // group_k), so a tile narrower than the block (e.g. N=64 with + # block_n=128) is valid -- and often faster at small batch. An exact + # multiple was required before, which dropped those smaller tiles entirely. if block_quant_shape is not None and not use_fp16: block_n, block_k = block_quant_shape[0], block_quant_shape[1] for config in configs[:]: - if ( - config["BLOCK_SIZE_K"] % block_k != 0 - or config["BLOCK_SIZE_N"] % block_n != 0 - ): + bn, bk = config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"] + n_aligned = bn % block_n == 0 or block_n % bn == 0 + k_aligned = bk % block_k == 0 or block_k % bk == 0 + if not (n_aligned and k_aligned): configs.remove(config) return configs diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 386f9e30c77..5c19446601e 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -166,12 +166,13 @@ elseif (S390_FOUND) "-mtune=native") elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64") message(STATUS "RISC-V detected") - if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0) + if(DEFINED VLLM_RVV_VLEN AND VLLM_RVV_VLEN LESS 0) message(FATAL_ERROR - "VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'") + "VLLM_RVV_VLEN must be zero or a positive integer; got '${VLLM_RVV_VLEN}'") endif() # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo - # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256. + # by default; set -DVLLM_RVV_VLEN=0 to force scalar RISC-V build. + # Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV. if(NOT DEFINED VLLM_RVV_VLEN) # Auto-detect: find the largest zvlb in /proc/cpuinfo isa line. if(EXISTS /proc/cpuinfo) diff --git a/csrc/cpu/cpu_types_riscv_defs.hpp b/csrc/cpu/cpu_types_riscv_defs.hpp index 650dc5bcc79..16475505d9f 100644 --- a/csrc/cpu/cpu_types_riscv_defs.hpp +++ b/csrc/cpu/cpu_types_riscv_defs.hpp @@ -3,13 +3,17 @@ // VLEN-to-LMUL mapping for RISC-V Vector extension. // -// LMUL_ expands to the LMUL suffix giving N total bits of vector data: -// VLEN=128: LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8 -// VLEN=256: LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4 +// LMUL_ expands to the LMUL suffix giving N total bits of vector data. +// LMUL_64 is used by 8-lane int8/uint8 vectors. +// VLEN=128: +// LMUL_64=mf2, LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8 +// VLEN=256: +// LMUL_64=mf4, LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4 #include #if __riscv_v_min_vlen == 128 + #define LMUL_64 mf2 #define LMUL_128 m1 #define LMUL_256 m2 #define LMUL_512 m4 @@ -17,6 +21,7 @@ #define BOOL_256 b16 #define BOOL_512 b8 #elif __riscv_v_min_vlen == 256 + #define LMUL_64 mf4 #define LMUL_128 mf2 #define LMUL_256 m1 #define LMUL_512 m2 @@ -41,6 +46,16 @@ // ---- Semantic fixed-vector typedefs (named by element count) ---- +// uint8 / int8 +typedef RVVTYPE(vuint8, LMUL_64, _t) fixed_u8x8_t + __attribute__((riscv_rvv_vector_bits(64))); +typedef RVVTYPE(vint8, LMUL_64, _t) fixed_i8x8_t + __attribute__((riscv_rvv_vector_bits(64))); + +// int16 +typedef RVVTYPE(vint16, LMUL_128, _t) fixed_i16x8_t + __attribute__((riscv_rvv_vector_bits(128))); + // float16 typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t __attribute__((riscv_rvv_vector_bits(128))); diff --git a/csrc/cpu/sgl-kernels/conv.cpp b/csrc/cpu/sgl-kernels/conv.cpp index 15114732aac..b918aed8bff 100644 --- a/csrc/cpu/sgl-kernels/conv.cpp +++ b/csrc/cpu/sgl-kernels/conv.cpp @@ -289,19 +289,18 @@ void causal_conv1d_fwd_kernel_impl( } } -#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \ - tinygemm_kernel::apply( \ - input + batch_offset * dim + mb_start * dim + nb_start, \ - weight + nb_start * width, \ - out + batch_offset * dim + mb_start * dim + nb_start, \ - has_bias ? bias + nb_start : nullptr, \ - nullptr, \ - false, \ - mb_size, \ - dim, \ +#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \ + tinygemm_kernel::apply( \ + input + batch_offset * dim + mb_start * dim + nb_start, \ + weight + nb_start * width, \ + out + batch_offset * dim + mb_start * dim + nb_start, \ + has_bias ? bias + nb_start : nullptr, \ + has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \ + has_initial_states_value, \ + mb_size, \ + dim, \ mb_start == 0); -// TODO: add `has_initial_state` support for varlen kernel template void causal_conv1d_fwd_varlen_kernel_impl( scalar_t* __restrict__ out, @@ -343,6 +342,9 @@ void causal_conv1d_fwd_varlen_kernel_impl( int64_t nb_start = nb * BLOCK_N; int64_t nb_size = std::min(dim - nb_start, BLOCK_N); + const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false; + int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs; + switch (width << 4 | nb_size >> 4) { case 0x42: LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32); @@ -373,7 +375,7 @@ void causal_conv1d_fwd_varlen_kernel_impl( width, dim, seqlen, - /* has_initial_state */ false); + has_initial_state[bs]); } }); } diff --git a/csrc/cpu/sgl-kernels/gemm_int4.cpp b/csrc/cpu/sgl-kernels/gemm_int4.cpp index 1fec14c956f..6dbd09080d0 100644 --- a/csrc/cpu/sgl-kernels/gemm_int4.cpp +++ b/csrc/cpu/sgl-kernels/gemm_int4.cpp @@ -285,6 +285,125 @@ inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f); } +#if defined(CPU_CAPABILITY_RVV) +template +inline fixed_i8x8_t load_uint4_as_int8_rvv(const uint8_t* __restrict__ B, int64_t k) { + constexpr int64_t n_group_size = 8; + constexpr int64_t vnni_size = 4; + static_assert(N == 32); + static_assert(ldb == N / 2); + static_assert(group >= 0 && group < N / n_group_size); + + // Unpack: gather 8 packed int4 values from the VNNI4 layout. + const int64_t ki = k % vnni_size; + const int64_t k_base = k - ki; + constexpr int64_t packed_group = group / 2; + const uint8_t* packed_ptr = B + k_base * ldb + packed_group * n_group_size * vnni_size + ki; + + fixed_u8x8_t packed = RVVI(__riscv_vlse8_v_u8, LMUL_64)(packed_ptr, vnni_size, n_group_size); + if constexpr (group % 2 == 1) { + packed = RVVI(__riscv_vsrl_vx_u8, LMUL_64)(packed, 4, n_group_size); + } + fixed_u8x8_t nibbles = RVVI(__riscv_vand_vx_u8, LMUL_64)(packed, 0x0f, n_group_size); + return RVVI4(__riscv_vreinterpret_v_u8, LMUL_64, _i8, LMUL_64)(nibbles); +} + +inline fixed_i32x8_t gemm_accum_uint8_int8_rvv(fixed_i32x8_t acc, uint8_t a, fixed_i8x8_t b) { + constexpr int64_t vl = 8; + fixed_i16x8_t b_i16 = RVVI(__riscv_vsext_vf2_i16, LMUL_128)(b, vl); + return RVVI(__riscv_vwmacc_vx_i32, LMUL_256)(acc, static_cast(a), b_i16, vl); +} + +template +inline fixed_i32x8_t gemm_accum_uint4_rvv( + fixed_i32x8_t acc, + const uint8_t* __restrict__ B, + const int8_t* __restrict__ qzeros_b, + uint8_t a, + int64_t k) { + constexpr int64_t n_group_size = 8; + fixed_i8x8_t b = load_uint4_as_int8_rvv(B, k); + fixed_i8x8_t qzeros = + RVVI(__riscv_vle8_v_i8, LMUL_64)(qzeros_b + group * n_group_size, n_group_size); + b = RVVI(__riscv_vsub_vv_i8, LMUL_64)(b, qzeros, n_group_size); + return gemm_accum_uint8_int8_rvv(acc, a, b); +} + +template +inline void _dequant_and_store_rvv( + float* __restrict__ C, + fixed_i32x8_t acc, + const float* __restrict__ scales_a, + const int32_t* __restrict__ qzeros_a, + const float* __restrict__ scales_b, + const int32_t* __restrict__ compensation, + int64_t m, + int64_t ldc) { + constexpr int64_t n_group_size = 8; + constexpr int64_t n = group * n_group_size; + constexpr int64_t vl = n_group_size; + + // Dequant compensation: remove activation zero-point contribution. + fixed_i32x8_t comp = RVVI(__riscv_vle32_v_i32, LMUL_256)(compensation + n, vl); + fixed_i32x8_t zp_comp = RVVI(__riscv_vmul_vx_i32, LMUL_256)(comp, qzeros_a[m], vl); + acc = RVVI(__riscv_vsub_vv_i32, LMUL_256)(acc, zp_comp, vl); + + // Scale: convert int32 accumulators to fp32 and apply activation/weight scales. + fixed_fp32x8_t acc_f = RVVI(__riscv_vfcvt_f_x_v_f32, LMUL_256)(acc, vl); + acc_f = RVVI(__riscv_vfmul_vf_f32, LMUL_256)(acc_f, scales_a[m], vl); + fixed_fp32x8_t scale_b = RVVI(__riscv_vle32_v_f32, LMUL_256)(scales_b + n, vl); + acc_f = RVVI(__riscv_vfmul_vv_f32, LMUL_256)(acc_f, scale_b, vl); + + // Store: accumulate into the float scratch buffer that already holds bias/zero. + float* c_ptr = C + m * ldc + n; + fixed_fp32x8_t c_old = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr, vl); + fixed_fp32x8_t c_new = RVVI(__riscv_vfadd_vv_f32, LMUL_256)(c_old, acc_f, vl); + RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr, c_new, vl); +} + +template +void _dequant_gemm_accum_rvv( + float* __restrict__ C, + const uint8_t* __restrict__ A, + const float* __restrict__ scales_a, + const int32_t* __restrict__ qzeros_a, + const uint8_t* __restrict__ B, + const float* __restrict__ scales_b, + const int8_t* __restrict__ qzeros_b, + const int32_t* __restrict__ compensation, + int64_t M, + int64_t K, + int64_t lda, + int64_t ldc) { + static_assert(N == 32); + static_assert(ldb == N / 2); + constexpr int64_t vl = 8; + + // Accumulate one C row over the 32-column block. + for (int64_t m = 0; m < M; ++m) { + fixed_i32x8_t acc0 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl); + fixed_i32x8_t acc1 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl); + fixed_i32x8_t acc2 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl); + fixed_i32x8_t acc3 = RVVI(__riscv_vmv_v_x_i32, LMUL_256)(0, vl); + // A[m][k] @ B[k][0:32] -> acc[m][0:32] + for (int64_t k = 0; k < K; ++k) { + // GEMM K step: one scalar activation updates four 8-column RVV tiles. + const uint8_t a = A[m * lda + k]; + acc0 = gemm_accum_uint4_rvv(acc0, B, qzeros_b, a, k); + acc1 = gemm_accum_uint4_rvv(acc1, B, qzeros_b, a, k); + acc2 = gemm_accum_uint4_rvv(acc2, B, qzeros_b, a, k); + acc3 = gemm_accum_uint4_rvv(acc3, B, qzeros_b, a, k); + } + + // Dequant/scale/store each 8-column group back into C. + _dequant_and_store_rvv<0>(C, acc0, scales_a, qzeros_a, scales_b, compensation, m, ldc); + _dequant_and_store_rvv<1>(C, acc1, scales_a, qzeros_a, scales_b, compensation, m, ldc); + _dequant_and_store_rvv<2>(C, acc2, scales_a, qzeros_a, scales_b, compensation, m, ldc); + _dequant_and_store_rvv<3>(C, acc3, scales_a, qzeros_a, scales_b, compensation, m, ldc); + } +} +#endif + template void _dequant_gemm_accum( float* C, @@ -336,6 +455,11 @@ void _dequant_gemm_accum( _dequant_and_store( C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/); } else +#elif defined(CPU_CAPABILITY_RVV) + if constexpr (!sym_quant_act && N == BLOCK_N && ldb == BLOCK_N / 2) { + _dequant_gemm_accum_rvv(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, M, K, lda, ldc); + return; + } else #endif { for (int64_t m = 0; m < M; ++m) { diff --git a/csrc/cpu/sgl-kernels/vec.h b/csrc/cpu/sgl-kernels/vec.h index 72143fedc69..407cfe60434 100644 --- a/csrc/cpu/sgl-kernels/vec.h +++ b/csrc/cpu/sgl-kernels/vec.h @@ -9,11 +9,19 @@ #define CPU_CAPABILITY_AVX512 #endif +#if defined(__riscv_v_min_vlen) && (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256) +#define CPU_CAPABILITY_RVV +#endif + #include #include #if defined(CPU_CAPABILITY_AVX512) #include #endif + +#if defined(CPU_CAPABILITY_RVV) +#include "../cpu_types_riscv_defs.hpp" +#endif namespace { using namespace at::vec; diff --git a/csrc/cpu/spec_decode_utils.cpp b/csrc/cpu/spec_decode_utils.cpp index a76b8bc6937..30192196b95 100644 --- a/csrc/cpu/spec_decode_utils.cpp +++ b/csrc/cpu/spec_decode_utils.cpp @@ -208,6 +208,89 @@ void copy_and_expand_eagle_inputs_kernel_impl( } } +void copy_and_expand_dflash_inputs_kernel_impl( + const torch::Tensor& next_token_ids, const torch::Tensor& target_positions, + torch::Tensor& out_input_ids, torch::Tensor& out_context_positions, + torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping, + torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices, + const torch::Tensor& block_table, const torch::Tensor& query_start_loc, + const std::optional& num_rejected_tokens, + const int64_t parallel_drafting_token_id, const int64_t block_size, + const int64_t num_query_per_req, const int64_t num_speculative_tokens, + const int64_t total_input_tokens, const bool has_num_rejected) { + const int64_t num_reqs = query_start_loc.size(0) - 1; + + const int64_t* next_ids_ptr = next_token_ids.data_ptr(); + const int64_t* target_pos_ptr = target_positions.data_ptr(); + const int32_t* block_table_ptr = block_table.data_ptr(); + const int32_t* query_start_ptr = query_start_loc.data_ptr(); + const int64_t* rejected_ptr = + has_num_rejected && num_rejected_tokens.has_value() + ? num_rejected_tokens.value().data_ptr() + : nullptr; + + int64_t* out_ids_ptr = out_input_ids.data_ptr(); + int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr(); + int64_t* out_query_pos_ptr = out_query_positions.data_ptr(); + int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr(); + int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr(); + int32_t* out_token_idx_ptr = out_token_indices.data_ptr(); + + const int64_t block_table_stride = block_table.stride(0); + +#pragma omp parallel for + for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { + int32_t ctx_start = query_start_ptr[req_idx]; + int32_t ctx_end = query_start_ptr[req_idx + 1]; + int64_t num_ctx = ctx_end - ctx_start; + int64_t valid_ctx_end = ctx_end; + if (rejected_ptr != nullptr) { + valid_ctx_end -= rejected_ptr[req_idx]; + } + // Guard against out-of-bounds: ensure valid_ctx_end > ctx_start so that + // valid_ctx_end - 1 never reads before the request's context range. + valid_ctx_end = + std::max(valid_ctx_end, static_cast(ctx_start + 1)); + + int64_t last_pos = target_pos_ptr[valid_ctx_end - 1]; + + for (int64_t j = 0; j < num_ctx; ++j) { + int64_t ctx_idx = ctx_start + j; + int64_t ctx_pos_idx = std::min(ctx_idx, total_input_tokens - 1); + int64_t position = target_pos_ptr[ctx_pos_idx]; + int64_t block_num = position / block_size; + block_num = std::min(block_num, block_table_stride - 1); + int32_t block_id = + block_table_ptr[req_idx * block_table_stride + block_num]; + int64_t slot = block_id * block_size + (position % block_size); + + out_ctx_pos_ptr[ctx_idx] = position; + out_ctx_slot_ptr[ctx_idx] = slot; + } + + for (int64_t query_off = 0; query_off < num_query_per_req; ++query_off) { + int64_t query_out = req_idx * num_query_per_req + query_off; + int64_t position = last_pos + 1 + query_off; + int64_t block_num = position / block_size; + block_num = std::min(block_num, block_table_stride - 1); + int32_t block_id = + block_table_ptr[req_idx * block_table_stride + block_num]; + int64_t slot = block_id * block_size + (position % block_size); + + out_query_pos_ptr[query_out] = position; + out_query_slot_ptr[query_out] = slot; + out_ids_ptr[query_out] = + query_off == 0 ? next_ids_ptr[req_idx] : parallel_drafting_token_id; + + if (query_off > 0) { + int64_t sample_out_idx = + req_idx * num_speculative_tokens + (query_off - 1); + out_token_idx_ptr[sample_out_idx] = query_out; + } + } + } +} + void rejection_greedy_sample_kernel_impl( torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax, diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 9cef2d0d535..bc02511eb80 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -237,6 +237,16 @@ void copy_and_expand_eagle_inputs_kernel_impl( const int64_t padding_token_id, const int64_t parallel_drafting_token_id, const int64_t total_input_tokens, const int64_t num_padding_slots_per_request, const bool shift_input_ids); +void copy_and_expand_dflash_inputs_kernel_impl( + const torch::Tensor& next_token_ids, const torch::Tensor& target_positions, + torch::Tensor& out_input_ids, torch::Tensor& out_context_positions, + torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping, + torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices, + const torch::Tensor& block_table, const torch::Tensor& query_start_loc, + const std::optional& num_rejected_tokens, + const int64_t parallel_drafting_token_id, const int64_t block_size, + const int64_t num_query_per_req, const int64_t num_speculative_tokens, + const int64_t total_input_tokens, const bool has_num_rejected); void rejection_greedy_sample_kernel_impl( torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax, @@ -599,6 +609,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "SymInt total_input_tokens, SymInt num_padding_slots_per_request, " "bool shift_input_ids) -> ()", &cpu_utils::copy_and_expand_eagle_inputs_kernel_impl); + ops.def( + "copy_and_expand_dflash_inputs_kernel_impl(" + "Tensor next_token_ids, Tensor target_positions, " + "Tensor(a2!) out_input_ids, Tensor(a3!) out_context_positions, " + "Tensor(a4!) out_query_positions, " + "Tensor(a5!) out_context_slot_mapping, " + "Tensor(a6!) out_query_slot_mapping, " + "Tensor(a7!) out_token_indices, Tensor block_table, " + "Tensor query_start_loc, Tensor? num_rejected_tokens, " + "SymInt parallel_drafting_token_id, SymInt block_size, " + "SymInt num_query_per_req, SymInt num_speculative_tokens, " + "SymInt total_input_tokens, bool has_num_rejected) -> ()", + &cpu_utils::copy_and_expand_dflash_inputs_kernel_impl); ops.def( "rejection_greedy_sample_kernel_impl(" "Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, " diff --git a/csrc/custom_quickreduce.cu b/csrc/custom_quickreduce.cu index 33d0d4a7226..d4e5d179a54 100644 --- a/csrc/custom_quickreduce.cu +++ b/csrc/custom_quickreduce.cu @@ -97,18 +97,28 @@ int64_t qr_max_size() { cast_bf2half>; \ template struct quickreduce::AllReduceTwoshot, cast_bf2half>; + // INT3 (CodecQ3) is restricted to TP2 only, so we only instantiate the + // world_size == 2 kernel for it. + #define INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(T, Codec, cast_bf2half) \ + template struct quickreduce::AllReduceTwoshot, cast_bf2half>; + INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, false) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, false) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, false) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, false) +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16, + quickreduce::CodecQ3, false) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecFP, true) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ4, true) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ6, true) INSTANTIATE_FOR_WORLDSIZE(quickreduce::nv_bfloat16, quickreduce::CodecQ8, true) +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(quickreduce::nv_bfloat16, + quickreduce::CodecQ3, true) INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecFP, false) INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ4, false) INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ6, false) INSTANTIATE_FOR_WORLDSIZE(half, quickreduce::CodecQ8, false) +INSTANTIATE_FOR_WORLDSIZE_TP2_ONLY(half, quickreduce::CodecQ3, false) #endif // USE_ROCM \ No newline at end of file diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu index 8bdb4f56795..51f84d2ffd9 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu @@ -180,7 +180,7 @@ bool cutlass_group_gemm_supported(int64_t cuda_device_capability) { #if defined CUDA_VERSION #if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100 - if (cuda_device_capability >= 100 && cuda_device_capability < 110) { + if (cuda_device_capability >= 100 && cuda_device_capability < 120) { return CUDA_VERSION >= 12080; } #endif diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu index f3a9122b9ee..050c4322fa9 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu @@ -186,6 +186,10 @@ __global__ void per_token_group_quant_8bit_kernel( static_cast(output_q) + block_group_offset; scale_element_t* scale_output; +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + if constexpr (IS_COLUMN_MAJOR) { const int num_elems_per_pack = static_cast(sizeof(scale_packed_t) / sizeof(scale_element_t)); @@ -220,6 +224,10 @@ __global__ void per_token_group_quant_8bit_kernel( QuantizeGroup(smem_group, group_output, group_size, lane_id, threads_per_group, y_s, min_8bit, max_8bit); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif } inline int GetGroupsPerBlock(int64_t num_groups) { @@ -358,45 +366,56 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, (input.scalar_type() == torch::headeronly::ScalarType::Half || input.scalar_type() == torch::headeronly::ScalarType::BFloat16); -#define LAUNCH_KERNEL(T, DST_DTYPE) \ - do { \ - dim3 grid(num_blocks); \ - dim3 block(num_threads); \ - size_t smem_bytes = \ - static_cast(groups_per_block) * group_size * sizeof(T); \ - if (is_column_major) { \ - if (scale_ue8m0) { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit, scale_num_rows, scale_stride); \ - } else { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit, scale_num_rows, scale_stride); \ - } \ - } else { \ - if (scale_ue8m0) { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit); \ - } else { \ - per_token_group_quant_8bit_kernel \ - <<>>( \ - static_cast(input.data_ptr()), output_q.data_ptr(), \ - static_cast(output_s.data_ptr()), group_size, \ - num_groups, groups_per_block, (float)eps, (float)min_8bit, \ - (float)max_8bit); \ - } \ - } \ +#ifndef USE_ROCM + #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \ + do { \ + cudaLaunchConfig_t config = {}; \ + config.gridDim = dim3(num_blocks); \ + config.blockDim = dim3(num_threads); \ + config.dynamicSmemBytes = (SMEM_BYTES); \ + config.stream = stream; \ + cudaLaunchAttribute attrs[1]; \ + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \ + attrs[0].val.programmaticStreamSerializationAllowed = 1; \ + config.numAttrs = 1; \ + config.attrs = attrs; \ + cudaLaunchKernelEx( \ + &config, \ + per_token_group_quant_8bit_kernel, \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + static_cast(output_s.data_ptr()), group_size, num_groups, \ + groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \ + scale_num_rows, scale_stride); \ + } while (0) +#else + #define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \ + do { \ + per_token_group_quant_8bit_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + static_cast(output_s.data_ptr()), group_size, \ + num_groups, groups_per_block, (float)eps, (float)min_8bit, \ + (float)max_8bit, scale_num_rows, scale_stride); \ + } while (0) +#endif + +#define LAUNCH_KERNEL(T, DST_DTYPE) \ + do { \ + size_t smem_bytes = \ + static_cast(groups_per_block) * group_size * sizeof(T); \ + if (is_column_major) { \ + if (scale_ue8m0) { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes); \ + } else { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes); \ + } \ + } else { \ + if (scale_ue8m0) { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes); \ + } else { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes); \ + } \ + } \ } while (0) #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, IS_COLUMN_MAJOR, SCALE_UE8M0, KX, \ @@ -484,6 +503,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, #undef LAUNCH_REG_KERNEL_INST #undef LAUNCH_REG_KERNEL_FLAGS #undef LAUNCH_KERNEL +#undef LAUNCH_KERNEL_INST } // Register-resident fast path for group_size==128. @@ -528,12 +548,12 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel( const int mn_idx = blockIdx.x * kRowsPerBlock + row_local; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); + cudaGridDependencySynchronize(); #endif if (mn_idx >= tma_aligned_mn) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif return; } @@ -607,7 +627,7 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel( *reinterpret_cast(group_output) = packed_out; #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); + cudaTriggerProgrammaticLaunchCompletion(); #endif } diff --git a/csrc/quickreduce/base.h b/csrc/quickreduce/base.h index a2170e48320..6c3456d06f2 100644 --- a/csrc/quickreduce/base.h +++ b/csrc/quickreduce/base.h @@ -283,6 +283,29 @@ __quickreduce_device_inline__ int packed_rcp(int a) { return R.i; } +template +__quickreduce_device_inline__ int packed_from_int16_pair(int16_t low, + int16_t high); + +template <> +__quickreduce_device_inline__ int packed_from_int16_pair(int16_t low, + int16_t high) { + // Convert two signed integers to one fp16x2 packed 32-bit lane. + half2 h = __halves2half2(__int2half_rn(static_cast(low)), + __int2half_rn(static_cast(high))); + return __builtin_bit_cast(int, h); +} + +template <> +__quickreduce_device_inline__ int packed_from_int16_pair( + int16_t low, int16_t high) { + // Convert two signed integers to one bf16x2 packed 32-bit lane. + nv_bfloat16 bf_low = __float2bfloat16(static_cast(low)); + nv_bfloat16 bf_high = __float2bfloat16(static_cast(high)); + nv_bfloat162 bf2 = __halves2bfloat162(bf_low, bf_high); + return *reinterpret_cast(&bf2); +} + // changes dtype __quickreduce_device_inline__ float T2float_cast(half a) { return __half2float(a); diff --git a/csrc/quickreduce/quick_reduce.h b/csrc/quickreduce/quick_reduce.h index 4cc35300bf8..7506329972b 100644 --- a/csrc/quickreduce/quick_reduce.h +++ b/csrc/quickreduce/quick_reduce.h @@ -59,11 +59,30 @@ allreduce_prototype_twoshot(T const* A, T* B, uint32_t N, uint32_t num_blocks, flag_color, this->kMaxProblemSize); \ } +// INT3 only retains good performance on TP2 (world_size == 2). On TP4/TP8 +// the 3-bit codec's pack/unpack overhead outweighs the reduced communication +// volume, so INT3 is restricted to a TP2-only dispatch here. +#define TWOSHOT_DISPATCH_TP2_ONLY(__codec) \ + if (world_size == 2) { \ + using LineCodec = __codec; \ + using AllReduceKernel = AllReduceTwoshot; \ + hipLaunchKernelGGL((allreduce_prototype_twoshot), \ + dim3(grid), dim3(kBlockTwoShot), 0, stream, A, B, N, \ + num_blocks, rank, dbuffer_list, data_offset, \ + flag_color, this->kMaxProblemSize); \ + } else { \ + throw std::runtime_error( \ + "INT3 quick all-reduce is only supported for world_size == 2 " \ + "(TP2); use INT4/NONE for larger world sizes."); \ + } + enum QuickReduceQuantLevel { - F16 = 0, - INT8 = 1, - INT6 = 2, - INT4 = 3, + // Keep these ids in sync with Python QuickReduceRegime enum. + F16 = 0, // full-precision fp16/bf16 communication + INT8 = 1, // symmetric int8 + per-block scale + INT6 = 2, // symmetric int6 + per-block scale + INT4 = 3, // symmetric int4 + per-block scale + INT3 = 4, // symmetric int3 + per-block scale (TP2 only) }; struct DeviceComms { @@ -184,6 +203,9 @@ struct DeviceComms { case QuickReduceQuantLevel::INT4: TWOSHOT_DISPATCH(CodecQ4) break; + case QuickReduceQuantLevel::INT3: + TWOSHOT_DISPATCH_TP2_ONLY(CodecQ3) + break; default: TWOSHOT_DISPATCH(CodecFP) break; diff --git a/csrc/quickreduce/quick_reduce_impl.cuh b/csrc/quickreduce/quick_reduce_impl.cuh index 38dc9938fc8..e9586e40ff0 100644 --- a/csrc/quickreduce/quick_reduce_impl.cuh +++ b/csrc/quickreduce/quick_reduce_impl.cuh @@ -2,6 +2,7 @@ #include #include "base.h" +#include namespace quickreduce { @@ -206,6 +207,168 @@ struct CodecQ4 : public CodecBase { } }; +// Int3 symmetric quantization codec. +// We quantize the FP16 data to block-scaled Int3 in blocks of 4 * +// kThreadGroupSize. Uniform symmetric quantization (round-to-int + clip), +// matching the structure of CodecQ4. Signed range is [-4, +3]. +template +struct CodecQ3 : public CodecBase { + static constexpr int kWorldSize = world_size; + + // Layout per quantization block (32 values = 8 threads * 4 fp16x2 lanes): + // - each thread owns 8 values and writes: + // * q2 payload : 8 * 2 bits -> uint16 (2 bytes) + // * q1 payload : 8 * 1 bit -> uint8 (1 byte) + // - one scale is shared per 32 values and written by group leader. + // + // kRankTileStride is split as: + // [0 .. 511] : q2 payload region (256 threads * 2 bytes) + // [512 .. 767] : q1 payload region (256 threads * 1 byte) + // [768 .. 895] : scale region (32 groups * 4 bytes) + static constexpr int kRankAtoms = kAtoms / kWorldSize; + static constexpr int kRankTileStride = 896; + static constexpr int kRankTileQ1Offset = 512; + static constexpr int kRankTileScaleOffset = 768; + static constexpr int kRankTransmittedTileSize = kRankTileStride * kRankAtoms; + static_assert(kRankTransmittedTileSize % 16 == 0, + "kRankTransmittedTileSize must be 16B aligned."); + + static constexpr int kRankBufferTileStride = + kRankTileStride / sizeof(int32x4_t); + + static constexpr int kTransmittedTileSize = + kRankTransmittedTileSize * kWorldSize; + + // {-1/4.0h, -1/4.0h}, f16x2_t / bf16x2_t. Sign-flipped so absmax maps + // to -4; the sign cancels with decoding_scale on the recv side. + static constexpr int kScaleFactor = + std::is_same::value ? 0xB400B400 : 0xBE80BE80; + + // {1e-7, 1e-7}, f16x2_t + static constexpr int kScaleEpsilon = + std::is_same::value ? 0x00010001 : 0x33D733D7; + + // {-4, -4}, f16x2_t / bf16x2_t + static constexpr int kRangeMin = + std::is_same::value ? 0xC400C400 : 0xC080C080; + + // {+3, +3}, f16x2_t / bf16x2_t + static constexpr int kRangeMax = + std::is_same::value ? 0x42004200 : 0x40404040; + + // {+4, +4}, int16x2_t -- shifts signed [-4, +3] to unsigned [0, 7]. + static constexpr int kRangeBias = 0x00040004; + + __quickreduce_device_inline__ CodecQ3(int thread, int rank) + : CodecBase(thread, rank) {} + + __quickreduce_device_inline__ void send(int32x4_t* __restrict__ send_buffer, + const int32x4_t* __restrict__ data) { + for (int k = 0; k < kRankAtoms; k++) { + int32x4_t const atom = data[k]; + + // 1) Per-group dynamic scale (shared across 32 values). + int wblockmax = group_abs_max(atom); + int decoding_scale = packed_mul(wblockmax, kScaleFactor); + int encoding_scale = packed_add(decoding_scale, kScaleEpsilon); + encoding_scale = packed_rcp(encoding_scale); + + // 2) Scale + clip to signed int3 range [-4, +3]. + int32x4_t w; + for (int i = 0; i < 4; i++) { + w[i] = packed_mul(atom[i], encoding_scale); + w[i] = packed_max(w[i], kRangeMin); + w[i] = packed_min(w[i], kRangeMax); + } + + // 3) Round to integer and bias to unsigned domain [0, 7]. + int32x4_t q; + { + int16_t* qi = reinterpret_cast(&q); + T* wh = reinterpret_cast(&w); + for (int i = 0; i < 8; i++) qi[i] = (int16_t)rintf(T2float_cast(wh[i])); + + for (int i = 0; i < 4; i++) { + q[i] = packed_add(q[i], kRangeBias); + } + } + + // 4) Split each 3-bit unsigned value into low-2-bit and high-1-bit + // halves, packed into one uint16 (low 2 bits per value) plus one + // uint8 (high 1 bit per value). + uint16_t q2w = 0; + uint8_t q1w = 0; + { + int16_t* tw = reinterpret_cast(&q); +#pragma unroll + for (int i = 0; i < 8; i++) { + uint32_t v = static_cast(tw[i]) & 0x7u; + q2w |= static_cast((v & 0x3u) << (i * 2)); + q1w |= static_cast(((v >> 2) & 0x1u) << i); + } + } + + uint8_t* atom_ptr = + reinterpret_cast(send_buffer + k * kRankBufferTileStride); + uint16_t* q2w_ptr = reinterpret_cast(atom_ptr) + thread; + uint8_t* q1w_ptr = + reinterpret_cast(atom_ptr + kRankTileQ1Offset) + thread; + int* qs_ptr = reinterpret_cast(atom_ptr + kRankTileScaleOffset) + + (thread / 8); + + __builtin_nontemporal_store(q2w, q2w_ptr); + *q1w_ptr = q1w; + if (threadIdx.x == group_leader) { + __builtin_nontemporal_store(decoding_scale, qs_ptr); + } + } + } + + __quickreduce_device_inline__ void recv(int32x4_t** __restrict__ recv_buffer, + int32x4_t* __restrict__ data) { + for (int k = 0; k < kRankAtoms; k++) { + uint8_t* atom_ptr = reinterpret_cast(*recv_buffer); + uint16_t* q2w_ptr = reinterpret_cast(atom_ptr) + thread; + uint8_t* q1w_ptr = + reinterpret_cast(atom_ptr + kRankTileQ1Offset) + thread; + int* qs_ptr = reinterpret_cast(atom_ptr + kRankTileScaleOffset) + + (thread / 8); + + uint16_t q2w = __builtin_nontemporal_load(q2w_ptr); + uint8_t q1w = *q1w_ptr; + int qs = __builtin_nontemporal_load(qs_ptr); + + *recv_buffer += kRankBufferTileStride; + + // Unpack unsigned values [0, 7] then shift back to signed domain + // [-4, +3] by adding kRangeMin. + int32x4_t w; + { + int16_t qv[8]; +#pragma unroll + for (int i = 0; i < 8; i++) { + uint32_t low2 = (q2w >> (2 * i)) & 0x3u; + uint32_t high1 = (q1w >> i) & 0x1u; + qv[i] = static_cast(low2 | (high1 << 2)); + } + +#pragma unroll + for (int i = 0; i < 4; i++) { + int qpack = packed_from_int16_pair(qv[2 * i], qv[2 * i + 1]); + w[i] = packed_add(qpack, kRangeMin); + } + } + + // Apply decode scale to reconstruct fp16/bf16 lanes. + for (int i = 0; i < 4; i++) { + w[i] = packed_mul(w[i], qs); + } + + data[k] = w; + } + } +}; + // Int6 symmetric quantization codec. // We quantize the FP16 data to block-scaled Int6 in blocks of 4 * // kThreadGroupSize. @@ -377,7 +540,6 @@ struct CodecQ6 : public CodecBase { w[i] = packed_mul(w[i], qs); } - // That's pretty much it... data[k] = w; } } diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 9e6c0726d19..4ac255d0a75 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -1045,7 +1045,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( const scalar_t* q_ptr = q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE; const _B16x8* q_ptrh8 = reinterpret_cast(q_ptr); - const int qhead_elemh8 = laneid / 4; + const int qhead_elemh8 = MIN(laneid / 4, HEAD_SIZE / 8 - 1); for (int h = 0; h < QHLOOP - 1; h++) { const int qhead_idx = h * 4 + lane4id; diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 10e3cbf2e0b..615cdabed58 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -70,6 +70,15 @@ bool on_gfx12() { return result; } +bool on_gfx1151() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx1151") != std::string::npos; + }(); + return result; +} + #if defined(NDEBUG) #undef NDEBUG #include @@ -1237,6 +1246,45 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \ } +// WVSPLITK_CFG arguments are: (THRDS, WVPRGRP, YTILE, UNRL, N). +// THRDS = wavefront width (32 on GFX11/GFX12, 64 on GFX9) +// WVPRGRP= waves per group (always 16) +// YTILE = output rows per thread tile +// UNRL = K-loop unroll factor +// N = batch size (passed through from the switch in wvSplitK) +#define WVSPLIT_TILE(_sYT, __N) \ + { \ + if (on_gfx1151()) { \ + bool fit_lds = (Kbp_in * N_in <= max_lds_len); \ + if (_sYT <= 1) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \ + __N) \ + else if ((K_in % 1024 == 512) && K_in >= 1536 && \ + (_sYT >= 40 || K_in >= 4096)) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/4, /*UNRL=*/1, \ + __N) \ + else if (K_in < 1024) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/2, /*UNRL=*/4, \ + __N) \ + else if (K_in <= 2048 && (__N >= 2 || _sYT <= 26)) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \ + __N) \ + else if (__N >= 2 && !fit_lds) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/4, \ + __N) \ + else if (__N == 1) \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/2, \ + __N) \ + else \ + WVSPLITK_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, /*YTILE=*/1, /*UNRL=*/1, \ + __N) \ + } else if (on_gfx1x()) { /* gfx1100/gfx1150/GFX12, wave32 */ \ + WVSPLIT_TILE_CFG(/*THRDS=*/32, /*WVPRGRP=*/16, _sYT, __N) \ + } else { /* GFX9, wave64 */ \ + WVSPLIT_TILE_CFG(/*THRDS=*/64, /*WVPRGRP=*/16, _sYT, __N) \ + } \ + } + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] { using fptype = typename scalar::type; fptype* af4 = reinterpret_cast(in_a.data_ptr()); @@ -1251,37 +1299,21 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, // then cut the active waves to balance their distribution... int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); - const bool use_wave32 = on_gfx1x(); switch (N_in) { case 1: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 1) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 1) + WVSPLIT_TILE(sYT, 1) break; case 2: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 2) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 2) + WVSPLIT_TILE(sYT, 2) break; case 3: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 3) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 3) + WVSPLIT_TILE(sYT, 3) break; case 4: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 4) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 4) + WVSPLIT_TILE(sYT, 4) break; case 5: - if (use_wave32) - WVSPLIT_TILE_CFG(32, 16, sYT, 5) - else - WVSPLIT_TILE_CFG(64, 16, sYT, 5) + WVSPLIT_TILE(sYT, 5) break; default: throw std::runtime_error( diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index a3b2a539bd9..fbd5e1e60e3 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" -ARG AITER_BRANCH="v0.1.13.post1" +ARG AITER_BRANCH="v0.1.16.post2" ARG AITER_REPO="https://github.com/ROCm/aiter.git" ARG MORI_BRANCH="v1.1.0" ARG MORI_REPO="https://github.com/ROCm/mori.git" @@ -244,7 +244,7 @@ RUN pip install pyyaml && cd aiter \ export HIP_CLANG_PATH=/opt/sccache-wrappers \ && sccache --show-stats; \ fi \ - && PREBUILD_KERNELS=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + && PREBUILD_KERNELS=1 AITER_USE_SYSTEM_TRITON=1 GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ && ls /app/aiter/dist/*.whl RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index f965127cbfb..d268d5b4db2 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -159,7 +159,7 @@ Priority is **1 = highest** (tried first). | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ | -| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A | +| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ✅ | ❌ | ❌ | All | N/A | | `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x | | `FLASHINFER` | TRTLLM† | 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 | diff --git a/docs/design/cuda_graphs_multimodal.md b/docs/design/cuda_graphs_multimodal.md index 1db82ffa688..264b1f139f6 100644 --- a/docs/design/cuda_graphs_multimodal.md +++ b/docs/design/cuda_graphs_multimodal.md @@ -134,7 +134,8 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra | `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ | | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ | | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ | -| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ | +| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ | | `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ❌︎ | !!! note diff --git a/docs/design/hybrid_kv_cache_manager.md b/docs/design/hybrid_kv_cache_manager.md index 8f17b473adc..82d54e9b5c1 100644 --- a/docs/design/hybrid_kv_cache_manager.md +++ b/docs/design/hybrid_kv_cache_manager.md @@ -159,7 +159,7 @@ For simplicity, we assume `block_size=1` in this section. ### High level idea -The block pool uses a dict similar to `tuple(block_hash, group_id) -> block` to catch the full blocks. That means the same tokens of different groups are cached and evicted independently. +The block pool uses a dict similar to `tuple(block_hash, group_id) -> block` to cache the full blocks. That means the same tokens of different groups are cached and evicted independently. When a new request comes in, we check the cache hit prefix of each group, and return the intersection of these groups as the cached prefix of the request. See below for the detailed algorithm for checking the cache hit of one group & performing the intersection. diff --git a/docs/design/paged_attention.md b/docs/design/paged_attention.md index 7c0132cd2a2..f4742c7faaa 100644 --- a/docs/design/paged_attention.md +++ b/docs/design/paged_attention.md @@ -52,7 +52,7 @@ __device__ void paged_attention_kernel( ) ``` -There are also a list of template arguments above the function +There is also a list of template arguments above the function signature that are determined during compilation time. `scalar_t` represents the data type of the query, key, and value data elements, such as FP16. `HEAD_SIZE` indicates the number of elements in each @@ -178,7 +178,7 @@ const scalar_t* k_ptr = k_cache + physical_block_number * kv_block_stride + physical_block_offset * x; ``` -Unlike to `q_ptr`, `k_ptr` in each thread will point to different +Unlike `q_ptr`, `k_ptr` in each thread will point to different key token at different iterations. As shown above, that `k_ptr` points to key token data based on `k_cache` at assigned block, assigned head and assigned token. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 82022a08608..294c0c6b3f2 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -366,7 +366,6 @@ th { | ------------ | ------ | ----------------- | -------------------- | ------------------------- | | `AfmoeForCausalLM` | Afmoe | TBA | ✅︎ | ✅︎ | | `ApertusForCausalLM` | Apertus | `swiss-ai/Apertus-8B-2509`, `swiss-ai/Apertus-70B-Instruct-2509`, etc. | ✅︎ | ✅︎ | -| `AquilaForCausalLM` | Aquila, Aquila2 | `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc. | ✅︎ | ✅︎ | | `ArceeForCausalLM` | Arcee (AFM) | `arcee-ai/AFM-4.5B-Base`, etc. | ✅︎ | ✅︎ | | `ArcticForCausalLM` | Arctic | `Snowflake/snowflake-arctic-base`, `Snowflake/snowflake-arctic-instruct`, etc. | | ✅︎ | | `AXK1ForCausalLM` | A.X-K1 | `skt/A.X-K1`, etc. | | ✅︎ | diff --git a/pyproject.toml b/pyproject.toml index 031f8d1a0a2..249832ff2e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,8 +129,8 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html", "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*", "rust/src/chat/src/renderer/deepseek_v32/fixtures/*", - "rust/src/tool-parser/src/gemma4.rs", "rust/src/text/src/output/decoded.rs", - "rust/src/tokenizer/src/incremental.rs", "rust/src/reasoning-parser/src/tests.rs"] + "rust/src/parser/src/tool/gemma4.rs", "rust/src/text/src/output/decoded.rs", + "rust/src/tokenizer/src/incremental.rs", "rust/src/parser/src/reasoning/tests.rs"] ignore-hidden = false [tool.typos.default] diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 743633b447e..70c325c152d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5733,10 +5733,9 @@ dependencies = [ "uuid", "vllm-engine-core-client", "vllm-llm", - "vllm-reasoning-parser", + "vllm-parser", "vllm-text", "vllm-tokenizer", - "vllm-tool-parser", "xgrammar-structural-tag", "zeromq", ] @@ -5872,11 +5871,22 @@ dependencies = [ ] [[package]] -name = "vllm-reasoning-parser" +name = "vllm-parser" version = "0.1.0" dependencies = [ + "criterion", + "easy-ext", + "expect-test", + "futures", + "openai-protocol", + "serde", + "serde_json", "thiserror 2.0.18", + "thiserror-ext", + "tool-parser", "vllm-tokenizer", + "winnow", + "xgrammar-structural-tag", ] [[package]] @@ -5978,24 +5988,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "vllm-tool-parser" -version = "0.1.0" -dependencies = [ - "criterion", - "easy-ext", - "expect-test", - "futures", - "openai-protocol", - "serde", - "serde_json", - "thiserror 2.0.18", - "thiserror-ext", - "tool-parser", - "winnow", - "xgrammar-structural-tag", -] - [[package]] name = "vllm-tool-parser-py" version = "0.1.0" @@ -6004,7 +5996,7 @@ dependencies = [ "pythonize", "serde_json", "thiserror-ext", - "vllm-tool-parser", + "vllm-parser", ] [[package]] diff --git a/rust/Cargo.toml b/rust/Cargo.toml index e31bf07bbd4..dc3895c372d 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -7,12 +7,11 @@ members = [ "src/managed-engine", "src/metrics", "src/mock-engine", - "src/reasoning-parser", + "src/parser", + "src/parser/python", "src/server", "src/text", "src/tokenizer", - "src/tool-parser", - "src/tool-parser/python", ] resolver = "3" @@ -117,11 +116,10 @@ vllm-engine-core-client = { path = "src/engine-core-client" } vllm-llm = { path = "src/llm" } vllm-managed-engine = { path = "src/managed-engine" } vllm-metrics = { path = "src/metrics" } -vllm-reasoning-parser = { path = "src/reasoning-parser" } +vllm-parser = { path = "src/parser" } vllm-server = { path = "src/server" } vllm-text = { path = "src/text" } vllm-tokenizer = { path = "src/tokenizer" } -vllm-tool-parser = { path = "src/tool-parser" } winnow = "1.0.2" xgrammar-structural-tag = "0.1.0" zeromq = { version = "0.6.0", default-features = false, features = [ diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index 85368ad98b9..f5860c18597 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -30,10 +30,9 @@ trait-set.workspace = true uuid.workspace = true vllm-engine-core-client.workspace = true vllm-llm.workspace = true -vllm-reasoning-parser.workspace = true +vllm-parser.workspace = true vllm-text.workspace = true vllm-tokenizer.workspace = true -vllm-tool-parser.workspace = true xgrammar-structural-tag.workspace = true [dev-dependencies] diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index dbc9cc05a51..8c9a4362d1c 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -1,44 +1,34 @@ //! Default output processing pipeline. -mod reasoning; mod structural_tag; -mod tool; +mod unified; use std::sync::Once; -use futures::{Stream, StreamExt as _}; +use futures::StreamExt as _; use tracing::info; -use trait_set::trait_set; +use vllm_parser::unified::{CombinedParser, UnifiedParser}; use vllm_text::tokenizer::DynTokenizer; -use self::reasoning::reasoning_event_stream; use self::structural_tag::apply_structural_tag_constraint; -use self::tool::tool_event_stream; +use self::unified::unified_event_stream; use super::structured::structured_chat_event_stream; use crate::error::Result; -use crate::output::{ - AssistantEvent, ChatOutputProcessor, ContentEvent, DynChatEventStream, - DynDecodedTextEventStream, -}; +use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEventStream}; use crate::parser::ParserSelection; use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory}; use crate::parser::tool::{ToolParser, ToolParserFactory}; use crate::request::ChatRequest; use crate::{Error, Result as ChatResult}; -trait_set! { - trait ContentEventStream = Stream> + Send + 'static; -} - /// Default request-scoped output processor used by Hugging Face style chat /// backends. /// /// This implementation assumes the backend already emitted decoded text deltas, -/// then optionally layers reasoning parsing and tool-call parsing before +/// then optionally layers unified reasoning and tool-call parsing before /// assembling final structured chat events. pub struct DefaultChatOutputProcessor { - reasoning_parser: Option>, - tool_parser: Option>, + parser: Box, parallel_tool_calls: bool, } @@ -66,16 +56,17 @@ impl DefaultChatOutputProcessor { } else { None }; - let reasoning_parser = Self::resolve_optional_reasoning_parser( - request, - model_id, - tokenizer, - reasoning_parser, - )?; + let reasoning_parser = + Self::resolve_optional_reasoning_parser(model_id, tokenizer, reasoning_parser)?; + let parser: Box = + Box::new(CombinedParser::new(reasoning_parser, tool_parser)); + + if parser.preserve_special_tokens() { + request.decode_options.skip_special_tokens = false; + } Ok(Self { - reasoning_parser, - tool_parser, + parser, parallel_tool_calls: request.parallel_tool_calls, }) } @@ -87,8 +78,7 @@ impl DefaultChatOutputProcessor { /// content is treated as opaque text. pub fn plain_text_only() -> Self { Self { - reasoning_parser: None, - tool_parser: None, + parser: Box::new(CombinedParser::plain_text_only()), parallel_tool_calls: true, } } @@ -112,10 +102,6 @@ impl DefaultChatOutputProcessor { let parser = factory.create(parser_name, &request.tools)?; - if parser.preserve_special_tokens() { - request.decode_options.skip_special_tokens = false; - } - apply_structural_tag_constraint(request, parser.as_ref())?; TOOL_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using tool parser")); @@ -123,7 +109,6 @@ impl DefaultChatOutputProcessor { } fn resolve_optional_reasoning_parser( - request: &mut ChatRequest, model_id: &str, tokenizer: DynTokenizer, selection: &ParserSelection, @@ -142,10 +127,6 @@ impl DefaultChatOutputProcessor { let parser = factory.create(parser_name, tokenizer)?; - if parser.preserve_special_tokens() { - request.decode_options.skip_special_tokens = false; - } - REASONING_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using reasoning parser")); Ok(Some(parser)) } @@ -156,16 +137,14 @@ static REASONING_PARSER_LOG_ONCE: Once = Once::new(); impl ChatOutputProcessor for DefaultChatOutputProcessor { /// Transforms a raw generate-output token stream into structured chat - /// events through three sequential stages once text decoding has + /// events through two sequential stages once text decoding has /// already happened: /// - /// 1. [`reasoning_event_stream`] — reasoning/content separation - /// 2. [`tool_event_stream`] — tool-call parsing - /// 3. [`structured_chat_event_stream`] — final block assembly + /// 1. [`unified_event_stream`] — reasoning and tool-call parsing + /// 2. [`structured_chat_event_stream`] — final block assembly fn process(self: Box, decoded: DynDecodedTextEventStream) -> Result { - let reasoning = reasoning_event_stream(decoded, self.reasoning_parser); - let tool = tool_event_stream(reasoning, self.tool_parser); - let structured = structured_chat_event_stream(tool, self.parallel_tool_calls); + let parsed = unified_event_stream(decoded, self.parser); + let structured = structured_chat_event_stream(parsed, self.parallel_tool_calls); Ok(structured.boxed()) } diff --git a/rust/src/chat/src/output/default/reasoning.rs b/rust/src/chat/src/output/default/reasoning.rs deleted file mode 100644 index faa9d7894bb..00000000000 --- a/rust/src/chat/src/output/default/reasoning.rs +++ /dev/null @@ -1,509 +0,0 @@ -//! Adapts decoded text updates into reasoning-aware assistant deltas. -//! -//! This stage sits between low-level token decoding and final block assembly. -//! It is the only place in the new pipeline that understands reasoning -//! separation: `decoded.rs` still only produces plain text deltas, while later -//! stages consume the semantic `Text` / `Reasoning` split emitted here. - -use asynk_strim_attr::{TryYielder, try_stream}; -use futures::{StreamExt as _, pin_mut}; -use thiserror_ext::AsReport; -use tracing::warn; -use vllm_text::output::DecodedTextEvent; - -use super::ContentEvent; -use crate::Result; -use crate::error::Error; -use crate::event::AssistantBlockKind; -use crate::output::DecodedTextEventStream; -use crate::parser::reasoning::{ReasoningDelta, ReasoningParser}; - -/// Per-stream reasoning parsing state. -struct ReasoningState { - /// Reasoning parser for the current model family. - parser: Box, - /// Whether reasoning parsing has already failed for this stream. - parser_failed: bool, -} - -impl ReasoningState { - /// Create one fresh reasoning-adaptation state for a new streamed response. - fn new(parser: Box) -> Self { - Self { - parser, - parser_failed: false, - } - } - - /// Convert one decoded text delta into zero or more semantic assistant - /// deltas. - fn process_delta(&mut self, delta: String) -> Vec { - // If the parser has already failed, skip parsing and return plain text deltas. - if self.parser_failed { - return vec![ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta, - }]; - } - - let mut events = Vec::new(); - - match self.parser.push(&delta) { - Ok(result) => { - push_reasoning_delta(&mut events, result); - } - Err(error) => { - if !self.parser_failed { - warn!( - error = %error.as_report(), - "reasoning parser failed; falling back to plain text deltas" - ); - self.parser_failed = true; - } - push_text_delta(&mut events, AssistantBlockKind::Text, delta); - } - } - - events - } - - /// Initialize parser state once prompt token IDs are available. - fn initialize(&mut self, prompt_token_ids: &[u32]) { - if self.parser_failed { - return; - } - - match self.parser.initialize(prompt_token_ids) { - Ok(()) => {} - Err(error) => { - warn!( - error = %error.as_report(), - "failed to initialize reasoning parser; falling back to plain text deltas" - ); - self.parser_failed = true; - } - } - } - - /// Flush any parser-held partial delimiter state at end of stream. - fn finish(&mut self) -> Vec { - if self.parser_failed { - return Vec::new(); - } - - match self.parser.finish() { - Ok(result) => { - let mut events = Vec::new(); - push_reasoning_delta(&mut events, result); - events - } - Err(error) => { - warn!(error = %error.as_report(), "failed to flush reasoning parser state"); - Vec::new() - } - } - } -} - -/// Push one semantic text delta if it is non-empty. -fn push_text_delta(events: &mut Vec, kind: AssistantBlockKind, delta: String) { - if delta.is_empty() { - return; - } - events.push(ContentEvent::TextDelta { kind, delta }); -} - -/// Convert one parsed reasoning delta into zero or more content events. -fn push_reasoning_delta(events: &mut Vec, delta: ReasoningDelta) { - if let Some(reasoning) = delta.reasoning { - push_text_delta(events, AssistantBlockKind::Reasoning, reasoning); - } - if let Some(content) = delta.content { - push_text_delta(events, AssistantBlockKind::Text, content); - } -} - -/// Wrap one decoded-text stream into the internal reasoning event stream. -#[try_stream] -pub(crate) async fn reasoning_event_stream( - decoded_stream: impl DecodedTextEventStream, - reasoning_parser: Option>, - mut y: TryYielder, -) -> Result<()> { - pin_mut!(decoded_stream); - - // Without a parser, pass through as plain text deltas. - let Some(reasoning_parser) = reasoning_parser else { - while let Some(event) = decoded_stream.next().await.transpose()? { - for next in ContentEvent::from_decoded_plain_text(event) { - y.yield_ok(next).await; - } - } - return Ok(()); - }; - - let mut state = ReasoningState::new(reasoning_parser); - - while let Some(event) = decoded_stream.next().await.transpose()? { - match event { - DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - } => { - state.initialize(&prompt_token_ids); - y.yield_ok(ContentEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) - .await; - } - DecodedTextEvent::TextDelta { - delta, - token_ids, - logprobs, - finished, - } => { - for next in state.process_delta(delta) { - y.yield_ok(next).await; - } - if logprobs.is_some() || !token_ids.is_empty() { - y.yield_ok(ContentEvent::LogprobsDelta { - logprobs, - token_ids, - }) - .await; - } - if let Some(finished) = finished { - for next in state.finish() { - y.yield_ok(next).await; - } - y.yield_ok(ContentEvent::Done { - usage: finished.usage, - finish_reason: finished.finish_reason, - kv_transfer_params: finished.kv_transfer_params, - }) - .await; - } - } - } - } - Ok(()) -} - -#[cfg(test)] -mod tests { - - use std::sync::Arc; - - use futures::{StreamExt as _, stream}; - use vllm_llm::FinishReason; - use vllm_text::output::{ - DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEvent, DecodedTokenLogprob, - }; - use vllm_tokenizer::{DynTokenizer, Tokenizer}; - - use super::super::ContentEvent; - use super::reasoning_event_stream; - use crate::event::AssistantBlockKind; - use crate::parser::reasoning::{ - ReasoningDelta, ReasoningError, ReasoningParser, ReasoningParserFactory, names, - }; - - struct FakeTokenizer; - - impl Tokenizer for FakeTokenizer { - fn encode( - &self, - text: &str, - _add_special_tokens: bool, - ) -> vllm_tokenizer::Result> { - Ok(text.chars().map(u32::from).collect()) - } - - fn decode( - &self, - token_ids: &[u32], - _skip_special_tokens: bool, - ) -> vllm_tokenizer::Result { - Ok(token_ids - .iter() - .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) - .collect()) - } - - fn token_to_id(&self, token: &str) -> Option { - match token { - "" => Some(1), - "" => Some(2), - _ => None, - } - } - } - - struct FailingReasoningParser { - fail_next: bool, - } - - impl ReasoningParser for FailingReasoningParser { - fn create(_tokenizer: DynTokenizer) -> Result, ReasoningError> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { fail_next: true })) - } - - fn push(&mut self, _text: &str) -> Result { - if self.fail_next { - self.fail_next = false; - return Err(ReasoningError::MissingToken { - token: "".to_string(), - }); - } - Ok(ReasoningDelta::default()) - } - } - - fn test_reasoning_parser(factory: &mut ReasoningParserFactory) -> Box { - factory.register_parser::("failing"); - - factory.create("failing", Arc::new(FakeTokenizer)).unwrap() - } - - #[tokio::test] - async fn reasoning_parser_failure_falls_back_to_plain_text() { - let mut factory = ReasoningParserFactory::new(); - let events = stream::iter(vec![ - Ok(DecodedTextEvent::Start { - prompt_token_ids: vec![1, 2, 3].into(), - prompt_logprobs: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "abc".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "def".to_string(), - token_ids: vec![], - logprobs: None, - finished: Some(vllm_text::Finished { - usage: vllm_llm::TokenUsage { - prompt_token_count: 3, - output_token_count: 0, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - }), - ]); - - let collected = reasoning_event_stream(events, Some(test_reasoning_parser(&mut factory))) - .collect::>() - .await; - - let events = collected - .into_iter() - .collect::>>() - .expect("reasoning stream should not fail"); - - assert_eq!( - events, - vec![ - ContentEvent::Start { - prompt_token_ids: vec![1, 2, 3].into(), - prompt_logprobs: None, - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "abc".to_string(), - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "def".to_string(), - }, - ContentEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 3, - output_token_count: 0, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }, - ] - ); - } - - #[tokio::test] - async fn reasoning_stream_preserves_logprobs_delta() { - let events = stream::iter(vec![ - Ok(DecodedTextEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "abc".to_string(), - token_ids: vec![], - logprobs: Some(DecodedLogprobs { - positions: vec![DecodedPositionLogprobs { - entries: vec![DecodedTokenLogprob { - token_id: 0, - token: "a".to_string(), - logprob: -0.1, - rank: 1, - }], - }], - }), - finished: None, - }), - ]); - - let collected = reasoning_event_stream(events, None) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert_eq!( - collected, - vec![ - ContentEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "abc".to_string(), - }, - ContentEvent::LogprobsDelta { - logprobs: Some(DecodedLogprobs { - positions: vec![DecodedPositionLogprobs { - entries: vec![DecodedTokenLogprob { - token_id: 0, - token: "a".to_string(), - logprob: -0.1, - rank: 1, - }], - }], - }), - token_ids: vec![], - }, - ] - ); - } - - #[tokio::test] - async fn qwen3_parser_uses_prompt_end_marker_to_switch_to_content() { - let tokenizer = Arc::new(FakeTokenizer); - let events = stream::iter(vec![ - Ok(DecodedTextEvent::Start { - prompt_token_ids: vec![2].into(), - prompt_logprobs: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "thought ".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "doneOK".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - ]); - - let factory = ReasoningParserFactory::new(); - let collected = reasoning_event_stream( - events, - Some(factory.create(names::QWEN3, tokenizer).unwrap()), - ) - .collect::>() - .await; - - let events = collected - .into_iter() - .collect::>>() - .expect("reasoning stream should not fail"); - - assert_eq!( - events, - vec![ - ContentEvent::Start { - prompt_token_ids: vec![2].into(), - prompt_logprobs: None, - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "thought ".to_string(), - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "doneOK".to_string(), - }, - ] - ); - } - - #[tokio::test] - async fn qwen3_parser_tolerates_prompt_prefill_reasoning() { - let tokenizer = Arc::new(FakeTokenizer); - let events = stream::iter(vec![ - Ok(DecodedTextEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "thought ".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - Ok(DecodedTextEvent::TextDelta { - delta: "doneOK".to_string(), - token_ids: vec![], - logprobs: None, - finished: None, - }), - ]); - - let factory = ReasoningParserFactory::new(); - let collected = reasoning_event_stream( - events, - Some(factory.create(names::QWEN3, tokenizer).unwrap()), - ) - .collect::>() - .await; - - let events = collected - .into_iter() - .collect::>>() - .expect("reasoning stream should not fail"); - - assert_eq!( - events, - vec![ - ContentEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Reasoning, - delta: "thought ".to_string(), - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Reasoning, - delta: "done".to_string(), - }, - ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "OK".to_string(), - }, - ] - ); - } -} diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs index eb1fa3d1436..bdc5d8e14c2 100644 --- a/rust/src/chat/src/output/default/structural_tag.rs +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -77,7 +77,7 @@ fn structural_tag_tool_choice(request: &ChatRequest) -> Option, - /// Whether tool parsing has already failed for this stream. - parser_failed: bool, - /// The parser-local index of the currently open tool call, if any. - // NOTE: We only allow single open tool call at a time right now, since that's what all - // supported parsers currently emit. Change this to a `BTreeMap` if we need to support multiple - // interleaved calls in the future. - open_call_index: Option, -} - -impl ToolState { - /// Create one fresh tool-parsing state for a new streamed response. - fn new(parser: Box) -> Self { - Self { - parser, - parser_failed: false, - open_call_index: None, - } - } - - /// Convert one semantic assistant text delta into zero or more tool-aware - /// internal events. - fn process_text_delta( - &mut self, - kind: AssistantBlockKind, - delta: String, - ) -> Result> { - let mut events = Vec::new(); - - // Only normal assistant text is eligible for tool parsing. Reasoning - // blocks and plain-text fallback should pass through unchanged. - if kind != AssistantBlockKind::Text || self.parser_failed { - self.open_call_index = None; - events.push(AssistantEvent::TextDelta { kind, delta }); - return Ok(events); - } - - let mut output = ToolParserOutput::default(); - let parse_result = self.parser.parse_into(&delta, &mut output); - - match parse_result { - Ok(()) => self.process_parser_output(kind, output, &mut events)?, - Err(error) => { - warn!( - error = %error.as_report(), - "tool parser failed; falling back to plain text deltas" - ); - // Permanently mark this parser as failed. - // TODO: we may consider recovering from parsing errors in the future. - self.parser_failed = true; - - // On parsing failure, we still apply the partial parser output if any, but we close - // any open tool calls and emit the remaining buffered text as a plain-text delta to - // preserve as much of the output as possible. - self.process_parser_output(kind, output, &mut events)?; - self.open_call_index = None; - push_text_delta(&mut events, kind, self.parser.reset()); - } - } - - Ok(events) - } - - /// Apply one parsed tool output to the current stream state. - fn process_parser_output( - &mut self, - kind: AssistantBlockKind, - output: ToolParserOutput, - events: &mut Vec, - ) -> Result<()> { - // When we are not currently streaming a tool call, preserve plain - // text first and then surface any new tool call items. - if self.open_call_index.is_none() { - push_text_delta(events, kind, output.normal_text); - self.process_tool_items(output.calls, events)?; - } else { - // Once a tool call is open, prioritize tool deltas first. If the - // parser emits normal text again, close the tool call and resume - // plain text output. - self.process_tool_items(output.calls, events)?; - if !output.normal_text.is_empty() { - self.open_call_index = None; - push_text_delta(events, kind, output.normal_text); - } - } - Ok(()) - } - - /// Apply one batch of parsed tool-call deltas emitted by the parser. - fn process_tool_items( - &mut self, - items: Vec, - events: &mut Vec, - ) -> Result<()> { - for item in items { - if let Some(name) = item.name { - let is_new_tool = match self.open_call_index { - Some(open_call_index) => open_call_index != item.tool_index, - None => true, - }; - if is_new_tool { - let id = self - .parser - .tool_call_id(item.tool_index) - .map(str::to_string) - .unwrap_or_else(generate_tool_call_id); - self.open_call_index = Some(item.tool_index); - events.push(AssistantEvent::ToolCallStart { id, name }); - } - } - - if item.arguments.is_empty() { - // No arguments delta to apply. - continue; - } - let Some(open_call_index) = self.open_call_index else { - return Err(Error::ToolCallStreamInvariant { - message: format!( - "received arguments for tool index {} before any tool-call start", - item.tool_index - ), - }); - }; - if open_call_index != item.tool_index { - return Err(Error::ToolCallStreamInvariant { - message: format!( - "received arguments for tool index {} while tool index {} is open", - item.tool_index, open_call_index - ), - }); - } - - events.push(AssistantEvent::ToolCallArgumentsDelta { - delta: item.arguments, - }); - } - Ok(()) - } - - /// Flush parser state at end-of-stream and close any remaining open calls. - fn finish(&mut self) -> Result> { - let mut events = Vec::new(); - - if self.parser_failed { - return Ok(events); - } - - match self.parser.finish() { - Ok(output) => { - self.process_parser_output(AssistantBlockKind::Text, output, &mut events)? - } - Err(error) => { - warn!( - error = %error.as_report(), - "tool parser finish failed; closing open tool calls with buffered state" - ); - self.parser_failed = true; - } - } - - Ok(events) - } -} - -/// Push one plain-text delta if it is non-empty. -fn push_text_delta(events: &mut Vec, kind: AssistantBlockKind, delta: String) { - if delta.is_empty() { - return; - } - events.push(AssistantEvent::TextDelta { kind, delta }); -} - -/// Wrap one semantic assistant stream into the internal tool-aware assistant -/// stream. -#[try_stream] -pub(crate) async fn tool_event_stream( - stream: impl ContentEventStream, - parser: Option>, - mut y: TryYielder, -) -> Result<()> { - // Without a parser, pass through the input stream unchanged. - let Some(parser) = parser else { - pin_mut!(stream); - while let Some(event) = stream.next().await.transpose()? { - y.yield_ok(event.into()).await; - } - return Ok(()); - }; - - pin_mut!(stream); - let mut state = ToolState::new(parser); - - while let Some(event) = stream.next().await.transpose()? { - match event { - ContentEvent::Start { - prompt_token_ids, - prompt_logprobs, - } => { - y.yield_ok(AssistantEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) - .await; - } - ContentEvent::TextDelta { kind, delta } => { - for next in state.process_text_delta(kind, delta)? { - y.yield_ok(next).await; - } - } - ContentEvent::LogprobsDelta { - logprobs, - token_ids, - } => { - y.yield_ok(AssistantEvent::LogprobsDelta { - logprobs, - token_ids, - }) - .await; - } - ContentEvent::Done { - usage, - finish_reason, - kv_transfer_params, - } => { - for next in state.finish()? { - y.yield_ok(next).await; - } - - y.yield_ok(AssistantEvent::Done { - usage, - finish_reason, - kv_transfer_params, - }) - .await; - } - } - } - Ok(()) -} - -#[cfg(test)] -mod tests { - - use futures::{StreamExt as _, stream}; - use vllm_llm::FinishReason; - use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob}; - use vllm_tool_parser::Result; - - use super::super::{AssistantEvent, ContentEvent}; - use super::tool_event_stream; - use crate::error::Error; - use crate::event::{AssistantBlockKind, AssistantMessageExt as _}; - use crate::output::structured::structured_chat_event_stream; - use crate::parser::tool::{ - DeepSeekV4ToolParser, ToolParser, ToolParserError, ToolParserOutput, - }; - use crate::request::ChatTool; - use crate::stream::{ChatEventStream, CollectedAssistantMessage}; - - struct FailingParser { - fail_next: bool, - buffered: String, - } - - struct ScriptedParser { - push_outputs: Vec, - finish_output: ToolParserOutput, - } - - struct PartialThenFailParser { - buffered: String, - } - - struct IdScriptedParser { - output: ToolParserOutput, - tool_call_id: Option, - } - - impl ToolParser for FailingParser { - fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { - fail_next: false, - buffered: String::new(), - })) - } - - fn parse_into(&mut self, chunk: &str, _output: &mut ToolParserOutput) -> Result<()> { - self.buffered.push_str(chunk); - if self.fail_next { - self.fail_next = false; - return Err(ToolParserError::ParsingFailed { - message: "boom".to_string(), - }); - } - - self.buffered.clear(); - Ok(()) - } - - fn finish(&mut self) -> Result { - Ok(ToolParserOutput::default()) - } - - fn reset(&mut self) -> String { - std::mem::take(&mut self.buffered) - } - } - - impl ToolParser for ScriptedParser { - fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { - push_outputs: Vec::new(), - finish_output: ToolParserOutput::default(), - })) - } - - fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> { - let mut next = self.push_outputs.pop().unwrap_or_default(); - output.normal_text.push_str(&next.normal_text); - output.calls.append(&mut next.calls); - Ok(()) - } - - fn finish(&mut self) -> Result { - Ok(std::mem::take(&mut self.finish_output)) - } - - fn reset(&mut self) -> String { - String::new() - } - } - - impl ToolParser for IdScriptedParser { - fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { - output: ToolParserOutput::default(), - tool_call_id: None, - })) - } - - fn tool_call_id(&self, tool_index: usize) -> Option<&str> { - (tool_index == 0).then_some(self.tool_call_id.as_deref()).flatten() - } - - fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> { - output.append(std::mem::take(&mut self.output)); - Ok(()) - } - - fn finish(&mut self) -> Result { - Ok(ToolParserOutput::default()) - } - - fn reset(&mut self) -> String { - String::new() - } - } - - impl ToolParser for PartialThenFailParser { - fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self { - buffered: String::new(), - })) - } - - fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> { - output.calls.extend([ - crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("get_weather".to_string()), - arguments: String::new(), - }, - crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: None, - arguments: r#"{"location":"SF"}"#.to_string(), - }, - ]); - self.buffered.push_str(" trailing text"); - Err(ToolParserError::ParsingFailed { - message: "boom".to_string(), - }) - } - - fn finish(&mut self) -> Result { - Ok(ToolParserOutput::default()) - } - - fn reset(&mut self) -> String { - std::mem::take(&mut self.buffered) - } - } - - fn deepseek_v4_test_tools() -> Vec { - vec![ - ChatTool { - name: "get_weather".to_string(), - description: None, - parameters: serde_json::json!({ - "type": "object", - "properties": { - "location": { "type": "string" } - } - }), - strict: None, - }, - ChatTool { - name: "add".to_string(), - description: None, - parameters: serde_json::json!({ - "type": "object", - "properties": { - "x": { "type": "integer" }, - "y": { "type": "integer" } - } - }), - strict: None, - }, - ] - } - - async fn collect_deepseek_v4_message(chunks: Vec) -> CollectedAssistantMessage { - let events = chunks - .into_iter() - .map(|delta| { - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta, - }) - }) - .chain(std::iter::once(Ok(ContentEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 1, - output_token_count: 1, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }))); - let parser = DeepSeekV4ToolParser::create(&deepseek_v4_test_tools()).unwrap(); - let assistant_events = tool_event_stream(stream::iter(events), Some(parser)); - let chat_events = structured_chat_event_stream(assistant_events, true); - - ChatEventStream::new("req_deepseek_v4".to_string(), Box::pin(chat_events)) - .collect_message() - .await - .unwrap() - } - - fn message_tool_projection( - message: &CollectedAssistantMessage, - ) -> (String, Vec<(String, serde_json::Value)>) { - ( - message.message.text(), - message - .message - .tool_calls() - .map(|call| { - ( - call.name.clone(), - serde_json::from_str(&call.arguments).unwrap(), - ) - }) - .collect(), - ) - } - - #[tokio::test] - async fn tool_parser_error_preserves_partial_output_and_flushes_buffer() { - let events = stream::iter(vec![ - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - }), - Ok(ContentEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 1, - output_token_count: 1, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - - let events = tool_event_stream( - events, - Some(Box::new(PartialThenFailParser { - buffered: String::new(), - })), - ) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert!(matches!( - &events[0], - AssistantEvent::ToolCallStart { name, .. } if name == "get_weather" - )); - assert!(matches!( - &events[1], - AssistantEvent::ToolCallArgumentsDelta { delta } if delta == r#"{"location":"SF"}"# - )); - assert_eq!( - events[2], - AssistantEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: " trailing text".to_string(), - } - ); - assert!(matches!(events[3], AssistantEvent::Done { .. })); - } - - #[tokio::test] - async fn tool_stream_preserves_parser_provided_tool_call_id() { - let events = stream::iter(vec![Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - })]); - let parser = IdScriptedParser { - output: ToolParserOutput { - normal_text: String::new(), - calls: vec![crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("get_weather".to_string()), - arguments: "{}".to_string(), - }], - }, - tool_call_id: Some("functions.get_weather:0".to_string()), - }; - - let events = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert!(matches!( - &events[0], - AssistantEvent::ToolCallStart { id, name } - if id == "functions.get_weather:0" && name == "get_weather" - )); - } - - #[tokio::test] - async fn tool_stream_generates_tool_call_id_when_parser_omits_one() { - let events = stream::iter(vec![Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - })]); - let parser = IdScriptedParser { - output: ToolParserOutput { - normal_text: String::new(), - calls: vec![crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("get_weather".to_string()), - arguments: "{}".to_string(), - }], - }, - tool_call_id: None, - }; - - let events = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert!(matches!( - &events[0], - AssistantEvent::ToolCallStart { id, name } - if id.starts_with("call_") && name == "get_weather" - )); - } - - #[tokio::test] - async fn real_buffered_parser_error_matches_streaming_and_non_streaming() { - let prefix = "I will check both.\n"; - let first_tool_call = concat!( - "<｜DSML｜tool_calls>\n", - "<｜DSML｜invoke name=\"get_weather\">\n", - "<｜DSML｜parameter name=\"location\" string=\"true\">Tokyo\n", - "", - ); - let malformed_second_tool_call = concat!( - "\n<｜DSML｜invoke name=\"add\">\n", - "not a parameter\n", - "\n", - "", - ); - let streaming_chunks = vec![ - prefix.to_string(), - first_tool_call.to_string(), - malformed_second_tool_call.to_string(), - ]; - let full_output = streaming_chunks.concat(); - - let streaming = collect_deepseek_v4_message(streaming_chunks).await; - let non_streaming = collect_deepseek_v4_message(vec![full_output]).await; - - let expected = ( - format!("{prefix}{malformed_second_tool_call}"), - vec![( - "get_weather".to_string(), - serde_json::json!({ "location": "Tokyo" }), - )], - ); - assert_eq!(message_tool_projection(&streaming), expected); - assert_eq!(message_tool_projection(&non_streaming), expected); - } - - #[tokio::test] - async fn tool_parser_failure_falls_back_to_plain_text() { - let events = stream::iter(vec![ - Ok(ContentEvent::Start { - prompt_token_ids: vec![1, 2, 3].into(), - prompt_logprobs: None, - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "abc".to_string(), - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "def".to_string(), - }), - Ok(ContentEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 3, - output_token_count: 0, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - - let collected = tool_event_stream( - events, - Some(Box::new(FailingParser { - fail_next: true, - buffered: String::new(), - })), - ) - .collect::>() - .await; - - let events = collected - .into_iter() - .collect::>>() - .expect("tool stream should not fail"); - - assert_eq!( - events, - vec![ - AssistantEvent::Start { - prompt_token_ids: vec![1, 2, 3].into(), - prompt_logprobs: None, - }, - AssistantEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "abc".to_string(), - }, - AssistantEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "def".to_string(), - }, - AssistantEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 3, - output_token_count: 0, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }, - ] - ); - - let message = ChatEventStream::new( - "req_fallback".to_string(), - Box::pin(structured_chat_event_stream( - stream::iter(events.into_iter().map(Ok)), - true, - )), - ) - .collect_message() - .await - .expect("collect_message should succeed"); - assert_eq!(message.message.text(), "abcdef"); - assert!(message.message.tool_calls().next().is_none()); - } - - #[tokio::test] - async fn tool_stream_preserves_logprobs_delta() { - let events = stream::iter(vec![ - Ok(ContentEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }), - Ok(ContentEvent::LogprobsDelta { - logprobs: Some(DecodedLogprobs { - positions: vec![DecodedPositionLogprobs { - entries: vec![DecodedTokenLogprob { - token_id: 0, - token: "a".to_string(), - logprob: -0.2, - rank: 1, - }], - }], - }), - token_ids: vec![], - }), - Ok(ContentEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 1, - output_token_count: 0, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - let events = tool_event_stream( - events, - Some(Box::new(FailingParser { - fail_next: false, - buffered: String::new(), - })), - ) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert_eq!( - events, - vec![ - AssistantEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }, - AssistantEvent::LogprobsDelta { - logprobs: Some(DecodedLogprobs { - positions: vec![DecodedPositionLogprobs { - entries: vec![DecodedTokenLogprob { - token_id: 0, - token: "a".to_string(), - logprob: -0.2, - rank: 1, - }], - }], - }), - token_ids: vec![], - }, - AssistantEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 1, - output_token_count: 0, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }, - ] - ); - } - - #[tokio::test] - async fn tool_stream_rejects_interleaved_tool_indices() { - let events = stream::iter(vec![ - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - }), - Ok(ContentEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 1, - output_token_count: 1, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - - let parser = ScriptedParser { - push_outputs: vec![ToolParserOutput { - normal_text: String::new(), - calls: vec![ - crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("first".to_string()), - arguments: String::new(), - }, - crate::parser::tool::ToolCallDelta { - tool_index: 1, - name: None, - arguments: "{}".to_string(), - }, - ], - }], - finish_output: ToolParserOutput::default(), - }; - - let err = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .find_map(|output| output.err()) - .expect("expected invariant error"); - - assert!(matches!(err, Error::ToolCallStreamInvariant { .. })); - } - - #[tokio::test] - async fn tool_stream_resets_open_tool_when_normal_text_interrupts_it() { - let events = stream::iter(vec![ - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "start".to_string(), - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "text".to_string(), - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "args".to_string(), - }), - ]); - - let parser = ScriptedParser { - push_outputs: vec![ - ToolParserOutput { - normal_text: String::new(), - calls: vec![crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: None, - arguments: "}".to_string(), - }], - }, - ToolParserOutput { - normal_text: "plain text".to_string(), - calls: Vec::new(), - }, - ToolParserOutput { - normal_text: String::new(), - calls: vec![crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("first".to_string()), - arguments: "{".to_string(), - }], - }, - ], - finish_output: ToolParserOutput::default(), - }; - - let err = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .find_map(|output| output.err()) - .expect("expected invariant error"); - - assert!(matches!( - err, - Error::ToolCallStreamInvariant { message } - if message == "received arguments for tool index 0 before any tool-call start" - )); - } - - #[tokio::test] - async fn tool_stream_emits_start_and_args_for_terminal_text() { - let events = stream::iter(vec![ - Ok(ContentEvent::Start { - prompt_token_ids: vec![1].into(), - prompt_logprobs: None, - }), - Ok(ContentEvent::TextDelta { - kind: AssistantBlockKind::Text, - delta: "ignored".to_string(), - }), - Ok(ContentEvent::Done { - usage: vllm_llm::TokenUsage { - prompt_token_count: 1, - output_token_count: 1, - cached_token_count: 0, - }, - finish_reason: FinishReason::stop_eos(), - kv_transfer_params: None, - }), - ]); - - let parser = ScriptedParser { - push_outputs: vec![ToolParserOutput { - normal_text: String::new(), - calls: vec![ - crate::parser::tool::ToolCallDelta { - tool_index: 0, - name: Some("first".to_string()), - arguments: r#"{"a":1}"#.to_string(), - }, - crate::parser::tool::ToolCallDelta { - tool_index: 1, - name: Some("second".to_string()), - arguments: r#"{"b":2}"#.to_string(), - }, - ], - }], - finish_output: ToolParserOutput::default(), - }; - - let events = tool_event_stream(events, Some(Box::new(parser))) - .collect::>() - .await - .into_iter() - .collect::>>() - .unwrap(); - - assert!(matches!(events[1], AssistantEvent::ToolCallStart { .. })); - assert!(matches!( - events[2], - AssistantEvent::ToolCallArgumentsDelta { .. } - )); - assert!(matches!(events[3], AssistantEvent::ToolCallStart { .. })); - assert!(matches!( - events[4], - AssistantEvent::ToolCallArgumentsDelta { .. } - )); - let collected = ChatEventStream::new( - "req_final_only".to_string(), - Box::pin(structured_chat_event_stream( - stream::iter(events.into_iter().map(Ok)), - true, - )), - ) - .collect_message() - .await - .unwrap(); - let tool_calls = collected.message.tool_calls().collect::>(); - assert_eq!(tool_calls.len(), 2); - assert_eq!(tool_calls[0].name, "first"); - assert_eq!(tool_calls[1].name, "second"); - } -} diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs new file mode 100644 index 00000000000..66d2c5a204e --- /dev/null +++ b/rust/src/chat/src/output/default/unified.rs @@ -0,0 +1,626 @@ +//! Adapts decoded text updates into parsed assistant deltas. +//! +//! This stage sits between low-level token decoding and final block assembly. +//! It drives one unified parser that may emit normal text, reasoning text, or +//! tool-call deltas, then normalizes those parser events into internal +//! assistant events. + +use asynk_strim_attr::{TryYielder, try_stream}; +use futures::{StreamExt as _, pin_mut}; +use thiserror_ext::AsReport; +use tracing::warn; +use vllm_parser::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; +use vllm_text::output::DecodedTextEvent; + +use crate::Result; +use crate::error::Error; +use crate::event::AssistantBlockKind; +use crate::output::{AssistantEvent, DecodedTextEventStream, generate_tool_call_id}; + +/// Per-stream unified parsing state. +struct UnifiedParserState { + /// Parser for the current request stream. + parser: Box, + /// Whether unified parsing has already failed for this stream. + parser_failed: bool, + /// The parser-local index of the currently open tool call, if any. + /// + /// Supported parsers currently emit at most one active tool call at a time. + /// Change this to an indexed map if a model needs interleaved calls later. + open_call_index: Option, +} + +impl UnifiedParserState { + /// Create one fresh unified parsing state for a new streamed response. + fn new(parser: Box) -> Self { + Self { + parser, + parser_failed: false, + open_call_index: None, + } + } + + /// Initialize parser state once prompt token IDs are available. + fn initialize(&mut self, prompt_token_ids: &[u32]) { + if self.parser_failed { + return; + } + + match self.parser.initialize(prompt_token_ids) { + Ok(()) => {} + Err(error) => { + warn!( + error = %error.as_report(), + "failed to initialize unified parser; falling back to plain text deltas" + ); + self.parser_failed = true; + self.open_call_index = None; + } + } + } + + /// Convert one decoded text delta into zero or more parsed assistant events. + fn process_delta(&mut self, delta: String) -> Result> { + if self.parser_failed { + self.open_call_index = None; + return Ok(text_event(AssistantBlockKind::Text, delta).into_iter().collect()); + } + + let mut output = UnifiedParserOutput::default(); + match self.parser.parse_into(&delta, &mut output) { + Ok(()) => { + let mut events = Vec::new(); + self.process_parser_output(output, &mut events)?; + Ok(events) + } + Err(error) => { + warn!( + error = %error.as_report(), + "unified parser failed; falling back to plain text deltas" + ); + self.parser_failed = true; + + let mut events = Vec::new(); + self.process_parser_output(output, &mut events)?; + self.open_call_index = None; + + let recovered = self.parser.reset(); + if recovered.is_empty() && events.is_empty() { + push_text_delta(&mut events, AssistantBlockKind::Text, delta); + } else { + push_text_delta(&mut events, AssistantBlockKind::Text, recovered); + } + Ok(events) + } + } + } + + /// Flush parser state at end-of-stream and close any remaining open calls. + fn finish(&mut self) -> Result> { + let mut events = Vec::new(); + + if self.parser_failed { + return Ok(events); + } + + match self.parser.finish() { + Ok(output) => self.process_parser_output(output, &mut events)?, + Err(error) => { + warn!( + error = %error.as_report(), + "unified parser finish failed; closing open parser state" + ); + self.parser_failed = true; + self.open_call_index = None; + // TODO: should we reset and emit the buffered text? + } + } + + Ok(events) + } + + /// Apply one parsed unified output to the current stream state. + fn process_parser_output( + &mut self, + output: UnifiedParserOutput, + events: &mut Vec, + ) -> Result<()> { + for event in output.events { + match event { + UnifiedParserEvent::Text(delta) => { + self.open_call_index = None; + push_text_delta(events, AssistantBlockKind::Text, delta); + } + UnifiedParserEvent::Reasoning(delta) => { + self.open_call_index = None; + push_text_delta(events, AssistantBlockKind::Reasoning, delta); + } + UnifiedParserEvent::ToolCall(item) => { + self.process_tool_item(item, events)?; + } + } + } + + Ok(()) + } + + /// Apply one parsed tool-call delta emitted by the parser. + fn process_tool_item( + &mut self, + item: vllm_parser::tool::ToolCallDelta, + events: &mut Vec, + ) -> Result<()> { + if let Some(name) = item.name { + let is_new_tool = match self.open_call_index { + Some(open_call_index) => open_call_index != item.tool_index, + None => true, + }; + if is_new_tool { + let id = self + .parser + .tool_call_id(item.tool_index) + .map(str::to_string) + .unwrap_or_else(generate_tool_call_id); + self.open_call_index = Some(item.tool_index); + events.push(AssistantEvent::ToolCallStart { id, name }); + } + } + + if item.arguments.is_empty() { + return Ok(()); + } + let Some(open_call_index) = self.open_call_index else { + return Err(Error::ToolCallStreamInvariant { + message: format!( + "received arguments for tool index {} before any tool-call start", + item.tool_index + ), + }); + }; + if open_call_index != item.tool_index { + return Err(Error::ToolCallStreamInvariant { + message: format!( + "received arguments for tool index {} while tool index {} is open", + item.tool_index, open_call_index + ), + }); + } + + events.push(AssistantEvent::ToolCallArgumentsDelta { + delta: item.arguments, + }); + Ok(()) + } +} + +/// Build one plain text event if `delta` is non-empty. +fn text_event(kind: AssistantBlockKind, delta: String) -> Option { + if delta.is_empty() { + return None; + } + Some(AssistantEvent::TextDelta { kind, delta }) +} + +/// Push one plain text delta if it is non-empty. +fn push_text_delta(events: &mut Vec, kind: AssistantBlockKind, delta: String) { + if let Some(event) = text_event(kind, delta) { + events.push(event); + } +} + +/// Wrap one decoded-text stream into the internal unified assistant stream. +#[try_stream] +pub(crate) async fn unified_event_stream( + decoded_stream: impl DecodedTextEventStream, + parser: Box, + mut y: TryYielder, +) -> Result<()> { + pin_mut!(decoded_stream); + + let mut state = UnifiedParserState::new(parser); + + while let Some(event) = decoded_stream.next().await.transpose()? { + match event { + DecodedTextEvent::Start { + prompt_token_ids, + prompt_logprobs, + } => { + state.initialize(&prompt_token_ids); + y.yield_ok(AssistantEvent::Start { + prompt_token_ids, + prompt_logprobs, + }) + .await; + } + DecodedTextEvent::TextDelta { + delta, + token_ids, + logprobs, + finished, + } => { + for next in state.process_delta(delta)? { + y.yield_ok(next).await; + } + if logprobs.is_some() || !token_ids.is_empty() { + y.yield_ok(AssistantEvent::LogprobsDelta { + logprobs, + token_ids, + }) + .await; + } + if let Some(finished) = finished { + for next in state.finish()? { + y.yield_ok(next).await; + } + y.yield_ok(AssistantEvent::Done { + usage: finished.usage, + finish_reason: finished.finish_reason, + kv_transfer_params: finished.kv_transfer_params, + }) + .await; + } + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::collections::VecDeque; + + use futures::{StreamExt as _, stream}; + use vllm_parser::reasoning::ReasoningError; + use vllm_parser::tool::ToolCallDelta; + use vllm_parser::unified::{UnifiedParserError, UnifiedParserOutput}; + + use super::unified_event_stream; + use crate::event::AssistantBlockKind; + use crate::output::AssistantEvent; + + enum ScriptedStep { + Output(UnifiedParserOutput), + Error { + committed: UnifiedParserOutput, + reset_text: String, + }, + } + + struct ScriptedParser { + steps: VecDeque, + reset_text: String, + tool_call_id: Option, + finish_error_reset_text: Option, + } + + impl ScriptedParser { + fn new(steps: impl IntoIterator) -> Self { + Self { + steps: steps.into_iter().collect(), + reset_text: String::new(), + tool_call_id: Some("call_test".to_string()), + finish_error_reset_text: None, + } + } + + fn with_finish_error(mut self, reset_text: &str) -> Self { + self.finish_error_reset_text = Some(reset_text.to_string()); + self + } + } + + impl vllm_parser::unified::UnifiedParser for ScriptedParser { + fn create( + _tools: &[vllm_parser::tool::Tool], + _tokenizer: vllm_tokenizer::DynTokenizer, + ) -> vllm_parser::unified::Result> + where + Self: Sized + 'static, + { + unreachable!("ScriptedParser is constructed directly in tests") + } + + fn parse_into( + &mut self, + _delta: &str, + output: &mut UnifiedParserOutput, + ) -> vllm_parser::unified::Result<()> { + match self.steps.pop_front().expect("unexpected parser call") { + ScriptedStep::Output(next) => { + output.append(next); + Ok(()) + } + ScriptedStep::Error { + committed, + reset_text, + } => { + output.append(committed); + self.reset_text = reset_text; + Err(UnifiedParserError::Reasoning( + ReasoningError::MissingToken { + token: "".to_string(), + }, + )) + } + } + } + + fn tool_call_id(&self, _tool_index: usize) -> Option<&str> { + self.tool_call_id.as_deref() + } + + fn finish(&mut self) -> vllm_parser::unified::Result { + if let Some(reset_text) = self.finish_error_reset_text.take() { + self.reset_text = reset_text; + return Err(UnifiedParserError::Reasoning( + ReasoningError::MissingToken { + token: "".to_string(), + }, + )); + } + Ok(UnifiedParserOutput::default()) + } + + fn reset(&mut self) -> String { + std::mem::take(&mut self.reset_text) + } + } + + fn decoded_delta(delta: &str) -> vllm_text::output::DecodedTextEvent { + vllm_text::output::DecodedTextEvent::TextDelta { + delta: delta.to_string(), + token_ids: Vec::new(), + logprobs: None, + finished: None, + } + } + + fn finished_delta(delta: &str) -> vllm_text::output::DecodedTextEvent { + vllm_text::output::DecodedTextEvent::TextDelta { + delta: delta.to_string(), + token_ids: Vec::new(), + logprobs: None, + finished: Some(vllm_text::output::Finished { + usage: vllm_llm::TokenUsage::default(), + finish_reason: crate::FinishReason::Stop(None), + kv_transfer_params: None, + }), + } + } + + async fn collect( + parser: ScriptedParser, + events: Vec, + ) -> Vec { + let stream = stream::iter(events.into_iter().map(Ok)); + unified_event_stream(stream, Box::new(parser)) + .collect::>() + .await + .into_iter() + .collect::>>() + .unwrap() + } + + fn text(delta: &str) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + output.push_text(delta.to_string()); + output + } + + fn reasoning(delta: &str) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + output.push_reasoning(delta.to_string()); + output + } + + fn tool_call(name: &str, arguments: &str) -> UnifiedParserOutput { + UnifiedParserOutput { + events: vec![vllm_parser::unified::UnifiedParserEvent::ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some(name.to_string()), + arguments: arguments.to_string(), + }, + )], + } + } + + fn tool_call_arguments(arguments: &str) -> UnifiedParserOutput { + UnifiedParserOutput { + events: vec![vllm_parser::unified::UnifiedParserEvent::ToolCall( + ToolCallDelta { + tool_index: 0, + name: None, + arguments: arguments.to_string(), + }, + )], + } + } + + fn combined(first: UnifiedParserOutput, second: UnifiedParserOutput) -> UnifiedParserOutput { + let mut output = first; + output.append(second); + output + } + + #[tokio::test] + async fn unified_stream_emits_reasoning_only_deltas() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(reasoning("thinking"))]), + vec![decoded_delta("raw")], + ) + .await; + + assert_eq!( + events, + vec![AssistantEvent::TextDelta { + kind: AssistantBlockKind::Reasoning, + delta: "thinking".to_string(), + }] + ); + } + + #[tokio::test] + async fn unified_stream_emits_tool_only_deltas() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(tool_call( + "get_weather", + r#"{"location":"Paris"}"#, + ))]), + vec![decoded_delta("raw")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::ToolCallStart { + id: "call_test".to_string(), + name: "get_weather".to_string(), + }, + AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"location":"Paris"}"#.to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_emits_reasoning_followed_by_tool_call() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(combined( + reasoning("thinking"), + tool_call("get_weather", r#"{"location":"Paris"}"#), + ))]), + vec![decoded_delta("raw")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Reasoning, + delta: "thinking".to_string(), + }, + AssistantEvent::ToolCallStart { + id: "call_test".to_string(), + name: "get_weather".to_string(), + }, + AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"location":"Paris"}"#.to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_emits_visible_text_followed_by_tool_call() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(combined( + text("visible "), + tool_call("get_weather", r#"{"location":"Paris"}"#), + ))]), + vec![decoded_delta("raw")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "visible ".to_string(), + }, + AssistantEvent::ToolCallStart { + id: "call_test".to_string(), + name: "get_weather".to_string(), + }, + AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"location":"Paris"}"#.to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_emits_tool_arguments_before_trailing_text() { + let events = collect( + ScriptedParser::new([ + ScriptedStep::Output(tool_call("get_weather", "")), + ScriptedStep::Output(combined( + tool_call_arguments(r#"{"location":"Paris"}"#), + text(" done"), + )), + ]), + vec![decoded_delta("start"), decoded_delta("finish")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::ToolCallStart { + id: "call_test".to_string(), + name: "get_weather".to_string(), + }, + AssistantEvent::ToolCallArgumentsDelta { + delta: r#"{"location":"Paris"}"#.to_string(), + }, + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: " done".to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_fallback_keeps_committed_output_and_disables_later_parsing() { + let events = collect( + ScriptedParser::new([ScriptedStep::Error { + committed: text("committed"), + reset_text: "buffered".to_string(), + }]), + vec![decoded_delta("bad"), decoded_delta("later")], + ) + .await; + + assert_eq!( + events, + vec![ + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "committed".to_string(), + }, + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "buffered".to_string(), + }, + AssistantEvent::TextDelta { + kind: AssistantBlockKind::Text, + delta: "later".to_string(), + }, + ] + ); + } + + #[tokio::test] + async fn unified_stream_finish_error_closes_parser_without_reset_text() { + let events = collect( + ScriptedParser::new([ScriptedStep::Output(UnifiedParserOutput::default())]) + .with_finish_error("buffered"), + vec![finished_delta("")], + ) + .await; + + assert_eq!( + events, + vec![AssistantEvent::Done { + usage: vllm_llm::TokenUsage::default(), + finish_reason: crate::FinishReason::Stop(None), + kv_transfer_params: None, + }] + ); + } +} diff --git a/rust/src/chat/src/output/mod.rs b/rust/src/chat/src/output/mod.rs index d7b73c4e5e2..836b199eb9b 100644 --- a/rust/src/chat/src/output/mod.rs +++ b/rust/src/chat/src/output/mod.rs @@ -2,7 +2,6 @@ use std::pin::Pin; use std::sync::Arc; use futures::Stream; -use subenum::subenum; use trait_set::trait_set; use uuid::Uuid; use vllm_llm::TokenUsage; @@ -22,23 +21,19 @@ pub(crate) use harmony::validate_harmony_parser_overrides; /// Internal assistant event before final assembly. /// -/// - [`ContentEvent`]: subenum after reasoning parsing, carries only text content. -/// - [`AssistantEvent`]: full event after tool parsing, adds tool-call variants. -#[subenum(ContentEvent)] +/// Unified parsing produces these events, and structured assembly consumes +/// them to build public chat events. #[derive(Debug, Clone, PartialEq)] pub(crate) enum AssistantEvent { - #[subenum(ContentEvent)] Start { prompt_token_ids: Arc<[u32]>, prompt_logprobs: Option, }, - #[subenum(ContentEvent)] TextDelta { kind: AssistantBlockKind, delta: String, }, /// Per-decoded-update sample metadata: logprobs and/or output token IDs. - #[subenum(ContentEvent)] LogprobsDelta { logprobs: Option, token_ids: Vec, @@ -48,7 +43,6 @@ pub(crate) enum AssistantEvent { /// A delta for the arguments of the currently open tool call. Must follow a /// `ToolCallStart`. ToolCallArgumentsDelta { delta: String }, - #[subenum(ContentEvent)] Done { usage: TokenUsage, finish_reason: FinishReason, @@ -57,50 +51,6 @@ pub(crate) enum AssistantEvent { }, } -impl ContentEvent { - /// Convert a [`DecodedTextEvent`] into one or more [`ContentEvent`] values - /// by treating all text as plain (non-reasoning) content. - fn from_decoded_plain_text(event: DecodedTextEvent) -> Vec { - match event { - DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - } => vec![Self::Start { - prompt_token_ids, - prompt_logprobs, - }], - DecodedTextEvent::TextDelta { - delta, - token_ids, - logprobs, - finished, - } => { - let mut events = Vec::new(); - if !delta.is_empty() { - events.push(Self::TextDelta { - kind: AssistantBlockKind::Text, - delta, - }); - } - if logprobs.is_some() || !token_ids.is_empty() { - events.push(Self::LogprobsDelta { - logprobs, - token_ids, - }); - } - if let Some(finished) = finished { - events.push(Self::Done { - usage: finished.usage, - finish_reason: finished.finish_reason, - kv_transfer_params: finished.kv_transfer_params, - }); - } - events - } - } - } -} - /// Boxed stream of decoded text events coming from [`vllm_text`]. pub type DynDecodedTextEventStream = Pin> + Send>>; /// Boxed stream of structured chat events exposed by [`crate::ChatLlm`]. diff --git a/rust/src/chat/src/parser/reasoning/mod.rs b/rust/src/chat/src/parser/reasoning/mod.rs index 7de8a9d5fa1..a414a8e3ba5 100644 --- a/rust/src/chat/src/parser/reasoning/mod.rs +++ b/rust/src/chat/src/parser/reasoning/mod.rs @@ -2,7 +2,7 @@ use std::sync::LazyLock; -pub use vllm_reasoning_parser::{ +pub use vllm_parser::reasoning::{ CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser, DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser, KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser, @@ -34,7 +34,7 @@ pub mod names { /// Constructor signature for one registered reasoning parser implementation. type ReasoningParserCreator = - fn(DynTokenizer) -> vllm_reasoning_parser::Result>; + fn(DynTokenizer) -> vllm_parser::reasoning::Result>; /// Registry and model matcher for reasoning parsers. pub type ReasoningParserFactory = ParserFactory; diff --git a/rust/src/chat/src/parser/tool/mod.rs b/rust/src/chat/src/parser/tool/mod.rs index 7561aa071ac..9884d1aca2a 100644 --- a/rust/src/chat/src/parser/tool/mod.rs +++ b/rust/src/chat/src/parser/tool/mod.rs @@ -2,13 +2,12 @@ use std::sync::LazyLock; -pub use vllm_tool_parser::{ +pub use vllm_parser::tool::{ DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser, Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser, HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser, - Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError, - ToolParserOutput, + Qwen3CoderToolParser, Qwen3XmlToolParser, ToolParser, ToolParserError, }; use crate::parser::ParserFactory; @@ -41,7 +40,7 @@ pub mod names { } /// Constructor signature for one registered tool parser implementation. -type ToolParserCreator = fn(&[ChatTool]) -> vllm_tool_parser::Result>; +type ToolParserCreator = fn(&[ChatTool]) -> vllm_parser::tool::Result>; /// Registry and model matcher for tool parsers. pub type ToolParserFactory = ParserFactory; diff --git a/rust/src/chat/src/parser/tool/tests.rs b/rust/src/chat/src/parser/tool/tests.rs index c40500adc74..a630f9a951a 100644 --- a/rust/src/chat/src/parser/tool/tests.rs +++ b/rust/src/chat/src/parser/tool/tests.rs @@ -1,6 +1,6 @@ -use vllm_tool_parser::Result; +use vllm_parser::tool::{Result, ToolParserOutput}; -use super::{ToolParser, ToolParserFactory, ToolParserOutput, names}; +use super::{ToolParser, ToolParserFactory, names}; use crate::Error; use crate::request::ChatTool; diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index 51b2efebd41..72de3d87663 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -4,9 +4,9 @@ use llm_multimodal::ImageDetail; use serde::{Deserialize, Serialize}; use serde_json::Value; use vllm_engine_core_client::protocol::lora::LoraRequest; +pub use vllm_parser::tool::Tool as ChatTool; pub use vllm_text::SamplingParams; use vllm_text::TextDecodeOptions; -pub use vllm_tool_parser::Tool as ChatTool; use crate::AssistantMessageExt; use crate::error::{Error, Result}; diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index 611dbe23973..5c4a2c29b7d 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -1180,7 +1180,7 @@ async fn chat_stream_parses_tool_calls_automatically() { ), request_output( "chat-tool", - bytes_to_token_ids( + bytes_with_special_stop_token( b"\"arguments\":{\"city\":\"Paris\"}}\n", ), Some(EngineCoreFinishReason::Stop), diff --git a/rust/src/tool-parser/Cargo.toml b/rust/src/parser/Cargo.toml similarity index 96% rename from rust/src/tool-parser/Cargo.toml rename to rust/src/parser/Cargo.toml index c4363906aa0..67c74bb5601 100644 --- a/rust/src/tool-parser/Cargo.toml +++ b/rust/src/parser/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "vllm-tool-parser" +name = "vllm-parser" version.workspace = true edition.workspace = true license.workspace = true @@ -13,6 +13,7 @@ serde.workspace = true serde_json.workspace = true thiserror.workspace = true thiserror-ext.workspace = true +vllm-tokenizer.workspace = true winnow.workspace = true xgrammar-structural-tag.workspace = true diff --git a/rust/src/tool-parser/benches/deepseek_v3.rs b/rust/src/parser/benches/deepseek_v3.rs similarity index 96% rename from rust/src/tool-parser/benches/deepseek_v3.rs rename to rust/src/parser/benches/deepseek_v3.rs index 75d2e417ace..4d1ea337a76 100644 --- a/rust/src/tool-parser/benches/deepseek_v3.rs +++ b/rust/src/parser/benches/deepseek_v3.rs @@ -2,8 +2,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::DeepSeekParser as ExternalDeepSeekParser; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{DeepSeekV3ToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{DeepSeekV3ToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; diff --git a/rust/src/tool-parser/benches/deepseek_v31.rs b/rust/src/parser/benches/deepseek_v31.rs similarity index 96% rename from rust/src/tool-parser/benches/deepseek_v31.rs rename to rust/src/parser/benches/deepseek_v31.rs index bb6d029baff..a6f17c9f017 100644 --- a/rust/src/tool-parser/benches/deepseek_v31.rs +++ b/rust/src/parser/benches/deepseek_v31.rs @@ -2,8 +2,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::DeepSeek31Parser as ExternalDeepSeek31Parser; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{DeepSeekV31ToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{DeepSeekV31ToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; diff --git a/rust/src/tool-parser/benches/deepseek_v32.rs b/rust/src/parser/benches/deepseek_v32.rs similarity index 96% rename from rust/src/tool-parser/benches/deepseek_v32.rs rename to rust/src/parser/benches/deepseek_v32.rs index c7a8346120d..1e770d9b136 100644 --- a/rust/src/tool-parser/benches/deepseek_v32.rs +++ b/rust/src/parser/benches/deepseek_v32.rs @@ -1,8 +1,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{DeepSeekV32ToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{DeepSeekV32ToolParser, Tool, ToolParser}; mod utils; use utils::feed_parser; diff --git a/rust/src/tool-parser/benches/gemma4.rs b/rust/src/parser/benches/gemma4.rs similarity index 97% rename from rust/src/tool-parser/benches/gemma4.rs rename to rust/src/parser/benches/gemma4.rs index c4e8f966c2b..761f8d4e235 100644 --- a/rust/src/tool-parser/benches/gemma4.rs +++ b/rust/src/parser/benches/gemma4.rs @@ -1,8 +1,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{Gemma4ToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{Gemma4ToolParser, Tool, ToolParser}; mod utils; use utils::feed_parser; diff --git a/rust/src/tool-parser/benches/glm45_moe.rs b/rust/src/parser/benches/glm45_moe.rs similarity index 97% rename from rust/src/tool-parser/benches/glm45_moe.rs rename to rust/src/parser/benches/glm45_moe.rs index 8486885eceb..a55a9e83ac0 100644 --- a/rust/src/tool-parser/benches/glm45_moe.rs +++ b/rust/src/parser/benches/glm45_moe.rs @@ -2,8 +2,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::Glm4MoeParser as ExternalGlm4MoeParser; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{Glm45MoeToolParser, Glm47MoeToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{Glm45MoeToolParser, Glm47MoeToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; diff --git a/rust/src/tool-parser/benches/kimi_k2.rs b/rust/src/parser/benches/kimi_k2.rs similarity index 97% rename from rust/src/tool-parser/benches/kimi_k2.rs rename to rust/src/parser/benches/kimi_k2.rs index 5a80f660673..ab4c98399aa 100644 --- a/rust/src/tool-parser/benches/kimi_k2.rs +++ b/rust/src/parser/benches/kimi_k2.rs @@ -2,8 +2,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::KimiK2Parser as ExternalKimiK2Parser; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{KimiK2ToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{KimiK2ToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; diff --git a/rust/src/tool-parser/benches/llama3_json.rs b/rust/src/parser/benches/llama3_json.rs similarity index 96% rename from rust/src/tool-parser/benches/llama3_json.rs rename to rust/src/parser/benches/llama3_json.rs index 03b5b54ee78..1126daf6f7d 100644 --- a/rust/src/tool-parser/benches/llama3_json.rs +++ b/rust/src/parser/benches/llama3_json.rs @@ -2,8 +2,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::LlamaParser as ExternalLlamaParser; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{Llama3JsonToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{Llama3JsonToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; diff --git a/rust/src/tool-parser/benches/minimax_m2.rs b/rust/src/parser/benches/minimax_m2.rs similarity index 97% rename from rust/src/tool-parser/benches/minimax_m2.rs rename to rust/src/parser/benches/minimax_m2.rs index 4ad20400934..734d7437fbc 100644 --- a/rust/src/tool-parser/benches/minimax_m2.rs +++ b/rust/src/parser/benches/minimax_m2.rs @@ -2,8 +2,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::MinimaxM2Parser as ExternalMinimaxM2Parser; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{MinimaxM2ToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{MinimaxM2ToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; diff --git a/rust/src/tool-parser/benches/qwen3_coder.rs b/rust/src/parser/benches/qwen3_coder.rs similarity index 97% rename from rust/src/tool-parser/benches/qwen3_coder.rs rename to rust/src/parser/benches/qwen3_coder.rs index b4f26ac5cdb..9d70937728f 100644 --- a/rust/src/tool-parser/benches/qwen3_coder.rs +++ b/rust/src/parser/benches/qwen3_coder.rs @@ -2,8 +2,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::QwenCoderParser as ExternalQwenCoderParser; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{Qwen3CoderToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; diff --git a/rust/src/tool-parser/benches/qwen3_xml.rs b/rust/src/parser/benches/qwen3_xml.rs similarity index 96% rename from rust/src/tool-parser/benches/qwen3_xml.rs rename to rust/src/parser/benches/qwen3_xml.rs index f2e37551dda..59ea0de47dd 100644 --- a/rust/src/tool-parser/benches/qwen3_xml.rs +++ b/rust/src/parser/benches/qwen3_xml.rs @@ -2,8 +2,8 @@ use std::time::Duration; use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; use tool_parser::parsers::QwenParser as ExternalQwenParser; -use vllm_tool_parser::test_utils::{split_by_chars, test_tools}; -use vllm_tool_parser::{Qwen3XmlToolParser, Tool, ToolParser}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{Qwen3XmlToolParser, Tool, ToolParser}; mod utils; use utils::{feed_external_parser, feed_parser, openai_tools}; diff --git a/rust/src/tool-parser/benches/utils/mod.rs b/rust/src/parser/benches/utils/mod.rs similarity index 91% rename from rust/src/tool-parser/benches/utils/mod.rs rename to rust/src/parser/benches/utils/mod.rs index a0ad768f115..914766a79aa 100644 --- a/rust/src/tool-parser/benches/utils/mod.rs +++ b/rust/src/parser/benches/utils/mod.rs @@ -3,8 +3,8 @@ use futures::FutureExt as _; use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool}; use tool_parser::traits::ToolParser as ExternalToolParser; -use vllm_tool_parser::test_utils::collect_stream; -use vllm_tool_parser::{Tool, ToolParser}; +use vllm_parser::tool::test_utils::collect_stream; +use vllm_parser::tool::{Tool, ToolParser}; pub(super) fn openai_tools(tools: &[Tool]) -> Vec { tools @@ -23,7 +23,7 @@ pub(super) fn openai_tools(tools: &[Tool]) -> Vec { pub(super) fn feed_parser(parser: &mut dyn ToolParser, chunks: &[&str]) -> (String, usize) { let result = collect_stream(parser, chunks); - (result.normal_text, result.calls.len()) + (result.normal_text(), result.calls().len()) } pub(super) fn feed_external_parser( diff --git a/rust/src/tool-parser/python/Cargo.toml b/rust/src/parser/python/Cargo.toml similarity index 91% rename from rust/src/tool-parser/python/Cargo.toml rename to rust/src/parser/python/Cargo.toml index c029ad90135..aadae5638f9 100644 --- a/rust/src/tool-parser/python/Cargo.toml +++ b/rust/src/parser/python/Cargo.toml @@ -13,7 +13,7 @@ pyo3.workspace = true pythonize = { workspace = true, features = ["serde_json"] } serde_json.workspace = true thiserror-ext.workspace = true -vllm-tool-parser.workspace = true +vllm-parser.workspace = true [lints] workspace = true diff --git a/rust/src/tool-parser/python/src/lib.rs b/rust/src/parser/python/src/lib.rs similarity index 93% rename from rust/src/tool-parser/python/src/lib.rs rename to rust/src/parser/python/src/lib.rs index e5ae0fa7b69..4567348bcd9 100644 --- a/rust/src/tool-parser/python/src/lib.rs +++ b/rust/src/parser/python/src/lib.rs @@ -1,4 +1,4 @@ -//! Thin PyO3 bindings for `vllm_tool_parser`. +//! Thin PyO3 bindings for `vllm_parser::tool`. //! //! This crate exposes the Rust tool parser trait and data shapes to Python //! while keeping parser state, grammar, and schema-aware argument conversion in @@ -11,7 +11,7 @@ use pyo3::types::{PyAny, PyModule}; use pythonize::{depythonize, pythonize}; use serde_json::Value; use thiserror_ext::AsReport as _; -use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput}; +use vllm_parser::tool::{Tool, ToolCallDelta, ToolParser, ToolParserOutput}; macro_rules! tool_parser_factory { ($($parser:ident),+ $(,)?) => { @@ -22,7 +22,7 @@ macro_rules! tool_parser_factory { match name { $( stringify!($parser) => { - ::create(tools) + ::create(tools) } )+ _ => { @@ -146,30 +146,30 @@ impl PyToolParserOutput { #[new] #[pyo3(signature = (normal_text="", calls=None))] fn new(py: Python<'_>, normal_text: &str, calls: Option>>) -> Self { - let calls = - calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect(); - Self(ToolParserOutput { - normal_text: normal_text.to_owned(), - calls, - }) + let mut output = ToolParserOutput::default(); + output.push_text(normal_text); + for call in calls.unwrap_or_default() { + output.push_call(call.borrow(py).0.clone()); + } + Self(output) } #[getter] - fn normal_text(&self) -> &str { - &self.0.normal_text + fn normal_text(&self) -> String { + self.0.normal_text() } #[getter] fn calls(&self) -> Vec { - self.0.calls.iter().cloned().map(PyToolCallDelta).collect() + self.0.calls().into_iter().cloned().map(PyToolCallDelta).collect() } fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) { self.0.append(other.0.clone()); } - fn coalesce_calls(&self) -> Self { - Self(self.0.clone().coalesce_calls()) + fn coalesce(&self) -> Self { + Self(self.0.clone().coalesce()) } } @@ -300,7 +300,7 @@ mod tests { } #[test] - fn output_append_and_coalesce_calls() { + fn output_append_and_coalesce() { with_python(|py| { let first = Py::new( py, @@ -311,7 +311,7 @@ mod tests { let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?; output.append(other.borrow(py)); - let coalesced = output.coalesce_calls(); + let coalesced = output.coalesce(); assert_eq!(coalesced.normal_text(), "text"); let calls = coalesced.calls(); assert_eq!(calls.len(), 1); @@ -334,7 +334,7 @@ mod tests { parser.parse_into_output(&build_call(), &mut output)?; let finish = Py::new(py, parser.finish()?)?; output.append(finish.borrow(py)); - let output = output.coalesce_calls(); + let output = output.coalesce(); assert_eq!(output.normal_text(), ""); let calls = output.calls(); diff --git a/rust/src/parser/src/lib.rs b/rust/src/parser/src/lib.rs new file mode 100644 index 00000000000..5ba2cf60edd --- /dev/null +++ b/rust/src/parser/src/lib.rs @@ -0,0 +1,5 @@ +//! Streaming parsers for chat completions. + +pub mod reasoning; +pub mod tool; +pub mod unified; diff --git a/rust/src/reasoning-parser/src/cohere_cmd.rs b/rust/src/parser/src/reasoning/cohere_cmd.rs similarity index 100% rename from rust/src/reasoning-parser/src/cohere_cmd.rs rename to rust/src/parser/src/reasoning/cohere_cmd.rs diff --git a/rust/src/reasoning-parser/src/deepseek_r1.rs b/rust/src/parser/src/reasoning/deepseek_r1.rs similarity index 100% rename from rust/src/reasoning-parser/src/deepseek_r1.rs rename to rust/src/parser/src/reasoning/deepseek_r1.rs diff --git a/rust/src/reasoning-parser/src/delimited.rs b/rust/src/parser/src/reasoning/delimited.rs similarity index 100% rename from rust/src/reasoning-parser/src/delimited.rs rename to rust/src/parser/src/reasoning/delimited.rs diff --git a/rust/src/reasoning-parser/src/gemma4.rs b/rust/src/parser/src/reasoning/gemma4.rs similarity index 99% rename from rust/src/reasoning-parser/src/gemma4.rs rename to rust/src/parser/src/reasoning/gemma4.rs index 86824f2ad40..ac5a6a17165 100644 --- a/rust/src/reasoning-parser/src/gemma4.rs +++ b/rust/src/parser/src/reasoning/gemma4.rs @@ -119,7 +119,7 @@ mod tests { use vllm_tokenizer::Tokenizer; use super::Gemma4ReasoningParser; - use crate::ReasoningParser; + use crate::reasoning::ReasoningParser; struct FakeTokenizer; diff --git a/rust/src/reasoning-parser/src/kimi.rs b/rust/src/parser/src/reasoning/kimi.rs similarity index 100% rename from rust/src/reasoning-parser/src/kimi.rs rename to rust/src/parser/src/reasoning/kimi.rs diff --git a/rust/src/reasoning-parser/src/minimax_m3.rs b/rust/src/parser/src/reasoning/minimax_m3.rs similarity index 100% rename from rust/src/reasoning-parser/src/minimax_m3.rs rename to rust/src/parser/src/reasoning/minimax_m3.rs diff --git a/rust/src/reasoning-parser/src/lib.rs b/rust/src/parser/src/reasoning/mod.rs similarity index 100% rename from rust/src/reasoning-parser/src/lib.rs rename to rust/src/parser/src/reasoning/mod.rs diff --git a/rust/src/reasoning-parser/src/qwen3.rs b/rust/src/parser/src/reasoning/qwen3.rs similarity index 100% rename from rust/src/reasoning-parser/src/qwen3.rs rename to rust/src/parser/src/reasoning/qwen3.rs diff --git a/rust/src/reasoning-parser/src/seed_oss.rs b/rust/src/parser/src/reasoning/seed_oss.rs similarity index 98% rename from rust/src/reasoning-parser/src/seed_oss.rs rename to rust/src/parser/src/reasoning/seed_oss.rs index f514b43a89f..eb996f8477c 100644 --- a/rust/src/reasoning-parser/src/seed_oss.rs +++ b/rust/src/parser/src/reasoning/seed_oss.rs @@ -49,7 +49,7 @@ mod tests { use std::sync::Arc; use super::SeedOssReasoningParser; - use crate::{ReasoningParser, tests::FakeTokenizer}; + use crate::reasoning::{ReasoningParser, tests::FakeTokenizer}; #[test] fn without_prompt_markers_expects_start_token() { diff --git a/rust/src/reasoning-parser/src/step3p5.rs b/rust/src/parser/src/reasoning/step3p5.rs similarity index 99% rename from rust/src/reasoning-parser/src/step3p5.rs rename to rust/src/parser/src/reasoning/step3p5.rs index e369531c92c..d66506538a1 100644 --- a/rust/src/reasoning-parser/src/step3p5.rs +++ b/rust/src/parser/src/reasoning/step3p5.rs @@ -127,7 +127,7 @@ mod tests { use std::sync::Arc; use super::Step3p5ReasoningParser; - use crate::{ReasoningParser, tests::FakeTokenizer}; + use crate::reasoning::{ReasoningParser, tests::FakeTokenizer}; #[test] fn picks_up_prompt_start_boundary() { diff --git a/rust/src/reasoning-parser/src/tests.rs b/rust/src/parser/src/reasoning/tests.rs similarity index 100% rename from rust/src/reasoning-parser/src/tests.rs rename to rust/src/parser/src/reasoning/tests.rs diff --git a/rust/src/tool-parser/src/deepseek_dsml/deepseek_v32.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs similarity index 82% rename from rust/src/tool-parser/src/deepseek_dsml/deepseek_v32.rs rename to rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs index bc636c6035a..7d3432f9e5c 100644 --- a/rust/src/tool-parser/src/deepseek_dsml/deepseek_v32.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v32.rs @@ -1,5 +1,5 @@ use super::{DeepSeekDsmlToolParser, DsmlTokens}; -use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; /// Tool parser for DeepSeek V3.2 models. /// @@ -67,8 +67,8 @@ mod tests { use thiserror_ext::AsReport; use super::DeepSeekV32ToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { let params = params @@ -90,8 +90,8 @@ mod tests { let mut parser = DeepSeekV32ToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -104,11 +104,11 @@ mod tests { )) .unwrap(); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF", "date": "2024-01-16" @@ -125,8 +125,8 @@ mod tests { ); let output = parser.parse_complete(&output).unwrap(); - assert_eq!(output.normal_text, "Thinking... "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Thinking... "); + assert_eq!(output.calls().len(), 1); } #[test] @@ -146,9 +146,9 @@ mod tests { ) .unwrap(); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "whole": 5.0, "flag": true, @@ -176,9 +176,9 @@ mod tests { ) .unwrap(); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "whole": "5.0", "flag": "true", @@ -206,7 +206,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "Hangzhou </｜DSML｜parameter></｜DSML｜invoke></｜DSML｜function_calls>", "date": "2026-05-08", @@ -228,11 +228,11 @@ mod tests { ], ); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF" }) ); } @@ -252,8 +252,8 @@ mod tests { ], ); - assert_eq!(output.normal_text, "Thinking... "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Thinking... "); + assert_eq!(output.calls().len(), 1); } #[test] @@ -261,8 +261,8 @@ mod tests { let mut parser = DeepSeekV32ToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &["Hello, ", "world!"]); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -278,17 +278,17 @@ mod tests { )], ); - assert_eq!(output.calls.len(), 2); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[1].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[1].tool_index, 1); + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[1].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[1].tool_index, 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF" }) ); assert_eq!( - serde_json::from_str::(&output.calls[1].arguments).unwrap(), + serde_json::from_str::(&output.calls()[1].arguments).unwrap(), json!({ "location": "NYC" }) ); } @@ -300,9 +300,9 @@ mod tests { let mut parser = DeepSeekV32ToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF" }) ); } @@ -337,11 +337,11 @@ mod tests { ], ); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "Beijing" }) ); } @@ -373,9 +373,9 @@ mod tests { ], ); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); } #[test] @@ -393,8 +393,8 @@ mod tests { ], ); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); } #[test] @@ -421,10 +421,10 @@ mod tests { .parse_complete(&build_tool_call("get_weather", &[("location", "NYC")])) .unwrap(); - assert_eq!(first.calls.len(), 1); - assert_eq!(second.calls.len(), 1); + assert_eq!(first.calls().len(), 1); + assert_eq!(second.calls().len(), 1); assert_eq!( - serde_json::from_str::(&second.calls[0].arguments).unwrap(), + serde_json::from_str::(&second.calls()[0].arguments).unwrap(), json!({ "location": "NYC" }) ); } @@ -439,7 +439,7 @@ mod tests { let mut parser = DeepSeekV32ToolParser::new(&test_tools()); let complete = parser.parse_complete(&full_text).unwrap(); - assert_eq!(streamed.normal_text, complete.normal_text); - assert_eq!(streamed.calls, complete.calls); + assert_eq!(streamed.normal_text(), complete.normal_text()); + assert_eq!(streamed.calls(), complete.calls()); } } diff --git a/rust/src/tool-parser/src/deepseek_dsml/deepseek_v4.rs b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs similarity index 84% rename from rust/src/tool-parser/src/deepseek_dsml/deepseek_v4.rs rename to rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs index 9047b24ced7..a0493932766 100644 --- a/rust/src/tool-parser/src/deepseek_dsml/deepseek_v4.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/deepseek_v4.rs @@ -1,5 +1,5 @@ use super::{DeepSeekDsmlToolParser, DsmlTokens}; -use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; /// Tool parser for DeepSeek V4 models. /// @@ -69,8 +69,8 @@ mod tests { use serde_json::{Value, json}; use super::DeepSeekV4ToolParser; - use crate::test_utils::{collect_stream, test_tools}; - use crate::{StructuralTagModel, ToolParser, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, test_tools}; + use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { let params = params @@ -107,11 +107,11 @@ mod tests { )) .unwrap(); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF", "date": "2024-01-16" @@ -137,11 +137,11 @@ mod tests { ], ); - assert_eq!(output.normal_text, "Thinking... "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.normal_text(), "Thinking... "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "Beijing" }) ); } diff --git a/rust/src/tool-parser/src/deepseek_dsml/mod.rs b/rust/src/parser/src/tool/deepseek_dsml/mod.rs similarity index 97% rename from rust/src/tool-parser/src/deepseek_dsml/mod.rs rename to rust/src/parser/src/tool/deepseek_dsml/mod.rs index 1a2031dd3d7..b49fb1de8b5 100644 --- a/rust/src/tool-parser/src/deepseek_dsml/mod.rs +++ b/rust/src/parser/src/tool/deepseek_dsml/mod.rs @@ -7,7 +7,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, ToolParserOutput}; -use crate::Tool; +use crate::tool::Tool; mod deepseek_v32; mod deepseek_v4; @@ -92,7 +92,7 @@ impl DeepSeekDsmlToolParser { fn apply_event(&mut self, event: DsmlEvent, output: &mut ToolParserOutput) -> Result<()> { match event { DsmlEvent::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } DsmlEvent::ToolCallsStart => { self.mode = DsmlMode::ToolBlock { @@ -116,7 +116,7 @@ impl DeepSeekDsmlToolParser { let arguments = serde_json::to_string(&arguments) .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index: self.emitted_invoke_count, name: Some(name), arguments, @@ -156,7 +156,7 @@ impl DeepSeekDsmlToolParser { fn finish(&mut self) -> Result { let mut output = ToolParserOutput::default(); match self.mode { - DsmlMode::Text => output.normal_text.push_str(&self.buffer), + DsmlMode::Text => output.push_text(&self.buffer), DsmlMode::Done => {} DsmlMode::ToolBlock { .. } => { return Err(parsing_failed!("incomplete DeepSeek DSML tool call")); diff --git a/rust/src/tool-parser/src/deepseek_json/deepseek_v3.rs b/rust/src/parser/src/tool/deepseek_json/deepseek_v3.rs similarity index 78% rename from rust/src/tool-parser/src/deepseek_json/deepseek_v3.rs rename to rust/src/parser/src/tool/deepseek_json/deepseek_v3.rs index 6d6062432ab..ea1a660ccec 100644 --- a/rust/src/tool-parser/src/deepseek_json/deepseek_v3.rs +++ b/rust/src/parser/src/tool/deepseek_json/deepseek_v3.rs @@ -1,5 +1,5 @@ use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser}; -use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; /// Tool parser for DeepSeek V3 JSON-fenced tool calls. /// @@ -55,12 +55,12 @@ mod tests { use thiserror_ext::AsReport; use super::DeepSeekV3ToolParser; - use crate::deepseek_json::{ + use crate::tool::deepseek_json::{ TOOL_CALL_SEPARATOR, TOOL_CALL_START, TOOL_CALLS_END, TOOL_CALLS_START, V3_ARGUMENT_END, V3_JSON_START, }; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; fn v3_tool_call(function_name: &str, arguments: &str) -> String { format!( @@ -77,8 +77,8 @@ mod tests { let mut parser = DeepSeekV3ToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -92,11 +92,11 @@ mod tests { )) .unwrap(); - assert_eq!(output.normal_text, "Let me check.\n"); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.normal_text(), "Let me check.\n"); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -107,7 +107,7 @@ mod tests { .parse_complete(&tool_section(&[v3_tool_call("get_weather", arguments)])) .unwrap(); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -132,7 +132,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -143,7 +143,7 @@ mod tests { assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]); assert_eq!( - output.coalesce_calls().calls[0].arguments, + output.coalesce().calls()[0].arguments, r#"{"location":"Beijing"}"# ); } @@ -159,9 +159,9 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "hello "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#); + assert_eq!(output.normal_text(), "hello "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } #[test] @@ -172,8 +172,8 @@ mod tests { let output = parser.parse_complete(&input).unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -189,22 +189,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] diff --git a/rust/src/tool-parser/src/deepseek_json/deepseek_v31.rs b/rust/src/parser/src/tool/deepseek_json/deepseek_v31.rs similarity index 76% rename from rust/src/tool-parser/src/deepseek_json/deepseek_v31.rs rename to rust/src/parser/src/tool/deepseek_json/deepseek_v31.rs index 33b362439a8..cf2ea196282 100644 --- a/rust/src/tool-parser/src/deepseek_json/deepseek_v31.rs +++ b/rust/src/parser/src/tool/deepseek_json/deepseek_v31.rs @@ -1,5 +1,5 @@ use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser}; -use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; /// Tool parser for DeepSeek V3.1 raw JSON tool calls. /// @@ -51,11 +51,11 @@ mod tests { use thiserror_ext::AsReport; use super::DeepSeekV31ToolParser; - use crate::deepseek_json::{ + use crate::tool::deepseek_json::{ TOOL_CALL_END, TOOL_CALL_SEPARATOR, TOOL_CALL_START, TOOL_CALLS_END, TOOL_CALLS_START, }; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; fn v31_tool_call(function_name: &str, arguments: &str) -> String { format!("{TOOL_CALL_START}{function_name}{TOOL_CALL_SEPARATOR}{arguments}{TOOL_CALL_END}") @@ -70,8 +70,8 @@ mod tests { let mut parser = DeepSeekV31ToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -85,11 +85,11 @@ mod tests { )) .unwrap(); - assert_eq!(output.normal_text, "Let me check."); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.normal_text(), "Let me check."); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -100,7 +100,7 @@ mod tests { .parse_complete(&tool_section(&[v31_tool_call("get_weather", arguments)])) .unwrap(); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -123,7 +123,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -134,7 +134,7 @@ mod tests { assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]); assert_eq!( - output.coalesce_calls().calls[0].arguments, + output.coalesce().calls()[0].arguments, r#"{"location":"Beijing"}"# ); } @@ -150,9 +150,9 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "hello "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#); + assert_eq!(output.normal_text(), "hello "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } #[test] @@ -163,8 +163,8 @@ mod tests { let output = parser.parse_complete(&input).unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -180,22 +180,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] @@ -212,9 +215,9 @@ mod tests { let output = collect_stream(&mut parser, &[&input]); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } #[test] diff --git a/rust/src/tool-parser/src/deepseek_json/mod.rs b/rust/src/parser/src/tool/deepseek_json/mod.rs similarity index 97% rename from rust/src/tool-parser/src/deepseek_json/mod.rs rename to rust/src/parser/src/tool/deepseek_json/mod.rs index 0f0d04f0428..c6fce9fec67 100644 --- a/rust/src/tool-parser/src/deepseek_json/mod.rs +++ b/rust/src/parser/src/tool/deepseek_json/mod.rs @@ -96,7 +96,7 @@ impl DeepSeekJsonToolParser { ) -> Result<()> { match event { DeepSeekJsonEvent::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } DeepSeekJsonEvent::ToolCallsStart => self.mode = DeepSeekJsonMode::ToolBlock, DeepSeekJsonEvent::ToolCallStart => self.mode = DeepSeekJsonMode::Header, @@ -107,7 +107,7 @@ impl DeepSeekJsonToolParser { self.mode = DeepSeekJsonMode::Arguments { json_scan: JsonObjectScanState::default(), }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: Some(function_name), arguments: String::new(), @@ -120,7 +120,7 @@ impl DeepSeekJsonToolParser { self.format.parser_name() )); }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: None, arguments: self.buffer[..consumed_len].to_string(), @@ -155,7 +155,7 @@ impl DeepSeekJsonToolParser { fn finish(&mut self) -> Result { let mut output = ToolParserOutput::default(); match &self.mode { - DeepSeekJsonMode::Text => output.normal_text.push_str(&self.buffer), + DeepSeekJsonMode::Text => output.push_text(&self.buffer), DeepSeekJsonMode::ToolBlock | DeepSeekJsonMode::Done => {} DeepSeekJsonMode::Header | DeepSeekJsonMode::Arguments { .. } => { return Err(parsing_failed!( diff --git a/rust/src/tool-parser/src/error.rs b/rust/src/parser/src/tool/error.rs similarity index 87% rename from rust/src/tool-parser/src/error.rs rename to rust/src/parser/src/tool/error.rs index 0ac4a02c658..6a64c257d8c 100644 --- a/rust/src/tool-parser/src/error.rs +++ b/rust/src/parser/src/tool/error.rs @@ -6,7 +6,7 @@ pub type Result = std::result::Result; /// Errors produced while creating or running tool parsers. #[derive(Debug, Error, Macro)] -#[thiserror_ext(macro(path = "crate::error"))] +#[thiserror_ext(macro(path = "crate::tool::error"))] pub enum ToolParserError { #[error("tool parser parsing failed: {message}")] ParsingFailed { message: String }, diff --git a/rust/src/tool-parser/src/gemma4.rs b/rust/src/parser/src/tool/gemma4.rs similarity index 96% rename from rust/src/tool-parser/src/gemma4.rs rename to rust/src/parser/src/tool/gemma4.rs index 2fad84574c0..e5a95485ce8 100644 --- a/rust/src/tool-parser/src/gemma4.rs +++ b/rust/src/parser/src/tool/gemma4.rs @@ -8,7 +8,7 @@ use winnow::token::{literal, take_till, take_until}; use super::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::Tool; +use crate::tool::Tool; const TOOL_CALL_START: &str = "<|tool_call>"; const TOOL_CALL_END: &str = ""; @@ -70,7 +70,7 @@ impl Gemma4ToolParser { fn apply_event(&mut self, event: Gemma4Event, output: &mut ToolParserOutput) -> Result<()> { match event { Gemma4Event::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header, Gemma4Event::ToolCallHeader { name } => { @@ -89,7 +89,7 @@ impl Gemma4ToolParser { let arguments = serde_json::to_string(&args) .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index: self.emitted_tool_count, name: Some(name), arguments, @@ -153,7 +153,7 @@ impl ToolParser for Gemma4ToolParser { let mut output = ToolParserOutput::default(); match &self.mode { - Gemma4Mode::Text => output.normal_text.push_str(&self.buffer), + Gemma4Mode::Text => output.push_text(&self.buffer), Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => { return Err(parsing_failed!("incomplete Gemma4 tool call")); } @@ -428,7 +428,7 @@ mod tests { Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_array_content, parse_gemma4_args, }; - use crate::{Tool, ToolParserTestExt as _}; + use crate::tool::{Tool, ToolParserTestExt as _}; fn parse_gemma4_array(array: &str) -> super::Result> { let mut input = array; @@ -501,11 +501,11 @@ mod tests { output.append(parser.parse_chunk(chunk).unwrap()); } output.append(parser.finish().unwrap()); - output.coalesce_calls() + output.coalesce() } - fn first_call(output: &ToolParserOutput) -> &ToolCallDelta { - output.calls.first().expect("expected one tool call") + fn first_call(output: &ToolParserOutput) -> ToolCallDelta { + (*output.calls().first().expect("expected one tool call")).clone() } #[test] @@ -547,8 +547,8 @@ mod tests { .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}") .unwrap(); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); assert_eq!( serde_json::from_str::(&first_call(&output).arguments).unwrap(), @@ -577,7 +577,7 @@ mod tests { "", ]); - assert!(output.normal_text.is_empty()); + assert!(output.normal_text().is_empty()); assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); assert_eq!( serde_json::from_str::(&first_call(&output).arguments).unwrap(), @@ -597,7 +597,7 @@ mod tests { "div>", ]); - assert_eq!(output.normal_text, "Let me check the weather.
"); + assert_eq!(output.normal_text(), "Let me check the weather.
"); assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); assert_eq!( serde_json::from_str::(&first_call(&output).arguments).unwrap(), @@ -616,11 +616,11 @@ mod tests { "location:<|\"|>Paris<|\"|>}", ] { output.append(parser.parse_chunk(chunk).unwrap()); - assert!(output.calls.is_empty()); + assert!(output.calls().is_empty()); } output.append(parser.parse_chunk("").unwrap()); - let output = output.coalesce_calls(); + let output = output.coalesce(); assert_eq!(first_call(&output).name.as_deref(), Some("get_weather")); assert_eq!( @@ -777,8 +777,8 @@ mod tests { let mut output = parser.parse_chunk("<").unwrap(); output.append(parser.finish().unwrap()); - assert_eq!(output.normal_text, "<"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "<"); + assert!(output.calls().is_empty()); } #[test] diff --git a/rust/src/tool-parser/src/glm_xml/glm45_moe.rs b/rust/src/parser/src/tool/glm_xml/glm45_moe.rs similarity index 94% rename from rust/src/tool-parser/src/glm_xml/glm45_moe.rs rename to rust/src/parser/src/tool/glm_xml/glm45_moe.rs index 2a2d2e03813..a8d1ea0f19e 100644 --- a/rust/src/tool-parser/src/glm_xml/glm45_moe.rs +++ b/rust/src/parser/src/tool/glm_xml/glm45_moe.rs @@ -1,5 +1,5 @@ use super::{GlmXmlToolParser, Separator}; -use crate::{Result, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, Tool, ToolParser, ToolParserOutput}; /// Tool parser for GLM-4.5/4.6 MoE XML-style tool calls. /// diff --git a/rust/src/tool-parser/src/glm_xml/glm47_moe.rs b/rust/src/parser/src/tool/glm_xml/glm47_moe.rs similarity index 77% rename from rust/src/tool-parser/src/glm_xml/glm47_moe.rs rename to rust/src/parser/src/tool/glm_xml/glm47_moe.rs index 74afd6c250b..ac1a9d6ac6d 100644 --- a/rust/src/tool-parser/src/glm_xml/glm47_moe.rs +++ b/rust/src/parser/src/tool/glm_xml/glm47_moe.rs @@ -1,5 +1,5 @@ use super::{GlmXmlToolParser, Separator}; -use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; /// Tool parser for GLM-4.7 MoE XML-style tool calls. /// @@ -44,8 +44,8 @@ mod tests { use serde_json::{Value, json}; use super::Glm47MoeToolParser; - use crate::ToolParserTestExt as _; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::ToolParserTestExt as _; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; fn glm47_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { let params = params @@ -69,11 +69,11 @@ mod tests { let output = parser.parse_complete(&output).unwrap(); - assert_eq!(output.normal_text, "Let me search for that.\n"); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.normal_text(), "Let me search for that.\n"); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({"city": "Beijing", "date": "2024-12-25"}) ); } @@ -90,12 +90,12 @@ mod tests { let chunks = split_by_chars(&output, 7); let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, ""); - assert_eq!(output.calls.len(), 2); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[1].name.as_deref(), Some("add")); + assert_eq!(output.normal_text(), ""); + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[1].name.as_deref(), Some("add")); assert_eq!( - serde_json::from_str::(&output.calls[1].arguments).unwrap(), + serde_json::from_str::(&output.calls()[1].arguments).unwrap(), json!({"x": 1, "y": 2}) ); } @@ -117,7 +117,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "whole": 42, "flag": true, @@ -134,10 +134,10 @@ mod tests { let output = parser.parse_complete("add").unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("add")); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("add")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({}) ); } diff --git a/rust/src/tool-parser/src/glm_xml/mod.rs b/rust/src/parser/src/tool/glm_xml/mod.rs similarity index 88% rename from rust/src/tool-parser/src/glm_xml/mod.rs rename to rust/src/parser/src/tool/glm_xml/mod.rs index ceeb9a75173..cc175aeb641 100644 --- a/rust/src/tool-parser/src/glm_xml/mod.rs +++ b/rust/src/parser/src/tool/glm_xml/mod.rs @@ -7,7 +7,7 @@ use winnow::token::{literal, rest, take_until, take_while}; use super::parameters::ToolSchemas; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; use super::{Result, ToolCallDelta, ToolParserOutput}; -use crate::Tool; +use crate::tool::Tool; mod glm45_moe; mod glm47_moe; @@ -79,7 +79,7 @@ impl GlmXmlToolParser { fn apply_event(&mut self, event: GlmEvent, output: &mut ToolParserOutput) -> Result<()> { match event { GlmEvent::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } GlmEvent::ToolCallStart => { self.mode = GlmMode::ToolCall { @@ -92,7 +92,7 @@ impl GlmXmlToolParser { let arguments = serde_json::to_string(&arguments) .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index: self.emitted_tool_count, name: Some(name), arguments, @@ -127,7 +127,7 @@ impl GlmXmlToolParser { let mut output = ToolParserOutput::default(); if !self.buffer.is_empty() { match self.mode { - GlmMode::Text => output.normal_text.push_str(&self.buffer), + GlmMode::Text => output.push_text(&self.buffer), GlmMode::ToolCall { .. } => { return Err(parsing_failed!("incomplete GLM MoE tool call")); } @@ -264,8 +264,8 @@ mod tests { use thiserror_ext::AsReport; use super::Glm45MoeToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserTestExt as _}; fn glm45_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { let params = params @@ -283,8 +283,8 @@ mod tests { let mut parser = Glm45MoeToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -300,11 +300,11 @@ mod tests { let output = parser.parse_complete(&output).unwrap(); - assert_eq!(output.normal_text, "Let me search for that.\n"); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.normal_text(), "Let me search for that.\n"); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({"city": "Beijing", "date": "2024-12-25"}) ); } @@ -321,12 +321,12 @@ mod tests { let chunks = split_by_chars(&output, 11); let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, ""); - assert_eq!(output.calls.len(), 2); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[1].name.as_deref(), Some("add")); + assert_eq!(output.normal_text(), ""); + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[1].name.as_deref(), Some("add")); assert_eq!( - serde_json::from_str::(&output.calls[1].arguments).unwrap(), + serde_json::from_str::(&output.calls()[1].arguments).unwrap(), json!({"x": 1, "y": 2}) ); } @@ -345,7 +345,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Paris </arg_value></tool_call>", "date": "2026-05-08", @@ -359,8 +359,8 @@ mod tests { let output = collect_stream(&mut parser, &["hello ", "world"]); - assert_eq!(output.normal_text, "hello world"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "hello world"); + assert!(output.calls().is_empty()); } #[test] @@ -375,8 +375,8 @@ mod tests { ], ); - assert_eq!(output.normal_text, "Prefix "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Prefix "); + assert_eq!(output.calls().len(), 1); } #[test] @@ -391,9 +391,9 @@ mod tests { ], ); - assert_eq!(output.normal_text, "hello "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.normal_text(), "hello "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); } #[test] @@ -402,8 +402,8 @@ mod tests { let output = parser.parse_chunk("get_weather\ncity").unwrap(); - assert_eq!(output.normal_text, ""); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), ""); + assert!(output.calls().is_empty()); } #[test] @@ -437,7 +437,7 @@ mod tests { )], ); - assert_eq!(output.normal_text, ""); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), ""); + assert_eq!(output.calls().len(), 1); } } diff --git a/rust/src/tool-parser/src/hy_v3.rs b/rust/src/parser/src/tool/hy_v3.rs similarity index 87% rename from rust/src/tool-parser/src/hy_v3.rs rename to rust/src/parser/src/tool/hy_v3.rs index c0cf9446348..566df28d320 100644 --- a/rust/src/tool-parser/src/hy_v3.rs +++ b/rust/src/parser/src/tool/hy_v3.rs @@ -7,7 +7,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::{StructuralTagModel, Tool}; +use crate::tool::{StructuralTagModel, Tool}; const TOOL_CALLS_START: &str = ""; const TOOL_CALLS_END: &str = ""; @@ -79,7 +79,7 @@ impl HyV3ToolParser { fn apply_event(&mut self, event: HyV3Event, output: &mut ToolParserOutput) -> Result<()> { match event { HyV3Event::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } HyV3Event::ToolBlockStart => { self.mode = HyV3Mode::ToolBlock { @@ -91,7 +91,7 @@ impl HyV3ToolParser { let arguments = serde_json::to_string(&arguments) .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index: self.emitted_tool_count, name: Some(name), arguments, @@ -133,7 +133,7 @@ impl ToolParser for HyV3ToolParser { fn finish(&mut self) -> Result { let mut output = ToolParserOutput::default(); match self.mode { - HyV3Mode::Text => output.normal_text.push_str(&self.buffer), + HyV3Mode::Text => output.push_text(&self.buffer), HyV3Mode::ToolBlock { .. } => return Err(parsing_failed!("incomplete HY3 tool call")), HyV3Mode::Done => {} } @@ -249,8 +249,8 @@ mod tests { use thiserror_ext::AsReport; use super::{HyV3ToolParser, ToolParser}; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParserOutput, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { let params = params @@ -266,7 +266,7 @@ mod tests { } fn parsed_arguments(output: &ToolParserOutput, index: usize) -> Value { - serde_json::from_str(&output.calls[index].arguments).unwrap() + serde_json::from_str(&output.calls()[index].arguments).unwrap() } #[test] @@ -281,8 +281,8 @@ mod tests { let mut parser = HyV3ToolParser::new(&test_tools()); let output = parser.parse_complete("This is a plain response.").unwrap(); - assert_eq!(output.normal_text, "This is a plain response."); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "This is a plain response."); + assert!(output.calls().is_empty()); } #[test] @@ -294,9 +294,9 @@ mod tests { ) .unwrap(); - assert_eq!(output.normal_text, ""); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date")); + assert_eq!(output.normal_text(), ""); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date")); assert_eq!(parsed_arguments(&output, 0), json!({})); } @@ -309,7 +309,7 @@ mod tests { ) .unwrap(); - assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date")); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date")); assert_eq!(parsed_arguments(&output, 0), json!({})); } @@ -354,8 +354,8 @@ mod tests { )) .unwrap(); - assert_eq!(output.normal_text, "Checking."); - assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date")); + assert_eq!(output.normal_text(), "Checking."); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date")); } #[test] @@ -376,22 +376,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"city\":\"Beijing\",\"date\":\"2026-03-30\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "get_weather", - ), - arguments: "{\"city\":\"Hangzhou\",\"date\":\"2026-03-30\"}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"city\":\"Beijing\",\"date\":\"2026-03-30\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "get_weather", + ), + arguments: "{\"city\":\"Hangzhou\",\"date\":\"2026-03-30\"}", + }, + ), ], } "#]] @@ -434,8 +437,8 @@ mod tests { output.append(parser.parse_chunk("response.").unwrap()); output.append(parser.finish().unwrap()); - assert_eq!(output.normal_text, "This is a plain response."); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "This is a plain response."); + assert!(output.calls().is_empty()); } #[test] @@ -452,8 +455,8 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date")); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date")); assert_eq!(parsed_arguments(&output, 0), json!({})); } @@ -475,8 +478,8 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( parsed_arguments(&output, 0), json!({ "city": "Beijing", "date": "2026-03-30" }) @@ -498,8 +501,8 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "Checking."); - assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date")); + assert_eq!(output.normal_text(), "Checking."); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_current_date")); } #[test] @@ -519,7 +522,7 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 2); + assert_eq!(output.calls().len(), 2); assert_eq!(parsed_arguments(&output, 0)["city"], json!("Beijing")); assert_eq!(parsed_arguments(&output, 1)["city"], json!("Hangzhou")); } @@ -535,8 +538,8 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "hello "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "hello "); + assert_eq!(output.calls().len(), 1); assert_eq!(parsed_arguments(&output, 0), json!({ "city": "Beijing" })); } @@ -552,8 +555,8 @@ mod tests { ) .unwrap(); - assert_eq!(output.normal_text, ""); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), ""); + assert!(output.calls().is_empty()); } #[test] diff --git a/rust/src/tool-parser/src/json/granite4.rs b/rust/src/parser/src/tool/json/granite4.rs similarity index 79% rename from rust/src/tool-parser/src/json/granite4.rs rename to rust/src/parser/src/tool/json/granite4.rs index a70c0645400..fe1cf190225 100644 --- a/rust/src/tool-parser/src/json/granite4.rs +++ b/rust/src/parser/src/tool/json/granite4.rs @@ -8,10 +8,10 @@ use super::{ JsonToolCallConfig, JsonToolCallEvent, JsonToolCallWhitespace, JsonToolInput, tool_call_header_event, }; -use crate::utils::{ +use crate::tool::utils::{ JsonObjectScanState, json_str, parse_buffered_event, safe_text_len, take_json_object, }; -use crate::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; +use crate::tool::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; const TOOL_CALL_START: &str = ""; const TOOL_CALL_END: &str = ""; @@ -86,14 +86,14 @@ impl Granite4ToolParser { /// Apply one parsed Granite 4 event to parser state and output. fn apply_event(&mut self, event: Granite4Event, output: &mut ToolParserOutput) -> Result<()> { match event { - Granite4Event::Text { len } => output.normal_text.push_str(&self.buffer[..len]), + Granite4Event::Text { len } => output.push_text(&self.buffer[..len]), Granite4Event::ToolCallStart => self.mode = Granite4Mode::Header, Granite4Event::ToolCallHeader { function_name } => { let tool_index = self.emitted_tool_count; self.emitted_tool_count += 1; self.active_tool_index = Some(tool_index); self.mode = Granite4Mode::Args { json_scan: None }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: Some(function_name), arguments: String::new(), @@ -125,7 +125,7 @@ impl Granite4ToolParser { "Granite4 arguments without an active tool call" )); }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: None, arguments, @@ -165,7 +165,7 @@ impl ToolParser for Granite4ToolParser { fn finish(&mut self) -> Result { let mut output = ToolParserOutput::default(); match &self.mode { - Granite4Mode::Text => output.normal_text.push_str(&self.buffer), + Granite4Mode::Text => output.push_text(&self.buffer), Granite4Mode::Header | Granite4Mode::Args { .. } | Granite4Mode::Close => { return Err(parsing_failed!("incomplete Granite4 tool call")); } @@ -279,16 +279,16 @@ mod tests { use thiserror_ext::AsReport; use super::Granite4ToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; #[test] fn granite4_parse_complete_without_tool_call_keeps_text() { let mut parser = Granite4ToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -300,10 +300,10 @@ mod tests { ) .unwrap(); - assert_eq!(output.normal_text, ""); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, r#"{"city":"Boston"}"#); + assert_eq!(output.normal_text(), ""); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, r#"{"city":"Boston"}"#); } #[test] @@ -317,9 +317,9 @@ mod tests { ) .unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, r#"{"city":"Boston"}"#); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, r#"{"city":"Boston"}"#); } #[test] @@ -333,22 +333,28 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "before middle after", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "find_bbox", - ), - arguments: "{\"x\":1}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "get_weather", - ), - arguments: "{\"city\":\"Boston\"}", - }, + events: [ + Text( + "before middle after", + ), + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "find_bbox", + ), + arguments: "{\"x\":1}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "get_weather", + ), + arguments: "{\"city\":\"Boston\"}", + }, + ), ], } "#]] @@ -363,10 +369,10 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "hello bye"); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, r#"{"city":"Tokyo"}"#); + assert_eq!(output.normal_text(), "hello bye"); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, r#"{"city":"Tokyo"}"#); } #[test] @@ -385,7 +391,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -396,7 +402,7 @@ mod tests { assert_eq!(observed_arguments, [r#"{"city":"#, r#""Beijing""#, r#"}"#]); assert_eq!( - output.coalesce_calls().calls[0].arguments, + output.coalesce().calls()[0].arguments, r#"{"city":"Beijing"}"# ); } @@ -409,9 +415,9 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("f")); - assert_eq!(output.calls[0].arguments, r#"{"a":1}"#); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("f")); + assert_eq!(output.calls()[0].arguments, r#"{"a":1}"#); } #[test] @@ -435,29 +441,37 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "Here goes the bbox call: \n Now the stock price call: \n Now another bbox call: \n See? I'm a helpful assistant.", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "find_bbox", - ), - arguments: "{\"coordinates\": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], \"coordinate_type\": \"latlong\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "get_stock_price", - ), - arguments: "{\"symbol\": \"AAPL\", \"start_date\": \"2021-01-01\", \"end_date\": \"2021-12-31\"}", - }, - ToolCallDelta { - tool_index: 2, - name: Some( - "find_bbox", - ), - arguments: "{\"coordinates\": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], \"coordinate_type\": \"latlong\"}", - }, + events: [ + Text( + "Here goes the bbox call: \n Now the stock price call: \n Now another bbox call: \n See? I'm a helpful assistant.", + ), + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "find_bbox", + ), + arguments: "{\"coordinates\": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], \"coordinate_type\": \"latlong\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "get_stock_price", + ), + arguments: "{\"symbol\": \"AAPL\", \"start_date\": \"2021-01-01\", \"end_date\": \"2021-12-31\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 2, + name: Some( + "find_bbox", + ), + arguments: "{\"coordinates\": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], \"coordinate_type\": \"latlong\"}", + }, + ), ], } "#]].assert_debug_eq(&output); diff --git a/rust/src/tool-parser/src/json/hermes.rs b/rust/src/parser/src/tool/json/hermes.rs similarity index 75% rename from rust/src/tool-parser/src/json/hermes.rs rename to rust/src/parser/src/tool/json/hermes.rs index 04635185176..817eaee91f1 100644 --- a/rust/src/tool-parser/src/json/hermes.rs +++ b/rust/src/parser/src/tool/json/hermes.rs @@ -1,5 +1,5 @@ use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; -use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig { parser_name: "Hermes", @@ -68,8 +68,8 @@ mod tests { use thiserror_ext::AsReport; use super::HermesToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, arguments: &str) -> String { format!(r#"{{"name":"{function_name}","arguments":{arguments}}}"#) @@ -80,8 +80,8 @@ mod tests { let mut parser = HermesToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -95,11 +95,11 @@ mod tests { )) .unwrap(); - assert_eq!(output.normal_text, "Let me check.\n"); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.normal_text(), "Let me check.\n"); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -112,8 +112,8 @@ mod tests { ) .unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); } #[test] @@ -122,7 +122,7 @@ mod tests { let arguments = r#"{"location":"Tokyo",}"#; let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap(); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -142,7 +142,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -152,9 +152,9 @@ mod tests { output.append(parser.finish().unwrap()); assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]); - assert_eq!(output.normal_text, "preface suffix"); + assert_eq!(output.normal_text(), "preface suffix"); assert_eq!( - output.coalesce_calls().calls[0].arguments, + output.coalesce().calls()[0].arguments, r#"{"location":"Beijing"}"# ); } @@ -170,9 +170,9 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "hello "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#); + assert_eq!(output.normal_text(), "hello "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } #[test] @@ -189,22 +189,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] diff --git a/rust/src/tool-parser/src/json/internlm2.rs b/rust/src/parser/src/tool/json/internlm2.rs similarity index 84% rename from rust/src/tool-parser/src/json/internlm2.rs rename to rust/src/parser/src/tool/json/internlm2.rs index 8284a4d0e1d..aae3b9f6a02 100644 --- a/rust/src/tool-parser/src/json/internlm2.rs +++ b/rust/src/parser/src/tool/json/internlm2.rs @@ -1,5 +1,5 @@ use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; -use crate::{Result, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, Tool, ToolParser, ToolParserOutput}; const INTERNLM2_CONFIG: JsonToolCallConfig = JsonToolCallConfig { parser_name: "InternLM2", @@ -123,8 +123,8 @@ mod tests { use thiserror_ext::AsReport; use super::Internlm2ToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; const ACTION_START: &str = "<|action_start|><|plugin|>"; const ACTION_END: &str = "<|action_end|>"; @@ -140,8 +140,8 @@ mod tests { let mut parser = Internlm2ToolParser::new(&test_tools()); let result = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(result.normal_text, "Hello, world!"); - assert!(result.calls.is_empty()); + assert_eq!(result.normal_text(), "Hello, world!"); + assert!(result.calls().is_empty()); } #[test] @@ -155,11 +155,11 @@ mod tests { )) .unwrap(); - assert_eq!(result.normal_text, "Let me check.\n"); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].tool_index, 0); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.normal_text(), "Let me check.\n"); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].tool_index, 0); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls()[0].arguments, arguments); } #[test] @@ -170,9 +170,9 @@ mod tests { .parse_complete(&build_tool_call("get_weather", "arguments", arguments)) .unwrap(); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls()[0].arguments, arguments); } #[test] @@ -185,8 +185,8 @@ mod tests { )) .unwrap(); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); } #[test] @@ -197,7 +197,7 @@ mod tests { .parse_complete(&build_tool_call("get_weather", "parameters", arguments)) .unwrap(); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.calls()[0].arguments, arguments); } #[test] @@ -218,7 +218,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -231,9 +231,9 @@ mod tests { observed_arguments, [r#"{"location":"#, r#""Beijing""#, r#"}"#] ); - assert_eq!(result.normal_text, "preface suffix"); + assert_eq!(result.normal_text(), "preface suffix"); assert_eq!( - result.coalesce_calls().calls[0].arguments, + result.coalesce().calls()[0].arguments, r#"{"location":"Beijing"}"# ); } @@ -249,9 +249,9 @@ mod tests { let result = collect_stream(&mut parser, &chunks); - assert_eq!(result.normal_text, "hello "); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#); + assert_eq!(result.normal_text(), "hello "); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } #[test] @@ -268,22 +268,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] @@ -298,8 +301,8 @@ mod tests { let result = parser.parse_complete(&input).unwrap(); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].arguments, arguments); } #[test] @@ -313,7 +316,7 @@ mod tests { let error = parser.finish().unwrap_err(); assert_eq!( - pre_finish.calls[0].name.as_deref(), + pre_finish.calls()[0].name.as_deref(), Some("get_weather"), "name delta is still emitted from parse_chunk() before truncation", ); diff --git a/rust/src/tool-parser/src/json/llama.rs b/rust/src/parser/src/tool/json/llama.rs similarity index 86% rename from rust/src/tool-parser/src/json/llama.rs rename to rust/src/parser/src/tool/json/llama.rs index d9456487d1f..b736f27e306 100644 --- a/rust/src/tool-parser/src/json/llama.rs +++ b/rust/src/parser/src/tool/json/llama.rs @@ -8,8 +8,8 @@ use super::{ JsonToolCallConfig, JsonToolCallEvent, JsonToolCallWhitespace, JsonToolInput, argument_delta_event, tool_call_header_event, }; -use crate::utils::{JsonObjectScanState, parse_buffered_event}; -use crate::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; +use crate::tool::utils::{JsonObjectScanState, parse_buffered_event}; +use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; #[derive(Debug, Clone, PartialEq, Eq)] enum LlamaJsonMode { @@ -87,7 +87,7 @@ impl Llama3JsonToolParser { self.mode = LlamaJsonMode::Arguments { json_scan: JsonObjectScanState::default(), }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: Some(function_name), arguments: String::new(), @@ -99,7 +99,7 @@ impl Llama3JsonToolParser { "Llama JSON arguments without an active tool call" )); }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: None, arguments: self.buffer[..consumed_len].to_string(), @@ -145,7 +145,7 @@ impl ToolParser for Llama3JsonToolParser { } if matches!(self.mode, LlamaJsonMode::Passthrough) { - output.normal_text.push_str(&self.buffer); + output.push_text(&self.buffer); self.buffer.clear(); return Ok(()); } @@ -164,7 +164,7 @@ impl ToolParser for Llama3JsonToolParser { let mut output = ToolParserOutput::default(); match &self.mode { LlamaJsonMode::Start | LlamaJsonMode::Passthrough => { - output.normal_text.push_str(&self.buffer); + output.push_text(&self.buffer); } LlamaJsonMode::AfterCall if self.buffer.trim().is_empty() => {} LlamaJsonMode::Header | LlamaJsonMode::Arguments { .. } => { @@ -256,8 +256,8 @@ mod tests { use thiserror_ext::AsReport; use super::Llama3JsonToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, parameters: &str) -> String { format!(r#"{{"name":"{function_name}","parameters":{parameters}}}"#) @@ -268,8 +268,8 @@ mod tests { let mut parser = Llama3JsonToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -284,10 +284,10 @@ mod tests { output.append(parser.finish().unwrap()); assert_eq!( - output.normal_text, + output.normal_text(), r#"plain text first {"name":"get_weather","parameters":{"location":"Tokyo"}}"# ); - assert!(output.calls.is_empty()); + assert!(output.calls().is_empty()); } #[test] @@ -299,8 +299,8 @@ mod tests { ); let output = parser.parse_complete(&input).unwrap(); - assert_eq!(output.normal_text, input); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), input); + assert!(output.calls().is_empty()); } #[test] @@ -312,8 +312,8 @@ mod tests { ); let output = parser.parse_complete(&input).unwrap(); - assert_eq!(output.normal_text, input); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), input); + assert!(output.calls().is_empty()); } #[test] @@ -322,10 +322,10 @@ mod tests { let arguments = r#"{ "location": "Tokyo", "days": 3 }"#; let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -353,22 +353,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] @@ -390,7 +393,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -401,7 +404,7 @@ mod tests { assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]); assert_eq!( - output.coalesce_calls().calls[0].arguments, + output.coalesce().calls()[0].arguments, r#"{"location":"Beijing"}"# ); } @@ -418,14 +421,14 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, ""); - assert_eq!(output.calls.len(), 2); + assert_eq!(output.normal_text(), ""); + assert_eq!(output.calls().len(), 2); assert_eq!( - output.calls[0].arguments, + output.calls()[0].arguments, r#"{"location":"Dallas","state":"TX"}"# ); - assert_eq!(output.calls[1].name.as_deref(), Some("add")); - assert_eq!(output.calls[1].arguments, r#"{"x":4,"y":5}"#); + assert_eq!(output.calls()[1].name.as_deref(), Some("add")); + assert_eq!(output.calls()[1].arguments, r#"{"x":4,"y":5}"#); } #[test] @@ -437,7 +440,7 @@ mod tests { }"#; let output = parser.parse_complete(&build_tool_call("convert", arguments)).unwrap(); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -450,8 +453,8 @@ mod tests { )) .unwrap(); - assert_eq!(output.normal_text, ""); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), ""); + assert_eq!(output.calls().len(), 1); } #[test] diff --git a/rust/src/tool-parser/src/json/mistral.rs b/rust/src/parser/src/tool/json/mistral.rs similarity index 76% rename from rust/src/tool-parser/src/json/mistral.rs rename to rust/src/parser/src/tool/json/mistral.rs index 9ca40fcaf97..8a20b4db7b8 100644 --- a/rust/src/tool-parser/src/json/mistral.rs +++ b/rust/src/parser/src/tool/json/mistral.rs @@ -1,5 +1,5 @@ use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; -use crate::{Result, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, Tool, ToolParser, ToolParserOutput}; const MISTRAL_CONFIG: JsonToolCallConfig = JsonToolCallConfig { parser_name: "Mistral", @@ -61,8 +61,8 @@ mod tests { use thiserror_ext::AsReport; use super::MistralToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, arguments: &str) -> String { format!(r#"{{"name":"{function_name}","arguments":{arguments}}}"#) @@ -77,8 +77,8 @@ mod tests { let mut parser = MistralToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -92,11 +92,11 @@ mod tests { )) .unwrap(); - assert_eq!(output.normal_text, "Let me check.\n"); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.normal_text(), "Let me check.\n"); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -115,22 +115,28 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "I'll help.\n", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"city\": \"Tokyo\", \"units\": \"celsius\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\": 1, \"y\": 2}", - }, + events: [ + Text( + "I'll help.\n", + ), + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"city\": \"Tokyo\", \"units\": \"celsius\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\": 1, \"y\": 2}", + }, + ), ], } "#]] @@ -148,7 +154,7 @@ mod tests { )])) .unwrap(); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -168,7 +174,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -178,9 +184,9 @@ mod tests { output.append(parser.finish().unwrap()); assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]); - assert_eq!(output.normal_text, "preface suffix"); + assert_eq!(output.normal_text(), "preface suffix"); assert_eq!( - output.coalesce_calls().calls[0].arguments, + output.coalesce().calls()[0].arguments, r#"{"location":"Beijing"}"# ); } @@ -196,9 +202,9 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "hello "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#); + assert_eq!(output.normal_text(), "hello "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } #[test] @@ -209,8 +215,8 @@ mod tests { .parse_complete(&build_tool_calls(&[build_tool_call("echo", arguments)])) .unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] diff --git a/rust/src/tool-parser/src/json/mod.rs b/rust/src/parser/src/tool/json/mod.rs similarity index 86% rename from rust/src/tool-parser/src/json/mod.rs rename to rust/src/parser/src/tool/json/mod.rs index 748f7e49e4d..6a701de435e 100644 --- a/rust/src/tool-parser/src/json/mod.rs +++ b/rust/src/parser/src/tool/json/mod.rs @@ -106,7 +106,7 @@ impl JsonToolCallParser { fn finish(&mut self) -> Result { let mut output = ToolParserOutput::default(); match &self.mode { - JsonToolCallMode::Text => output.normal_text.push_str(&self.buffer), + JsonToolCallMode::Text => output.push_text(&self.buffer), JsonToolCallMode::Header | JsonToolCallMode::Arguments { .. } => { return Err(parsing_failed!( "incomplete {} tool call", @@ -126,7 +126,7 @@ impl JsonToolCallParser { ) -> Result<()> { match event { JsonToolCallEvent::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } JsonToolCallEvent::ToolCallStart => self.mode = JsonToolCallMode::Header, JsonToolCallEvent::ToolCallHeader { function_name } => { @@ -136,7 +136,7 @@ impl JsonToolCallParser { self.mode = JsonToolCallMode::Arguments { json_scan: JsonObjectScanState::default(), }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: Some(function_name), arguments: String::new(), @@ -149,7 +149,7 @@ impl JsonToolCallParser { self.config.parser_name )); }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: None, arguments: self.buffer[..consumed_len].to_string(), @@ -374,7 +374,7 @@ mod tests { use expect_test::expect; use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; - use crate::ToolParserOutput; + use crate::tool::ToolParserOutput; const DELIMITED_CONFIG: JsonToolCallConfig = JsonToolCallConfig { parser_name: "Delimited JSON", @@ -400,7 +400,7 @@ mod tests { parser.parse_into(chunk, &mut output).unwrap(); } output.append(parser.finish().unwrap()); - output.coalesce_calls() + output.coalesce() } #[test] @@ -415,22 +415,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] @@ -451,22 +454,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] @@ -486,15 +492,19 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: " trailing text", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, + events: [ + Text( + " trailing text", + ), + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), ], } "#]] diff --git a/rust/src/tool-parser/src/json/phi4mini.rs b/rust/src/parser/src/tool/json/phi4mini.rs similarity index 77% rename from rust/src/tool-parser/src/json/phi4mini.rs rename to rust/src/parser/src/tool/json/phi4mini.rs index 463354b13c9..3f259c2d7fe 100644 --- a/rust/src/tool-parser/src/json/phi4mini.rs +++ b/rust/src/parser/src/tool/json/phi4mini.rs @@ -1,5 +1,5 @@ use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; -use crate::{Result, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, Tool, ToolParser, ToolParserOutput}; const PHI4MINI_CONFIG: JsonToolCallConfig = JsonToolCallConfig { parser_name: "Phi4Mini", @@ -69,8 +69,8 @@ mod tests { use thiserror_ext::AsReport; use super::Phi4MiniJsonToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserTestExt as _}; /// Build one phi-4-mini tool-call object: `{"name":..,"":}`. fn build_call(function_name: &str, args_key: &str, arguments: &str) -> String { @@ -87,8 +87,8 @@ mod tests { let mut parser = Phi4MiniJsonToolParser::new(&test_tools()); let result = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(result.normal_text, "Hello, world!"); - assert!(result.calls.is_empty()); + assert_eq!(result.normal_text(), "Hello, world!"); + assert!(result.calls().is_empty()); } #[test] @@ -99,10 +99,10 @@ mod tests { .parse_complete(&wrap(&[build_call("get_weather", "arguments", arguments)])) .unwrap(); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].tool_index, 0); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].tool_index, 0); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls()[0].arguments, arguments); } #[test] @@ -113,9 +113,9 @@ mod tests { .parse_complete(&wrap(&[build_call("get_weather", "parameters", arguments)])) .unwrap(); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls()[0].arguments, arguments); } #[test] @@ -130,22 +130,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] @@ -162,8 +165,8 @@ mod tests { .parse_complete(&wrap(&[build_call("convert", "arguments", arguments)])) .unwrap(); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].arguments, arguments); } /// Preface text before a tool call is preserved as normal_text, consistent @@ -182,8 +185,8 @@ mod tests { let result = parser.parse_complete(&input).unwrap(); - assert_eq!(result.normal_text, "Let me check.\n"); - assert_eq!(result.calls.len(), 1); + assert_eq!(result.normal_text(), "Let me check.\n"); + assert_eq!(result.calls().len(), 1); } #[test] @@ -194,7 +197,7 @@ mod tests { .parse_complete(&wrap(&[build_call("get_weather", "arguments", arguments)])) .unwrap(); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.calls()[0].arguments, arguments); } /// The bundled `tool_chat_template_phi4_mini.jinja` emits objects with @@ -208,9 +211,9 @@ mod tests { let result = parser.parse_complete(input).unwrap(); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(result.calls[0].arguments, r#"{"location": "Tokyo"}"#); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls()[0].arguments, r#"{"location": "Tokyo"}"#); } /// Argument deltas are streamed through the shared JSON core. @@ -230,10 +233,10 @@ mod tests { let result = collect_stream(&mut parser, &chunks); - assert_eq!(result.normal_text, "preface suffix"); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(result.calls[0].arguments, r#"{"location":"Beijing"}"#); + assert_eq!(result.normal_text(), "preface suffix"); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls()[0].arguments, r#"{"location":"Beijing"}"#); } #[test] @@ -251,9 +254,9 @@ mod tests { let result = collect_stream(&mut parser, &chunks); - assert_eq!(result.normal_text, "hello "); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].arguments, r#"{"location":"Tokyo"}"#); + assert_eq!(result.normal_text(), "hello "); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } #[test] @@ -287,9 +290,9 @@ mod tests { .parse_complete(&wrap(&[build_call("convert", "arguments", arguments)])) .unwrap(); - assert_eq!(result.calls.len(), 1); - assert_eq!(result.calls[0].name.as_deref(), Some("convert")); - assert_eq!(result.calls[0].arguments, arguments); + assert_eq!(result.calls().len(), 1); + assert_eq!(result.calls()[0].name.as_deref(), Some("convert")); + assert_eq!(result.calls()[0].arguments, arguments); } /// The chat template emits parallel calls as `},\n {` (comma + newline + @@ -307,9 +310,9 @@ mod tests { let result = parser.parse_complete(input).unwrap(); - assert_eq!(result.calls.len(), 2); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(result.calls[1].name.as_deref(), Some("add")); + assert_eq!(result.calls().len(), 2); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls()[1].name.as_deref(), Some("add")); } /// The shared core requires an object after the start marker. diff --git a/rust/src/tool-parser/src/json/qwen.rs b/rust/src/parser/src/tool/json/qwen.rs similarity index 78% rename from rust/src/tool-parser/src/json/qwen.rs rename to rust/src/parser/src/tool/json/qwen.rs index dd943dfffc7..7fa53c9007e 100644 --- a/rust/src/tool-parser/src/json/qwen.rs +++ b/rust/src/parser/src/tool/json/qwen.rs @@ -1,5 +1,5 @@ use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace}; -use crate::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; +use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput}; const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig { parser_name: "Qwen XML", @@ -70,8 +70,8 @@ mod tests { use thiserror_ext::AsReport; use super::Qwen3XmlToolParser; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParser, ToolParserOutput, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, arguments: &str) -> String { format!( @@ -84,8 +84,8 @@ mod tests { let mut parser = Qwen3XmlToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -99,11 +99,11 @@ mod tests { )) .unwrap(); - assert_eq!(output.normal_text, "Let me check.\n"); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.normal_text(), "Let me check.\n"); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -112,7 +112,7 @@ mod tests { let arguments = r#"{"location":"Tokyo",}"#; let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap(); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -132,7 +132,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -143,7 +143,7 @@ mod tests { assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]); assert_eq!( - output.coalesce_calls().calls[0].arguments, + output.coalesce().calls()[0].arguments, r#"{"location":"Beijing"}"# ); } @@ -159,9 +159,9 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "hello "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#); + assert_eq!(output.normal_text(), "hello "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Tokyo"}"#); } #[test] @@ -170,8 +170,8 @@ mod tests { let arguments = r#"{"text":"literal inside"}"#; let output = parser.parse_complete(&build_tool_call("echo", arguments)).unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -185,7 +185,7 @@ mod tests { ) .unwrap(); - assert_eq!(output.calls[0].name.as_deref(), Some("say_\"hi")); + assert_eq!(output.calls()[0].name.as_deref(), Some("say_\"hi")); } #[test] @@ -196,8 +196,8 @@ mod tests { let output = parser.parse_complete(input).unwrap(); - assert_eq!(output.normal_text, input); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), input); + assert!(output.calls().is_empty()); } #[test] @@ -227,22 +227,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] diff --git a/rust/src/tool-parser/src/kimi_k2.rs b/rust/src/parser/src/tool/kimi_k2.rs similarity index 88% rename from rust/src/tool-parser/src/kimi_k2.rs rename to rust/src/parser/src/tool/kimi_k2.rs index f83611ac79d..b692c1ec598 100644 --- a/rust/src/tool-parser/src/kimi_k2.rs +++ b/rust/src/parser/src/tool/kimi_k2.rs @@ -8,7 +8,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::{StructuralTagModel, Tool}; +use crate::tool::{StructuralTagModel, Tool}; const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>"; const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>"; @@ -81,7 +81,7 @@ impl KimiK2ToolParser { fn apply_event(&mut self, event: KimiK2Event, output: &mut ToolParserOutput) -> Result<()> { match event { KimiK2Event::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } KimiK2Event::ToolCallsStart => self.mode = KimiK2Mode::ToolBlock, KimiK2Event::ToolCallStart => self.mode = KimiK2Mode::Header, @@ -96,7 +96,7 @@ impl KimiK2ToolParser { json_scan: JsonObjectScanState::default(), }; self.call_ids.insert(tool_index, tool_call_id); - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: Some(function_name), arguments: String::new(), @@ -108,7 +108,7 @@ impl KimiK2ToolParser { "Kimi K2 arguments without an active tool call" )); }; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index, name: None, arguments: self.buffer[..consumed_len].to_string(), @@ -171,7 +171,7 @@ impl ToolParser for KimiK2ToolParser { fn finish(&mut self) -> Result { let mut output = ToolParserOutput::default(); match &self.mode { - KimiK2Mode::Text => output.normal_text.push_str(&self.buffer), + KimiK2Mode::Text => output.push_text(&self.buffer), KimiK2Mode::ToolBlock | KimiK2Mode::Done => {} KimiK2Mode::Header | KimiK2Mode::Arguments { .. } => { return Err(parsing_failed!("incomplete Kimi K2 tool call")); @@ -339,8 +339,8 @@ mod tests { KimiK2ToolParser, TOOL_CALL_ARGUMENT_START, TOOL_CALL_END, TOOL_CALL_START, TOOL_CALLS_END, TOOL_CALLS_START, ToolParser, tool_header, }; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParserOutput, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, index: usize, arguments: &str) -> String { format!( @@ -357,8 +357,8 @@ mod tests { let mut parser = KimiK2ToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -372,11 +372,11 @@ mod tests { )) .unwrap(); - assert_eq!(output.normal_text, "Checking. "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.normal_text(), "Checking. "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -391,7 +391,7 @@ mod tests { )])) .unwrap(); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -414,7 +414,7 @@ mod tests { for chunk in chunks { let next = parser.parse_chunk(chunk).unwrap(); observed_arguments.extend( - next.calls + next.calls() .iter() .filter(|call| call.name.is_none()) .map(|call| call.arguments.clone()), @@ -424,8 +424,8 @@ mod tests { output.append(parser.finish().unwrap()); assert_eq!(observed_arguments, ["{\"location\":", "\"Paris\"", "}"]); - let output = output.coalesce_calls(); - assert_eq!(output.calls[0].arguments, r#"{"location":"Paris"}"#); + let output = output.coalesce(); + assert_eq!(output.calls()[0].arguments, r#"{"location":"Paris"}"#); } #[test] @@ -445,9 +445,9 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.normal_text, "hello "); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, r#"{"location":"NYC"}"#); + assert_eq!(output.normal_text(), "hello "); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, r#"{"location":"NYC"}"#); } #[test] @@ -458,8 +458,8 @@ mod tests { let output = parser.parse_complete(&input).unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].arguments, arguments); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].arguments, arguments); } #[test] @@ -478,9 +478,9 @@ mod tests { let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - output.calls[0].arguments, + output.calls()[0].arguments, r#"{"text":"literal <|tool_call_end|> inside"}"# ); } @@ -498,22 +498,25 @@ mod tests { expect![[r#" ToolParserOutput { - normal_text: "", - calls: [ - ToolCallDelta { - tool_index: 0, - name: Some( - "get_weather", - ), - arguments: "{\"location\":\"Shanghai\"}", - }, - ToolCallDelta { - tool_index: 1, - name: Some( - "add", - ), - arguments: "{\"x\":1,\"y\":2}", - }, + events: [ + ToolCall( + ToolCallDelta { + tool_index: 0, + name: Some( + "get_weather", + ), + arguments: "{\"location\":\"Shanghai\"}", + }, + ), + ToolCall( + ToolCallDelta { + tool_index: 1, + name: Some( + "add", + ), + arguments: "{\"x\":1,\"y\":2}", + }, + ), ], } "#]] @@ -546,12 +549,12 @@ mod tests { "{TOOL_CALLS_START}{TOOL_CALL_START}api.tools.search:42{TOOL_CALL_ARGUMENT_START}{{}}{TOOL_CALL_END}{TOOL_CALLS_END}" ); - let output = parser.parse_chunk(&input).unwrap().coalesce_calls(); + let output = parser.parse_chunk(&input).unwrap().coalesce(); - assert_eq!(output.calls[0].tool_index, 42); + assert_eq!(output.calls()[0].tool_index, 42); assert_eq!(parser.tool_call_id(42), Some("api.tools.search:42")); - assert_eq!(output.calls[0].name.as_deref(), Some("search")); - assert_eq!(output.calls[0].arguments, "{}"); + assert_eq!(output.calls()[0].name.as_deref(), Some("search")); + assert_eq!(output.calls()[0].arguments, "{}"); } #[test] diff --git a/rust/src/tool-parser/src/minimax_m2.rs b/rust/src/parser/src/tool/minimax_m2.rs similarity index 86% rename from rust/src/tool-parser/src/minimax_m2.rs rename to rust/src/parser/src/tool/minimax_m2.rs index 16e2b85525f..5c5411775a9 100644 --- a/rust/src/tool-parser/src/minimax_m2.rs +++ b/rust/src/parser/src/tool/minimax_m2.rs @@ -7,7 +7,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::{StructuralTagModel, Tool}; +use crate::tool::{StructuralTagModel, Tool}; const TOOL_CALL_START: &str = ""; const TOOL_CALL_END: &str = ""; @@ -72,7 +72,7 @@ impl MinimaxM2ToolParser { fn apply_event(&mut self, event: MinimaxM2Event, output: &mut ToolParserOutput) -> Result<()> { match event { MinimaxM2Event::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } MinimaxM2Event::ToolBlockStart => { self.mode = MinimaxM2Mode::ToolBlock { @@ -84,7 +84,7 @@ impl MinimaxM2ToolParser { let arguments = serde_json::to_string(&arguments) .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index: self.emitted_tool_count, name: Some(name), arguments, @@ -133,7 +133,7 @@ impl ToolParser for MinimaxM2ToolParser { let mut output = ToolParserOutput::default(); match self.mode { MinimaxM2Mode::Text => { - output.normal_text.push_str(&self.buffer); + output.push_text(&self.buffer); } MinimaxM2Mode::ToolBlock { .. } => { return Err(parsing_failed!("incomplete MiniMax M2 tool call")); @@ -272,8 +272,8 @@ mod tests { use thiserror_ext::AsReport; use super::{MinimaxM2ToolParser, TOOL_CALL_END, TOOL_CALL_START, ToolParser}; - use crate::ToolParserTestExt as _; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::ToolParserTestExt as _; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; fn build_tool_block(invokes: &[(&str, Vec<(&str, &str)>)]) -> String { let invokes = invokes @@ -295,8 +295,8 @@ mod tests { let mut parser = MinimaxM2ToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -309,11 +309,11 @@ mod tests { )])) .unwrap(); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle", "days": 5 }) ); } @@ -327,8 +327,8 @@ mod tests { ); let output = parser.parse_complete(&output).unwrap(); - assert_eq!(output.normal_text, "Let me check. "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Let me check. "); + assert_eq!(output.calls().len(), 1); } #[test] @@ -341,15 +341,15 @@ mod tests { ])) .unwrap(); - assert_eq!(output.calls.len(), 2); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[1].tool_index, 1); + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[1].tool_index, 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle" }) ); assert_eq!( - serde_json::from_str::(&output.calls[1].arguments).unwrap(), + serde_json::from_str::(&output.calls()[1].arguments).unwrap(), json!({ "city": "NYC" }) ); } @@ -371,7 +371,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "whole": 5.0, "flag": true, @@ -395,7 +395,7 @@ mod tests { vec![("city", "Tom & Jerry <3")], )])) .unwrap(); - let args: Value = serde_json::from_str(&output.calls[0].arguments).unwrap(); + let args: Value = serde_json::from_str(&output.calls()[0].arguments).unwrap(); assert_eq!(args["city"], json!("Tom & Jerry <3")); } @@ -416,7 +416,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle </parameter></invoke></minimax:tool_call>", "days": 5, @@ -440,7 +440,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "shape": "\nrectangle\n", "dimensions": { "width": 10, "height": 20 }, @@ -462,11 +462,11 @@ mod tests { ], ); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle" }) ); } @@ -484,8 +484,8 @@ mod tests { ], ); - assert_eq!(output.normal_text, "Let me check. "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Let me check. "); + assert_eq!(output.calls().len(), 1); } #[test] @@ -493,8 +493,8 @@ mod tests { let mut parser = MinimaxM2ToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &["Hello, ", "world!"]); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -504,8 +504,8 @@ mod tests { let mut parser = MinimaxM2ToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 1); - assert!(output.normal_text.is_empty()); + assert_eq!(output.calls().len(), 1); + assert!(output.normal_text().is_empty()); } #[test] @@ -518,9 +518,9 @@ mod tests { let mut parser = MinimaxM2ToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 2); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[1].tool_index, 1); + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[1].tool_index, 1); } #[test] @@ -540,12 +540,12 @@ mod tests { let mut parser = MinimaxM2ToolParser::new(&test_tools()); let result = collect_stream(&mut parser, &chunks); - assert_eq!(result.normal_text, "I will call the tools.\n"); - assert_eq!(result.calls.len(), 2); - assert_eq!(result.calls[0].tool_index, 0); - assert_eq!(result.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(result.calls[1].tool_index, 1); - assert_eq!(result.calls[1].name.as_deref(), Some("get_weather")); + assert_eq!(result.normal_text(), "I will call the tools.\n"); + assert_eq!(result.calls().len(), 2); + assert_eq!(result.calls()[0].tool_index, 0); + assert_eq!(result.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(result.calls()[1].tool_index, 1); + assert_eq!(result.calls()[1].name.as_deref(), Some("get_weather")); } #[test] @@ -558,8 +558,8 @@ mod tests { let mut parser = MinimaxM2ToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); } #[test] @@ -568,8 +568,8 @@ mod tests { let output = parser.parse_chunk(r#""#).unwrap(); - assert!(output.normal_text.is_empty()); - assert!(output.calls.is_empty()); + assert!(output.normal_text().is_empty()); + assert!(output.calls().is_empty()); } #[test] diff --git a/rust/src/tool-parser/src/minimax_m3.rs b/rust/src/parser/src/tool/minimax_m3.rs similarity index 88% rename from rust/src/tool-parser/src/minimax_m3.rs rename to rust/src/parser/src/tool/minimax_m3.rs index ad40a7f18b7..a1ab375b731 100644 --- a/rust/src/tool-parser/src/minimax_m3.rs +++ b/rust/src/parser/src/tool/minimax_m3.rs @@ -8,7 +8,7 @@ use winnow::token::{literal, rest, take_until}; use super::parameters::{ParamElement, ParamInput, ToolSchemas}; use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker}; use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::Tool; +use crate::tool::Tool; const NAMESPACE: &str = "]<]minimax[>["; const TOOL_CALL_START: &str = "]<]minimax[>["; @@ -107,7 +107,7 @@ impl MinimaxM3ToolParser { fn apply_event(&mut self, event: MinimaxM3Event, output: &mut ToolParserOutput) -> Result<()> { match event { MinimaxM3Event::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } MinimaxM3Event::ToolBlockStart => { self.mode = MinimaxM3Mode::ToolBlock { @@ -119,7 +119,7 @@ impl MinimaxM3ToolParser { let arguments = serde_json::to_string(&arguments) .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index: self.emitted_tool_count, name: Some(name), arguments, @@ -158,7 +158,7 @@ impl ToolParser for MinimaxM3ToolParser { let mut output = ToolParserOutput::default(); match self.mode { MinimaxM3Mode::Text => { - output.normal_text.push_str(&self.buffer); + output.push_text(&self.buffer); } MinimaxM3Mode::ToolBlock { .. } => { if !self.buffer.trim_start().is_empty() { @@ -388,8 +388,8 @@ mod tests { ELEMENT_END_START, ELEMENT_START, INVOKE_END, INVOKE_START, MinimaxM3ToolParser, TOOL_CALL_END, TOOL_CALL_START, ToolParser, }; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{Tool, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{Tool, ToolParserEvent, ToolParserTestExt as _}; fn element(name: &str, body: &str) -> String { format!("{ELEMENT_START}{name}>{body}{ELEMENT_END_START}{name}>") @@ -510,8 +510,8 @@ mod tests { let mut parser = MinimaxM3ToolParser::new(&m3_test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -524,11 +524,11 @@ mod tests { )])) .unwrap(); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle", "days": 5 }) ); } @@ -542,8 +542,8 @@ mod tests { ); let output = parser.parse_complete(&output).unwrap(); - assert_eq!(output.normal_text, "Let me check. "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Let me check. "); + assert_eq!(output.calls().len(), 1); } #[test] @@ -556,15 +556,15 @@ mod tests { ])) .unwrap(); - assert_eq!(output.calls.len(), 2); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[1].tool_index, 1); + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[1].tool_index, 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle" }) ); assert_eq!( - serde_json::from_str::(&output.calls[1].arguments).unwrap(), + serde_json::from_str::(&output.calls()[1].arguments).unwrap(), json!({ "city": "NYC" }) ); } @@ -585,7 +585,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle" }) ); } @@ -608,7 +608,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "whole": 5.0, "flag": true, @@ -627,7 +627,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "user_id": 42, "urgent": true, @@ -677,7 +677,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "shape": "\nrectangle\n", "dimensions": { "width": 10, "height": 20 }, @@ -698,11 +698,11 @@ mod tests { ], ); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle" }) ); } @@ -720,8 +720,36 @@ mod tests { ], ); - assert_eq!(output.normal_text, "Let me check. "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Let me check. "); + assert_eq!(output.calls().len(), 1); + } + + #[test] + fn minimax_m3_streaming_preserves_ordered_events() { + let mut parser = MinimaxM3ToolParser::new(&m3_test_tools()); + let output = collect_stream( + &mut parser, + &[ + "Let me check. ", + TOOL_CALL_START, + &invoke("get_weather", &element("city", "Seattle")), + TOOL_CALL_END, + ], + ); + + assert_eq!(output.events.len(), 2); + assert_eq!( + output.events[0], + ToolParserEvent::Text("Let me check. ".to_string()) + ); + let ToolParserEvent::ToolCall(call) = &output.events[1] else { + panic!("expected tool-call event"); + }; + assert_eq!(call.name.as_deref(), Some("get_weather")); + assert_eq!( + serde_json::from_str::(&call.arguments).unwrap(), + json!({ "city": "Seattle" }) + ); } #[test] @@ -729,8 +757,8 @@ mod tests { let mut parser = MinimaxM3ToolParser::new(&m3_test_tools()); let output = collect_stream(&mut parser, &["Hello, ", "world!"]); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -740,8 +768,8 @@ mod tests { let mut parser = MinimaxM3ToolParser::new(&m3_test_tools()); let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 1); - assert!(output.normal_text.is_empty()); + assert_eq!(output.calls().len(), 1); + assert!(output.normal_text().is_empty()); } #[test] @@ -754,9 +782,9 @@ mod tests { let mut parser = MinimaxM3ToolParser::new(&m3_test_tools()); let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 2); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[1].tool_index, 1); + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[1].tool_index, 1); } #[test] @@ -768,8 +796,8 @@ mod tests { )) .unwrap(); - assert!(output.normal_text.is_empty()); - assert!(output.calls.is_empty()); + assert!(output.normal_text().is_empty()); + assert!(output.calls().is_empty()); } #[test] @@ -782,8 +810,8 @@ mod tests { let mut parser = MinimaxM3ToolParser::new(&m3_test_tools()); let output = collect_stream(&mut parser, &chunks); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); } #[test] @@ -804,8 +832,8 @@ mod tests { parser.parse_chunk(TOOL_CALL_START).unwrap(); let output = parser.finish().unwrap(); - assert!(output.normal_text.is_empty()); - assert!(output.calls.is_empty()); + assert!(output.normal_text().is_empty()); + assert!(output.calls().is_empty()); } #[test] @@ -819,9 +847,9 @@ mod tests { )) .unwrap(); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Seattle" }) ); } @@ -863,7 +891,7 @@ mod tests { let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "payload": { "child": "value", @@ -887,7 +915,7 @@ mod tests { let output = parser.parse_complete(&build_tool_block(&[("convert", body)])).unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "payload": { "$text": "child text", diff --git a/rust/src/tool-parser/src/lib.rs b/rust/src/parser/src/tool/mod.rs similarity index 66% rename from rust/src/tool-parser/src/lib.rs rename to rust/src/parser/src/tool/mod.rs index 6f785ee1d18..a27e202e660 100644 --- a/rust/src/tool-parser/src/lib.rs +++ b/rust/src/parser/src/tool/mod.rs @@ -1,9 +1,9 @@ //! Streaming tool parsers for chat completions. #[macro_use] -mod error; +pub(crate) mod error; mod deepseek_dsml; -mod deepseek_json; +pub(crate) mod deepseek_json; mod gemma4; mod glm_xml; mod hy_v3; @@ -15,7 +15,7 @@ mod parameters; mod qwen_coder; #[cfg(any(test, feature = "test-util"))] pub mod test_utils; -mod utils; +pub(crate) mod utils; use std::collections::{BTreeMap, btree_map}; @@ -57,55 +57,115 @@ pub struct ToolCallDelta { pub arguments: String, } +/// One ordered event emitted while parsing assistant text. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ToolParserEvent { + /// Plain assistant text that is not part of any tool call. + Text(String), + /// A tool-call update extracted from assistant text. + ToolCall(ToolCallDelta), +} + /// Result of advancing tool parsing with one assistant-text input. #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct ToolParserOutput { - /// Plain assistant text that is not part of any tool call. - pub normal_text: String, - /// Tool-call updates extracted from this input. - pub calls: Vec, + /// Ordered parser events committed by this input. + pub events: Vec, } impl ToolParserOutput { - /// Append another parser output onto this one. - /// - /// Note that this does not attempt to merge multiple deltas for the same - /// tool call into one complete item. Call `coalesce_calls()` after if - /// that behavior is desired. - pub fn append(&mut self, mut other: Self) { - self.normal_text.push_str(&other.normal_text); - self.calls.append(&mut other.calls); + /// Append one visible text event if `text` is non-empty. + pub fn push_text(&mut self, text: impl AsRef + Into) { + if text.as_ref().is_empty() { + return; + } + if let Some(ToolParserEvent::Text(last_text)) = self.events.last_mut() { + last_text.push_str(text.as_ref()); + return; + } + self.events.push(ToolParserEvent::Text(text.into())); } - /// Merge multiple deltas for the same tool call into one complete item. + /// Append one tool-call update event. + pub fn push_call(&mut self, call: ToolCallDelta) { + self.events.push(ToolParserEvent::ToolCall(call)); + } + + /// Return all plain assistant text committed by this output. + /// + /// Texts before and after tool calls will be concatenated into a single string. To preserve + /// the original order of the text and tool-call events, directly access `events` instead. + pub fn normal_text(&self) -> String { + self.events + .iter() + .filter_map(|event| match event { + ToolParserEvent::Text(text) => Some(text.as_str()), + ToolParserEvent::ToolCall(_) => None, + }) + .collect() + } + + /// Return all tool-call updates committed by this output. + pub fn calls(&self) -> Vec<&ToolCallDelta> { + self.events + .iter() + .filter_map(|event| match event { + ToolParserEvent::Text(_) => None, + ToolParserEvent::ToolCall(call) => Some(call), + }) + .collect() + } + + /// Append another parser output onto this one. + /// + /// Note that this keeps events exactly as they arrive. Call `coalesce()` + /// after if final text and tool-call fragments should be flattened. + pub fn append(&mut self, other: Self) { + for event in other.events { + match event { + ToolParserEvent::Text(text) => self.push_text(text), + ToolParserEvent::ToolCall(call) => self.push_call(call), + } + } + } + + /// Flatten text and merge deltas for the same tool call. + /// + /// All text events are concatenated into one leading text event. Tool-call + /// events follow that text event in first-seen tool index order, with + /// argument fragments for the same tool call concatenated together. /// /// This is primarily used by the default `parse_complete()` implementation, /// which delegates through the incremental parser lifecycle and then /// needs to collapse streaming-style argument fragments into one final /// tool call. - pub fn coalesce_calls(mut self) -> Self { + pub fn coalesce(self) -> Self { let mut merged = BTreeMap::::new(); let mut order = Vec::new(); + let normal_text = self.normal_text(); - for call in self.calls { + for call in self.calls() { match merged.entry(call.tool_index) { btree_map::Entry::Vacant(entry) => { order.push(call.tool_index); - entry.insert(call); + entry.insert(call.clone()); } btree_map::Entry::Occupied(mut entry) => { let existing = entry.get_mut(); if existing.name.is_none() { - existing.name = call.name; + existing.name = call.name.clone(); } existing.arguments.push_str(&call.arguments); } } } - self.calls = - order.into_iter().filter_map(|tool_index| merged.remove(&tool_index)).collect(); - self + let mut output = Self::default(); + output.push_text(normal_text); + for call in order.into_iter().filter_map(|tool_index| merged.remove(&tool_index)) { + output.push_call(call); + } + output } } @@ -183,7 +243,7 @@ impl T { pub fn parse_complete(&mut self, text: &str) -> Result { let mut output = self.parse_chunk(text)?; output.append(self.finish()?); - Ok(output.coalesce_calls()) + Ok(output.coalesce()) } } diff --git a/rust/src/tool-parser/src/parameters.rs b/rust/src/parser/src/tool/parameters.rs similarity index 99% rename from rust/src/tool-parser/src/parameters.rs rename to rust/src/parser/src/tool/parameters.rs index f857c147cb6..f5661456e3e 100644 --- a/rust/src/tool-parser/src/parameters.rs +++ b/rust/src/parser/src/tool/parameters.rs @@ -2,7 +2,7 @@ use std::collections::BTreeMap; use serde_json::{Map, Number, Value}; -use crate::Tool; +use crate::tool::Tool; /// Normalized parameter schemas for all tools in one request. #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -416,7 +416,7 @@ mod tests { use serde_json::{Value, json}; use super::{ParamElement, ParamInput, ToolSchema, ToolSchemas}; - use crate::Tool; + use crate::tool::Tool; fn test_tool(name: &str, parameters: serde_json::Value) -> Tool { Tool { diff --git a/rust/src/tool-parser/src/qwen_coder.rs b/rust/src/parser/src/tool/qwen_coder.rs similarity index 83% rename from rust/src/tool-parser/src/qwen_coder.rs rename to rust/src/parser/src/tool/qwen_coder.rs index 5e78d7ae520..d67a5c42f1c 100644 --- a/rust/src/tool-parser/src/qwen_coder.rs +++ b/rust/src/parser/src/tool/qwen_coder.rs @@ -7,7 +7,7 @@ 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; +use crate::tool::Tool; const TOOL_CALL_START: &str = ""; const TOOL_CALL_END: &str = ""; @@ -74,7 +74,7 @@ impl Qwen3CoderToolParser { fn apply_event(&mut self, event: QwenCoderEvent, output: &mut ToolParserOutput) -> Result<()> { match event { QwenCoderEvent::Text { len: consumed_len } => { - output.normal_text.push_str(&self.buffer[..consumed_len]); + output.push_text(&self.buffer[..consumed_len]); } QwenCoderEvent::ToolCallStart => { self.mode = QwenCoderMode::ToolCall { @@ -87,7 +87,7 @@ impl Qwen3CoderToolParser { let arguments = serde_json::to_string(&arguments) .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?; - output.calls.push(ToolCallDelta { + output.push_call(ToolCallDelta { tool_index: self.emitted_tool_count, name: Some(name), arguments, @@ -138,7 +138,7 @@ impl ToolParser for Qwen3CoderToolParser { { return Err(parsing_failed!("incomplete Qwen Coder tool call")); } - output.normal_text.push_str(&self.buffer); + output.push_text(&self.buffer); } let _ = self.reset(); Ok(output) @@ -241,8 +241,8 @@ mod tests { use thiserror_ext::AsReport; use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser}; - use crate::test_utils::{collect_stream, split_by_chars, test_tools}; - use crate::{ToolParserOutput, ToolParserTestExt as _}; + use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools}; + use crate::tool::{ToolParserOutput, ToolParserTestExt as _}; fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String { let params = params @@ -268,8 +268,8 @@ mod tests { let mut parser = Qwen3CoderToolParser::new(&test_tools()); let output = parser.parse_complete("Hello, world!").unwrap(); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -282,11 +282,11 @@ mod tests { )) .unwrap(); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF", "date": "2026-04-29" @@ -303,8 +303,8 @@ mod tests { ); let output = parser.parse_complete(&output).unwrap(); - assert_eq!(output.normal_text, "Thinking... "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Thinking... "); + assert_eq!(output.calls().len(), 1); } #[test] @@ -323,9 +323,9 @@ mod tests { )) .unwrap(); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "whole": 5.0, "flag": true, @@ -341,10 +341,10 @@ mod tests { let mut parser = Qwen3CoderToolParser::new(&test_tools()); let output = parser.parse_complete(&build_tool_call("get_weather", &[])).unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({}) ); } @@ -371,10 +371,10 @@ mod tests { ) .unwrap(); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("calculate_area")); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("calculate_area")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "shape": "rectangle", "dimensions": { "width": 10, "height": 20 }, @@ -396,9 +396,9 @@ mod tests { )) .unwrap(); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "payload": { "nested": { @@ -426,9 +426,9 @@ mod tests { )) .unwrap(); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "html_content": r#"
Hello
"#, "xml_snippet": r#""#, @@ -452,9 +452,9 @@ mod tests { )) .unwrap(); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "杭州 </parameter></function></tool_call>", "date": "2026-05-08", @@ -472,9 +472,9 @@ mod tests { )) .unwrap(); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "data": { "key": "value", "count": 42 }, }) @@ -495,11 +495,11 @@ mod tests { ], ); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF" }) ); } @@ -519,8 +519,8 @@ mod tests { ], ); - assert_eq!(output.normal_text, "Thinking... "); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.normal_text(), "Thinking... "); + assert_eq!(output.calls().len(), 1); } #[test] @@ -528,8 +528,8 @@ mod tests { let mut parser = Qwen3CoderToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &["Hello, ", "world!"]); - assert_eq!(output.normal_text, "Hello, world!"); - assert!(output.calls.is_empty()); + assert_eq!(output.normal_text(), "Hello, world!"); + assert!(output.calls().is_empty()); } #[test] @@ -542,17 +542,17 @@ mod tests { let mut parser = Qwen3CoderToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &[&text]); - assert_eq!(output.calls.len(), 2); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[1].name.as_deref(), Some("get_weather")); - assert_eq!(output.calls[0].tool_index, 0); - assert_eq!(output.calls[1].tool_index, 1); + assert_eq!(output.calls().len(), 2); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[1].name.as_deref(), Some("get_weather")); + assert_eq!(output.calls()[0].tool_index, 0); + assert_eq!(output.calls()[1].tool_index, 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF" }) ); assert_eq!( - serde_json::from_str::(&output.calls[1].arguments).unwrap(), + serde_json::from_str::(&output.calls()[1].arguments).unwrap(), json!({ "location": "NYC" }) ); } @@ -569,16 +569,16 @@ mod tests { let output = collect_stream(&mut parser, &chunks); assert_eq!( - output.normal_text, + output.normal_text(), "I'll check two cities.Between calls.Done." ); - assert_eq!(output.calls.len(), 2); + assert_eq!(output.calls().len(), 2); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "city": "Dallas", "state": "TX" }) ); assert_eq!( - serde_json::from_str::(&output.calls[1].arguments).unwrap(), + serde_json::from_str::(&output.calls()[1].arguments).unwrap(), json!({ "city": "Orlando", "state": "FL" }) ); } @@ -590,9 +590,9 @@ mod tests { let mut parser = Qwen3CoderToolParser::new(&test_tools()); let output = collect_stream(&mut parser, &chunks); - assert_eq!(output.calls.len(), 1); + assert_eq!(output.calls().len(), 1); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF" }) ); } @@ -610,19 +610,19 @@ mod tests { ) .unwrap(); - assert!(output.normal_text.is_empty()); - assert!(output.calls.is_empty()); + assert!(output.normal_text().is_empty()); + assert!(output.calls().is_empty()); let mut output = output; output.append(parser.parse_chunk("_call>").unwrap()); output.append(parser.finish().unwrap()); - let output = output.coalesce_calls(); + let output = output.coalesce(); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "SF" }) ); } @@ -641,20 +641,20 @@ mod tests { for chunk in chunks { let chunk_output = parser.parse_chunk(chunk).unwrap(); - assert!(chunk_output.normal_text.is_empty()); - assert!(chunk_output.calls.is_empty()); + assert!(chunk_output.normal_text().is_empty()); + assert!(chunk_output.calls().is_empty()); output.append(chunk_output); } output.append(parser.parse_chunk(end_suffix).unwrap()); output.append(parser.finish().unwrap()); - let output = output.coalesce_calls(); + let output = output.coalesce(); - assert!(output.normal_text.is_empty()); - assert_eq!(output.calls.len(), 1); - assert_eq!(output.calls[0].name.as_deref(), Some("get_weather")); + assert!(output.normal_text().is_empty()); + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("get_weather")); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": long_location }) ); } @@ -666,8 +666,8 @@ mod tests { .parse_chunk("\n\nSF") .unwrap(); - assert!(output.normal_text.is_empty()); - assert!(output.calls.is_empty()); + assert!(output.normal_text().is_empty()); + assert!(output.calls().is_empty()); } #[test] @@ -710,7 +710,7 @@ mod tests { .unwrap(); assert_eq!( - serde_json::from_str::(&output.calls[0].arguments).unwrap(), + serde_json::from_str::(&output.calls()[0].arguments).unwrap(), json!({ "location": "Hangzhou" }) ); } diff --git a/rust/src/tool-parser/src/test_utils.rs b/rust/src/parser/src/tool/test_utils.rs similarity index 96% rename from rust/src/tool-parser/src/test_utils.rs rename to rust/src/parser/src/tool/test_utils.rs index 70178756e4c..c160977479c 100644 --- a/rust/src/tool-parser/src/test_utils.rs +++ b/rust/src/parser/src/tool/test_utils.rs @@ -1,7 +1,7 @@ use serde_json::json; use super::{ToolParser, ToolParserOutput}; -use crate::{Tool, ToolParserTestExt as _}; +use crate::tool::Tool; /// Build a reusable set of function tools for parser unit tests. pub fn test_tools() -> Vec { @@ -87,10 +87,10 @@ pub fn test_tools() -> Vec { pub fn collect_stream(parser: &mut T, chunks: &[&str]) -> ToolParserOutput { let mut output = ToolParserOutput::default(); for chunk in chunks { - output.append(parser.parse_chunk(chunk).unwrap()); + parser.parse_into(chunk, &mut output).unwrap(); } output.append(parser.finish().unwrap()); - output.coalesce_calls() + output.coalesce() } /// Split text into chunks containing at most `chunk_chars` Unicode scalar diff --git a/rust/src/parser/src/tool/tests.rs b/rust/src/parser/src/tool/tests.rs new file mode 100644 index 00000000000..5a79e764203 --- /dev/null +++ b/rust/src/parser/src/tool/tests.rs @@ -0,0 +1,169 @@ +use super::{Result, Tool, ToolCallDelta, ToolParser, ToolParserEvent, ToolParserOutput}; +use crate::tool::ToolParserTestExt as _; + +struct DefaultParser; + +impl ToolParser for DefaultParser { + fn create(_tools: &[Tool]) -> Result> + where + Self: Sized + 'static, + { + Ok(Box::new(Self)) + } + + fn parse_into(&mut self, _chunk: &str, _output: &mut ToolParserOutput) -> Result<()> { + Ok(()) + } + + fn finish(&mut self) -> Result { + Ok(ToolParserOutput::default()) + } + + fn reset(&mut self) -> String { + String::new() + } +} + +#[test] +fn tool_parser_does_not_preserve_special_tokens_by_default() { + let parser = DefaultParser; + + assert!(!parser.preserve_special_tokens()); +} + +#[test] +fn tool_parser_output_coalesces_adjacent_text_events() { + let mut output = ToolParserOutput::default(); + output.push_text("hello"); + output.push_text(" "); + output.push_text("world"); + output.push_call(ToolCallDelta { + tool_index: 0, + name: Some("lookup".to_string()), + arguments: "{}".to_string(), + }); + output.push_text("!"); + + assert_eq!( + output.events, + vec![ + ToolParserEvent::Text("hello world".to_string()), + ToolParserEvent::ToolCall(ToolCallDelta { + tool_index: 0, + name: Some("lookup".to_string()), + arguments: "{}".to_string(), + }), + ToolParserEvent::Text("!".to_string()), + ] + ); +} + +#[test] +fn tool_parser_output_append_coalesces_adjacent_text_events() { + let mut output = ToolParserOutput::default(); + output.push_text("hello"); + + let mut other = ToolParserOutput::default(); + other.push_text(" "); + other.push_text("world"); + output.append(other); + + let mut after_call = ToolParserOutput::default(); + after_call.push_call(ToolCallDelta { + tool_index: 0, + name: Some("lookup".to_string()), + arguments: "{}".to_string(), + }); + after_call.push_text("!"); + output.append(after_call); + + assert_eq!( + output.events, + vec![ + ToolParserEvent::Text("hello world".to_string()), + ToolParserEvent::ToolCall(ToolCallDelta { + tool_index: 0, + name: Some("lookup".to_string()), + arguments: "{}".to_string(), + }), + ToolParserEvent::Text("!".to_string()), + ] + ); +} + +#[test] +fn default_parse_complete_delegates_through_parse_chunk_and_finish() { + struct StreamingParser; + + impl ToolParser for StreamingParser { + fn create(_tools: &[Tool]) -> Result> + where + Self: Sized + 'static, + { + Ok(Box::new(Self)) + } + + fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> { + output.push_text("prefix "); + for call in [ + ToolCallDelta { + tool_index: 0, + name: Some("weather".to_string()), + arguments: "{\"location\":".to_string(), + }, + ToolCallDelta { + tool_index: 0, + name: None, + arguments: "\"Paris\"".to_string(), + }, + ToolCallDelta { + tool_index: 1, + name: Some("time".to_string()), + arguments: "{\"timezone\":".to_string(), + }, + ] { + output.push_call(call); + } + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = ToolParserOutput::default(); + output.push_text("suffix"); + output.push_call(ToolCallDelta { + tool_index: 0, + name: None, + arguments: "}".to_string(), + }); + output.push_call(ToolCallDelta { + tool_index: 1, + name: None, + arguments: "\"UTC\"}".to_string(), + }); + Ok(output) + } + + fn reset(&mut self) -> String { + String::new() + } + } + + let mut parser = StreamingParser; + let output = parser.parse_complete("ignored").unwrap(); + assert_eq!(output.normal_text(), "prefix suffix"); + assert_eq!( + output.calls().into_iter().cloned().collect::>(), + vec![ + ToolCallDelta { + tool_index: 0, + name: Some("weather".to_string()), + arguments: "{\"location\":\"Paris\"}".to_string(), + }, + ToolCallDelta { + tool_index: 1, + name: Some("time".to_string()), + arguments: "{\"timezone\":\"UTC\"}".to_string(), + }, + ] + ); +} diff --git a/rust/src/tool-parser/src/utils.rs b/rust/src/parser/src/tool/utils.rs similarity index 100% rename from rust/src/tool-parser/src/utils.rs rename to rust/src/parser/src/tool/utils.rs diff --git a/rust/src/parser/src/unified/combined.rs b/rust/src/parser/src/unified/combined.rs new file mode 100644 index 00000000000..3f1c669013d --- /dev/null +++ b/rust/src/parser/src/unified/combined.rs @@ -0,0 +1,346 @@ +//! Adapter that combines reasoning and tool parsers. + +use vllm_tokenizer::DynTokenizer; + +use crate::reasoning::ReasoningParser; +use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput}; + +use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput}; + +/// Unified parser that composes existing reasoning and tool parsers. +pub struct CombinedParser { + reasoning: Option>, + tool: Option>, +} + +impl CombinedParser { + /// Create a combined parser from optional reasoning and tool parsers. + pub fn new( + reasoning: Option>, + tool: Option>, + ) -> Self { + Self { reasoning, tool } + } + + /// Create a text-only combined parser. + pub fn plain_text_only() -> Self { + Self { + reasoning: None, + tool: None, + } + } + + fn parse_tool(&mut self, content: &str, output: &mut UnifiedParserOutput) -> Result<()> { + let Some(tool) = self.tool.as_mut() else { + output.push_text(content); + return Ok(()); + }; + + // Preserve any tool output that was already produced before the error. + let mut tool_output = ToolParserOutput::default(); + let result = tool.parse_into(content, &mut tool_output); + output.append_tool_output(tool_output); + result?; + + Ok(()) + } + + fn flush_tool(&mut self) -> Result { + let Some(tool) = self.tool.as_mut() else { + return Ok(UnifiedParserOutput::default()); + }; + + let output = tool.finish()?; + let mut unified = UnifiedParserOutput::default(); + unified.append_tool_output(output); + Ok(unified) + } +} + +impl UnifiedParser for CombinedParser { + fn create(_tools: &[Tool], _tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static, + { + Err(UnifiedParserError::CombinedParserConstructor) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + if let Some(reasoning) = self.reasoning.as_mut() { + reasoning.initialize(prompt_token_ids)?; + } + Ok(()) + } + + fn preserve_special_tokens(&self) -> bool { + self.reasoning.as_ref().is_some_and(|parser| parser.preserve_special_tokens()) + || self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens()) + } + + fn structural_tag_model(&self) -> Option { + self.tool.as_ref().and_then(|parser| parser.structural_tag_model()) + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.tool.as_ref().and_then(|parser| parser.tool_call_id(tool_index)) + } + + fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()> { + let Some(reasoning) = self.reasoning.as_mut() else { + return self.parse_tool(delta, output); + }; + + let reasoning_delta = reasoning.push(delta)?; + if let Some(reasoning) = reasoning_delta.reasoning { + output.push_reasoning(reasoning); + } + if let Some(content) = reasoning_delta.content { + self.parse_tool(&content, output)?; + } + Ok(()) + } + + fn finish(&mut self) -> Result { + let mut output = UnifiedParserOutput::default(); + if let Some(reasoning) = self.reasoning.as_mut() { + let reasoning_delta = reasoning.finish()?; + if let Some(reasoning) = reasoning_delta.reasoning { + output.push_reasoning(reasoning); + } + if let Some(content) = reasoning_delta.content { + self.parse_tool(&content, &mut output)?; + } + } + output.append(self.flush_tool()?); + Ok(output) + } + + fn reset(&mut self) -> String { + self.tool.as_mut().map_or_else(String::new, |parser| parser.reset()) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use vllm_tokenizer::Tokenizer; + + use super::CombinedParser; + use crate::reasoning::{Qwen3ReasoningParser, ReasoningDelta, ReasoningParser}; + use crate::tool::{Qwen3XmlToolParser, Tool, ToolParser}; + use crate::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; + + struct FakeTokenizer; + + impl Tokenizer for FakeTokenizer { + fn encode( + &self, + text: &str, + _add_special_tokens: bool, + ) -> vllm_tokenizer::Result> { + Ok(text.chars().map(u32::from).collect()) + } + + fn decode( + &self, + token_ids: &[u32], + _skip_special_tokens: bool, + ) -> vllm_tokenizer::Result { + Ok(token_ids + .iter() + .map(|token_id| char::from_u32(*token_id).unwrap_or('\u{FFFD}')) + .collect()) + } + + fn token_to_id(&self, token: &str) -> Option { + match token { + "" => Some(1), + "" => Some(2), + _ => None, + } + } + } + + fn test_tools() -> Vec { + vec![Tool { + name: "get_weather".to_string(), + description: None, + parameters: serde_json::json!({ + "type": "object", + "properties": { + "location": { "type": "string" } + }, + }), + strict: None, + }] + } + + fn collect(parser: &mut dyn UnifiedParser, chunks: &[&str]) -> UnifiedParserOutput { + let mut output = UnifiedParserOutput::default(); + for chunk in chunks { + parser.parse_into(chunk, &mut output).unwrap(); + } + output.append(parser.finish().unwrap()); + output + } + + struct PreserveReasoningParser; + + impl ReasoningParser for PreserveReasoningParser { + fn create( + _tokenizer: vllm_tokenizer::DynTokenizer, + ) -> crate::reasoning::Result> + where + Self: Sized + 'static, + { + Ok(Box::new(Self)) + } + + fn preserve_special_tokens(&self) -> bool { + true + } + + fn push(&mut self, delta: &str) -> crate::reasoning::Result { + Ok(ReasoningDelta { + reasoning: None, + content: Some(delta.to_string()), + }) + } + } + + struct PreserveToolParser; + + impl ToolParser for PreserveToolParser { + fn create(_tools: &[Tool]) -> crate::tool::Result> + where + Self: Sized + 'static, + { + Ok(Box::new(Self)) + } + + fn preserve_special_tokens(&self) -> bool { + true + } + + fn parse_into( + &mut self, + chunk: &str, + output: &mut crate::tool::ToolParserOutput, + ) -> crate::tool::Result<()> { + output.push_text(chunk); + Ok(()) + } + + fn finish(&mut self) -> crate::tool::Result { + Ok(crate::tool::ToolParserOutput::default()) + } + + fn reset(&mut self) -> String { + String::new() + } + } + + struct PartialThenErrorToolParser; + + impl ToolParser for PartialThenErrorToolParser { + fn create(_tools: &[Tool]) -> crate::tool::Result> + where + Self: Sized + 'static, + { + Ok(Box::new(Self)) + } + + fn parse_into( + &mut self, + _chunk: &str, + output: &mut crate::tool::ToolParserOutput, + ) -> crate::tool::Result<()> { + output.push_text("committed"); + Err(crate::tool::ToolParserError::ParsingFailed { + message: "synthetic failure".to_string(), + }) + } + + fn finish(&mut self) -> crate::tool::Result { + Ok(crate::tool::ToolParserOutput::default()) + } + + fn reset(&mut self) -> String { + String::new() + } + } + + #[test] + fn combined_parser_emits_reasoning_and_text() { + let tokenizer = Arc::new(FakeTokenizer); + let reasoning = Qwen3ReasoningParser::create(tokenizer).unwrap(); + let mut parser = CombinedParser::new(Some(reasoning), None); + + let output = collect(&mut parser, &["workanswer"]); + + assert_eq!( + output.events, + vec![ + UnifiedParserEvent::Reasoning("work".to_string()), + UnifiedParserEvent::Text("answer".to_string()), + ] + ); + } + + #[test] + 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) + )); + + let output = collect( + &mut parser, + &[r#" +{"name":"get_weather","arguments":{"location":"Paris"}} +"#], + ); + + assert_eq!( + output.events, + vec![ + UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta { + tool_index: 0, + name: Some("get_weather".to_string()), + arguments: String::new(), + }), + UnifiedParserEvent::ToolCall(crate::tool::ToolCallDelta { + tool_index: 0, + name: None, + arguments: r#"{"location":"Paris"}"#.to_string(), + }), + ] + ); + } + + #[test] + fn combined_parser_preserves_tool_output_on_parse_error() { + let mut parser = CombinedParser::new(None, Some(Box::new(PartialThenErrorToolParser))); + let mut output = UnifiedParserOutput::default(); + + let error = parser.parse_into("bad", &mut output).unwrap_err(); + + assert!(matches!(error, crate::unified::UnifiedParserError::Tool(_))); + assert_eq!( + output.events, + vec![UnifiedParserEvent::Text("committed".to_string())] + ); + } + + #[test] + fn combined_parser_preserves_special_tokens_when_either_inner_parser_needs_it() { + let mut parser = CombinedParser::new(Some(Box::new(PreserveReasoningParser)), None); + assert!(parser.preserve_special_tokens()); + + parser = CombinedParser::new(None, Some(Box::new(PreserveToolParser))); + assert!(parser.preserve_special_tokens()); + } +} diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs new file mode 100644 index 00000000000..49955b4d818 --- /dev/null +++ b/rust/src/parser/src/unified/mod.rs @@ -0,0 +1,200 @@ +//! Unified parser interface for reasoning and tool-call deltas. + +mod combined; + +use thiserror::Error; +use vllm_tokenizer::DynTokenizer; + +pub use combined::CombinedParser; + +use crate::reasoning::ReasoningError; +use crate::tool::{ + StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput, +}; + +/// Result alias for unified parser operations. +pub type Result = std::result::Result; + +/// One parsed event emitted by a unified parser. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum UnifiedParserEvent { + /// Normal assistant-visible text. + Text(String), + /// Reasoning text hidden from the normal content stream. + Reasoning(String), + /// A tool-call update extracted from visible assistant text. + ToolCall(ToolCallDelta), +} + +/// Result of advancing unified parsing with one assistant-text input. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct UnifiedParserOutput { + /// Ordered parser events committed by this input. + pub events: Vec, +} + +impl UnifiedParserOutput { + /// Append one visible text event if `delta` is non-empty. + pub fn push_text(&mut self, delta: impl AsRef + Into) { + if delta.as_ref().is_empty() { + return; + } + if let Some(UnifiedParserEvent::Text(last_text)) = self.events.last_mut() { + last_text.push_str(delta.as_ref()); + return; + } + self.events.push(UnifiedParserEvent::Text(delta.into())); + } + + /// Append one reasoning text event if `delta` is non-empty. + pub fn push_reasoning(&mut self, delta: impl AsRef + Into) { + if delta.as_ref().is_empty() { + return; + } + if let Some(UnifiedParserEvent::Reasoning(last_text)) = self.events.last_mut() { + last_text.push_str(delta.as_ref()); + return; + } + self.events.push(UnifiedParserEvent::Reasoning(delta.into())); + } + + /// Append one tool-call event. + pub fn push_call(&mut self, call: ToolCallDelta) { + self.events.push(UnifiedParserEvent::ToolCall(call)); + } + + /// Append parsed tool parser output as unified events. + pub fn append_tool_output(&mut self, output: ToolParserOutput) { + for event in output.events { + match event { + ToolParserEvent::Text(text) => self.push_text(text), + ToolParserEvent::ToolCall(call) => self.push_call(call), + } + } + } + + /// Append another parser output onto this one. + pub fn append(&mut self, other: Self) { + for event in other.events { + match event { + UnifiedParserEvent::Text(text) => self.push_text(text), + UnifiedParserEvent::Reasoning(reasoning) => self.push_reasoning(reasoning), + UnifiedParserEvent::ToolCall(call) => self.push_call(call), + } + } + } +} + +#[cfg(test)] +mod tests { + use super::{UnifiedParserEvent, UnifiedParserOutput}; + use crate::tool::ToolCallDelta; + + #[test] + fn unified_parser_output_coalesces_adjacent_text_events() { + let mut output = UnifiedParserOutput::default(); + output.push_text("hello"); + output.push_text(" "); + output.push_text("world"); + output.push_reasoning("think"); + output.push_reasoning("ing"); + output.push_call(ToolCallDelta { + tool_index: 0, + name: Some("lookup".to_string()), + arguments: "{}".to_string(), + }); + output.push_text("!"); + + assert_eq!( + output.events, + vec![ + UnifiedParserEvent::Text("hello world".to_string()), + UnifiedParserEvent::Reasoning("thinking".to_string()), + UnifiedParserEvent::ToolCall(ToolCallDelta { + tool_index: 0, + name: Some("lookup".to_string()), + arguments: "{}".to_string(), + }), + UnifiedParserEvent::Text("!".to_string()), + ] + ); + } + + #[test] + fn unified_parser_output_append_coalesces_adjacent_events() { + let mut output = UnifiedParserOutput::default(); + output.push_text("hello"); + + let mut other = UnifiedParserOutput::default(); + other.push_text(" "); + other.push_text("world"); + other.push_reasoning("think"); + output.append(other); + + let mut after_reasoning = UnifiedParserOutput::default(); + after_reasoning.push_reasoning("ing"); + after_reasoning.push_text("!"); + output.append(after_reasoning); + + assert_eq!( + output.events, + vec![ + UnifiedParserEvent::Text("hello world".to_string()), + UnifiedParserEvent::Reasoning("thinking".to_string()), + UnifiedParserEvent::Text("!".to_string()), + ] + ); + } +} + +/// Incremental parser that extracts reasoning and tool-call events from assistant output. +pub trait UnifiedParser: Send { + /// Construct a boxed parser instance for one request stream. + fn create(tools: &[Tool], tokenizer: DynTokenizer) -> Result> + where + Self: Sized + 'static; + + /// Initialize parser state from prompt token IDs before output deltas arrive. + fn initialize(&mut self, _prompt_token_ids: &[u32]) -> Result<()> { + Ok(()) + } + + /// Return whether decoded output must preserve tokenizer special tokens. + fn preserve_special_tokens(&self) -> bool { + false + } + + /// Return the xgrammar structural-tag model used for strict tool calling. + fn structural_tag_model(&self) -> Option { + None + } + + /// Return the parser-provided ID for a tool call by index, if the model emitted one. + fn tool_call_id(&self, _tool_index: usize) -> Option<&str> { + None + } + + /// Feed one decoded text delta into the parser, appending committed output into `output`. + fn parse_into(&mut self, delta: &str, output: &mut UnifiedParserOutput) -> Result<()>; + + /// Flush any buffered parser state at end of stream. + fn finish(&mut self) -> Result { + Ok(UnifiedParserOutput::default()) + } + + /// Clear parser state and return currently uncommitted buffered text. + fn reset(&mut self) -> String { + String::new() + } +} + +/// Errors produced while creating or running unified parsers. +#[derive(Debug, Error)] +pub enum UnifiedParserError { + #[error("combined parser is constructed from split parser instances")] + CombinedParserConstructor, + #[error(transparent)] + Reasoning(#[from] ReasoningError), + #[error(transparent)] + Tool(#[from] ToolParserError), +} diff --git a/rust/src/reasoning-parser/Cargo.toml b/rust/src/reasoning-parser/Cargo.toml deleted file mode 100644 index d6500a7b0c1..00000000000 --- a/rust/src/reasoning-parser/Cargo.toml +++ /dev/null @@ -1,12 +0,0 @@ -[package] -name = "vllm-reasoning-parser" -version.workspace = true -edition.workspace = true -license.workspace = true - -[dependencies] -thiserror.workspace = true -vllm-tokenizer.workspace = true - -[lints] -workspace = true diff --git a/rust/src/tool-parser/src/tests.rs b/rust/src/tool-parser/src/tests.rs deleted file mode 100644 index fb9c8e62bf3..00000000000 --- a/rust/src/tool-parser/src/tests.rs +++ /dev/null @@ -1,109 +0,0 @@ -use super::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; -use crate::ToolParserTestExt as _; - -struct DefaultParser; - -impl ToolParser for DefaultParser { - fn create(_tools: &[Tool]) -> Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self)) - } - - fn parse_into(&mut self, _chunk: &str, _output: &mut ToolParserOutput) -> Result<()> { - Ok(()) - } - - fn finish(&mut self) -> Result { - Ok(ToolParserOutput::default()) - } - - fn reset(&mut self) -> String { - String::new() - } -} - -#[test] -fn tool_parser_does_not_preserve_special_tokens_by_default() { - let parser = DefaultParser; - - assert!(!parser.preserve_special_tokens()); -} - -#[test] -fn default_parse_complete_delegates_through_parse_chunk_and_finish() { - struct StreamingParser; - - impl ToolParser for StreamingParser { - fn create(_tools: &[Tool]) -> Result> - where - Self: Sized + 'static, - { - Ok(Box::new(Self)) - } - - fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> { - output.normal_text.push_str("prefix "); - output.calls.extend([ - ToolCallDelta { - tool_index: 0, - name: Some("weather".to_string()), - arguments: "{\"location\":".to_string(), - }, - ToolCallDelta { - tool_index: 0, - name: None, - arguments: "\"Paris\"".to_string(), - }, - ToolCallDelta { - tool_index: 1, - name: Some("time".to_string()), - arguments: "{\"timezone\":".to_string(), - }, - ]); - Ok(()) - } - - fn finish(&mut self) -> Result { - Ok(ToolParserOutput { - normal_text: "suffix".to_string(), - calls: vec![ - ToolCallDelta { - tool_index: 0, - name: None, - arguments: "}".to_string(), - }, - ToolCallDelta { - tool_index: 1, - name: None, - arguments: "\"UTC\"}".to_string(), - }, - ], - }) - } - - fn reset(&mut self) -> String { - String::new() - } - } - - let mut parser = StreamingParser; - let output = parser.parse_complete("ignored").unwrap(); - assert_eq!(output.normal_text, "prefix suffix"); - assert_eq!( - output.calls, - vec![ - ToolCallDelta { - tool_index: 0, - name: Some("weather".to_string()), - arguments: "{\"location\":\"Paris\"}".to_string(), - }, - ToolCallDelta { - tool_index: 1, - name: Some("time".to_string()), - arguments: "{\"timezone\":\"UTC\"}".to_string(), - }, - ] - ); -} diff --git a/tests/basic_correctness/test_basic_correctness.py b/tests/basic_correctness/test_basic_correctness.py index 1a07ac6da6b..810a3a0aeed 100644 --- a/tests/basic_correctness/test_basic_correctness.py +++ b/tests/basic_correctness/test_basic_correctness.py @@ -29,7 +29,53 @@ MODELS = [ "meta-llama/Llama-3.2-1B-Instruct", ] -TARGET_TEST_SUITE = os.environ.get("TARGET_TEST_SUITE", "L4") +TARGET_TEST_SUITE_ENV = "VLLM_TARGET_TEST_SUITE" +LEGACY_TARGET_TEST_SUITE_ENV = "TARGET_TEST_SUITE" + +GENERIC_DISTRIBUTED_TEST_SUITES = ("L4", "MI250", "MI300", "MI325", "MI355") +ALL_DISTRIBUTED_TEST_SUITES = (*GENERIC_DISTRIBUTED_TEST_SUITES, "A100") + + +def _default_target_test_suite() -> str: + if not current_platform.is_rocm(): + return "L4" + + try: + device_name = current_platform.get_device_name().upper() + except Exception: + device_name = "" + + if "MI355" in device_name: + return "MI355" + if "MI300" in device_name: + return "MI300" + if "MI325" in device_name: + return "MI325" + if "MI250" in device_name: + return "MI250" + + try: + from vllm.platforms import rocm as rocm_platform + + if rocm_platform.on_gfx950(): + return "MI355" + if rocm_platform.on_gfx942(): + return "MI300" + except Exception: + pass + + return "MI250" + + +def _resolve_target_test_suite() -> str: + for env_name in (TARGET_TEST_SUITE_ENV, LEGACY_TARGET_TEST_SUITE_ENV): + value = os.environ.get(env_name, "").strip().upper() + if value: + return value + return _default_target_test_suite() + + +TARGET_TEST_SUITE = _resolve_target_test_suite() def test_vllm_gc_ed(): @@ -131,14 +177,27 @@ def test_models( @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize( - "model, distributed_executor_backend, attention_backend, test_suite, extra_env", + ( + "model, distributed_executor_backend, attention_backend, " + "target_test_suites, extra_env" + ), [ - ("facebook/opt-125m", "ray", "", "L4", {}), - ("facebook/opt-125m", "mp", "", "L4", {}), - ("meta-llama/Llama-3.2-1B-Instruct", "ray", "", "L4", {}), - ("meta-llama/Llama-3.2-1B-Instruct", "mp", "", "L4", {}), - ("facebook/opt-125m", "ray", "", "A100", {}), - ("facebook/opt-125m", "mp", "", "A100", {}), + ("facebook/opt-125m", "ray", "", ALL_DISTRIBUTED_TEST_SUITES, {}), + ("facebook/opt-125m", "mp", "", ALL_DISTRIBUTED_TEST_SUITES, {}), + ( + "meta-llama/Llama-3.2-1B-Instruct", + "ray", + "", + GENERIC_DISTRIBUTED_TEST_SUITES, + {}, + ), + ( + "meta-llama/Llama-3.2-1B-Instruct", + "mp", + "", + GENERIC_DISTRIBUTED_TEST_SUITES, + {}, + ), ], ) @pytest.mark.parametrize("enable_prompt_embeds", [True, False]) @@ -150,19 +209,19 @@ def test_models_distributed( model: str, distributed_executor_backend: str, attention_backend: str, - test_suite: str, + target_test_suites: tuple[str, ...], extra_env: dict[str, str], enable_prompt_embeds: bool, ) -> None: - if test_suite != TARGET_TEST_SUITE: - pytest.skip(f"Skip test for {test_suite}") + if TARGET_TEST_SUITE and TARGET_TEST_SUITE not in target_test_suites: + pytest.skip(f"Skip test for {TARGET_TEST_SUITE}") with monkeypatch.context() as monkeypatch_context: if ( model == "meta-llama/Llama-3.2-1B-Instruct" and distributed_executor_backend == "ray" and attention_backend == "" - and test_suite == "L4" + and TARGET_TEST_SUITE == "L4" and enable_prompt_embeds ): # noqa pytest.skip("enable_prompt_embeds does not work with ray compiled dag.") diff --git a/tests/conftest.py b/tests/conftest.py index 4fc43ef04b1..4b92f285fac 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -64,6 +64,7 @@ from vllm.logprobs import Logprob from vllm.multimodal.media import MediaWithBytes from vllm.multimodal.utils import fetch_image from vllm.outputs import RequestOutput +from vllm.platforms import current_platform from vllm.sampling_params import BeamSearchParams from vllm.transformers_utils.utils import maybe_model_redirect from vllm.utils.collection_utils import is_list_of @@ -852,6 +853,24 @@ class HfRunner: return self.model.predict(prompts, *args, convert_to_tensor=True, **kwargs) def __enter__(self): + if current_platform.is_rocm(): + # Record starting memory usage stats on ROCm so that we can wait for + # memory to roughly settle back below these levels on shutdown. This is + # helpful in cases where the HfRunner is initialized after significant GPU + # memory is already occupied, e.g. in + # tests/basic_correctness/test_basic_correctness.py::test_models_distributed + from tests.utils import ( + get_physical_device_indices, + record_gpu_memory_usage_stats, + ) + + if (device_count := current_platform.device_count()) > 0: + devices = get_physical_device_indices(devices=list(range(device_count))) + mem_usage_stats = record_gpu_memory_usage_stats(devices=devices) + self.threshold_ratios = { + device: 0.05 + mem_used / mem_tot + for device, (mem_used, mem_tot) in mem_usage_stats.items() + } return self def __exit__(self, exc_type, exc_value, traceback): @@ -861,7 +880,11 @@ class HfRunner: cleanup_dist_env_and_memory() # ROCm frees VRAM lazily; wait so a runner started right after this HF # model exits does not OOM on its startup memory guard. - wait_for_rocm_memory_to_settle() + wait_for_rocm_memory_to_settle( + threshold_ratio=getattr(self, "threshold_ratios", None) + ) + if hasattr(self, "threshold_ratios"): + del self.threshold_ratios @pytest.fixture(scope="session") diff --git a/tests/distributed/test_context_parallel.py b/tests/distributed/test_context_parallel.py index a2863092177..484d29c5b53 100644 --- a/tests/distributed/test_context_parallel.py +++ b/tests/distributed/test_context_parallel.py @@ -13,13 +13,14 @@ import os from dataclasses import dataclass from typing import Literal, NamedTuple +import lm_eval import pytest import torch -from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k from tests.utils import RemoteOpenAIServer, create_new_process_for_each_test from vllm.config.model import RunnerOption from vllm.logger import init_logger +from vllm.platforms import current_platform from ..models.registry import HF_EXAMPLE_MODELS @@ -35,8 +36,10 @@ CP_TEST_MODELS = [ ] # GSM8K eval configuration -NUM_QUESTIONS = 256 # Fast eval for CI NUM_SHOTS = 5 # Few-shot examples +TASK = "gsm8k" +FILTER = "exact_match,strict-match" +NUM_CONCURRENT = 128 # tp accuracy with 2% buffer MIN_ACCURACY = { # .buildkite/lm-eval-harness/configs/DeepSeek-V2-Lite-Chat.yaml @@ -121,24 +124,34 @@ class CPTestSettings: ) -CP_TEXT_GENERATION_MODELS = { - "deepseek-ai/DeepSeek-V2-Lite-Chat": [ - CPTestSettings.detailed(dcp_multipliers=[1]), - CPTestSettings.detailed( - dcp_multipliers=[0.5], - cp_kv_cache_interleave_size=64, - attn_backend="FLASHMLA", - ), - ], - "Qwen/Qwen2.5-1.5B-Instruct": [ - CPTestSettings.detailed( - cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN" - ), - CPTestSettings.detailed( - cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER" - ), - ], -} +if current_platform.is_rocm(): + CP_TEXT_GENERATION_MODELS = { + "deepseek-ai/DeepSeek-V2-Lite-Chat": [ + CPTestSettings.detailed(dcp_multipliers=[1]), + ], + "Qwen/Qwen2.5-1.5B-Instruct": [ + CPTestSettings.detailed(dcp_multipliers=[1]), + ], + } +else: + CP_TEXT_GENERATION_MODELS = { + "deepseek-ai/DeepSeek-V2-Lite-Chat": [ + CPTestSettings.detailed(dcp_multipliers=[1]), + CPTestSettings.detailed( + dcp_multipliers=[0.5], + cp_kv_cache_interleave_size=64, + attn_backend="FLASHMLA", + ), + ], + "Qwen/Qwen2.5-1.5B-Instruct": [ + CPTestSettings.detailed( + cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN" + ), + CPTestSettings.detailed( + cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER" + ), + ], + } def _test_cp_gsm8k( @@ -227,19 +240,23 @@ def _test_cp_gsm8k( server_args, max_wait_seconds=720, ) as remote_server: - host = f"http://{remote_server.host}" - port = remote_server.port + url = f"{remote_server.url_for('v1')}/completions" - # Run GSM8K evaluation - results = evaluate_gsm8k( - num_questions=NUM_QUESTIONS, - num_shots=NUM_SHOTS, - host=host, - port=port, + model_args = ( + f"model={model_id}," + f"base_url={url}," + f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False" + ) + + results = lm_eval.simple_evaluate( + model="local-completions", + model_args=model_args, + tasks=TASK, + num_fewshot=NUM_SHOTS, ) # Validate accuracy is reasonable - accuracy = results["accuracy"] + accuracy = results["results"][TASK][FILTER] min_accuracy = MIN_ACCURACY[model_id] assert accuracy >= min_accuracy, ( f"TP+DCP accuracy too low: {accuracy:.3f} < {min_accuracy:.3f}" diff --git a/tests/distributed/test_pipeline_parallel.py b/tests/distributed/test_pipeline_parallel.py index 28c905baf73..75f05ca5069 100644 --- a/tests/distributed/test_pipeline_parallel.py +++ b/tests/distributed/test_pipeline_parallel.py @@ -102,8 +102,6 @@ class PPTestSettings: TEXT_GENERATION_MODELS = { # [Decoder-only] - # Uses Llama - # "BAAI/AquilaChat-7B": PPTestSettings.fast(), "Snowflake/snowflake-arctic-instruct": PPTestSettings.fast(load_format="dummy"), "bigscience/bloomz-1b1": PPTestSettings.fast(), "zai-org/chatglm3-6b": PPTestSettings.fast(), diff --git a/tests/distributed/test_pp_cudagraph.py b/tests/distributed/test_pp_cudagraph.py index 34ae305c2d2..2f0fc9a1b5d 100644 --- a/tests/distributed/test_pp_cudagraph.py +++ b/tests/distributed/test_pp_cudagraph.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from typing_extensions import LiteralString + +from vllm.platforms import current_platform from ..utils import compare_two_settings, create_new_process_for_each_test @@ -14,15 +15,13 @@ from ..utils import compare_two_settings, create_new_process_for_each_test ) @pytest.mark.parametrize( "ATTN_BACKEND", - [ - "FLASH_ATTN", - ], + [None] if current_platform.is_rocm() else ["FLASH_ATTN"], ) @create_new_process_for_each_test() def test_pp_cudagraph( PP_SIZE: int, MODEL_NAME: str, - ATTN_BACKEND: LiteralString, + ATTN_BACKEND: str | None, ): cudagraph_args = [ # use half precision for speed and memory savings in CI environment @@ -32,8 +31,10 @@ def test_pp_cudagraph( str(PP_SIZE), "--distributed-executor-backend", "mp", - f"--attention-backend={ATTN_BACKEND}", ] + # On ROCm, defer to the platform attention selector instead of forcing a backend. + if ATTN_BACKEND is not None: + cudagraph_args.append(f"--attention-backend={ATTN_BACKEND}") eager_args = cudagraph_args + ["--enforce-eager"] diff --git a/tests/distributed/test_quick_all_reduce.py b/tests/distributed/test_quick_all_reduce.py index 86eb82c962e..bfa28cc5c44 100644 --- a/tests/distributed/test_quick_all_reduce.py +++ b/tests/distributed/test_quick_all_reduce.py @@ -350,7 +350,7 @@ def bf16_cast_quickreduce( @pytest.mark.skipif( not current_platform.is_rocm(), reason="only test quick allreduce for rocm" ) -@pytest.mark.parametrize("quant_mode", ["FP", "INT8", "INT6", "INT4"]) +@pytest.mark.parametrize("quant_mode", ["FP", "INT8", "INT6", "INT4", "INT3"]) @pytest.mark.parametrize("tp_size", [2]) @pytest.mark.parametrize("pipeline_parallel_size", [1, 2]) @pytest.mark.parametrize("test_target", [graph_quickreduce, eager_quickreduce]) @@ -438,7 +438,7 @@ def qr_variable_input(rank, world_size): s2 = 2048 inp1 = torch.ones((s1, s2), dtype=dtype, device=device_idx) result = torch.empty_like(inp1) - # FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 NONE = 4 + # FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 INT3 = 4 ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True) try: if inp1[0, 0] == 0: diff --git a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py index 62e4965b8ed..33cb576f351 100644 --- a/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py +++ b/tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py @@ -24,6 +24,7 @@ tools = [ "description": "Get the current weather in a given location", "parameters": { "type": "object", + "strict": True, "properties": { "city": { "type": "string", @@ -215,80 +216,6 @@ async def test_function_tool_use( assert len(reasoning) > 0 -@pytest.fixture(scope="module") -def k2_server(): - args = [ - # use half precision for speed and memory savings in CI environment - "--dtype", - "half", - "--enable-auto-tool-choice", - "--structured-outputs-config.backend", - "xgrammar", - "--tool-call-parser", - "hermes", - "--reasoning-parser", - "qwen3", - "--gpu-memory-utilization", - "0.4", - ] + ROCM_EXTRA_ARGS - # Test kimi_k2 tool use tool_id format by overriding model_type. - # is_deepseek_mla safely returns False via getattr when kv_lora_rank - # is absent from the underlying config. - with RemoteOpenAIServer( - MODEL_NAME, - args, - env_dict=ROCM_ENV_OVERRIDES, - override_hf_configs={"model_type": "kimi_k2"}, - ) as remote_server: - yield remote_server - - -@pytest_asyncio.fixture -async def k2_client(k2_server): - async with k2_server.get_async_client() as async_client: - yield async_client - - -@pytest.mark.asyncio -@pytest.mark.skip(reason="Skipping Kimi K2 tool ID test") -@pytest.mark.parametrize("model_name", [MODEL_NAME]) -@pytest.mark.parametrize("stream", [True, False]) -@pytest.mark.parametrize("tool_choice", ["required"]) -async def test_tool_id_kimi_k2( - k2_client: openai.AsyncOpenAI, model_name: str, stream: bool, tool_choice: str -): - if not stream: - # Non-streaming test - chat_completion = await k2_client.chat.completions.create( - messages=messages, model=model_name, tools=tools, tool_choice=tool_choice - ) - assert chat_completion.choices[0].message.tool_calls is not None - assert len(chat_completion.choices[0].message.tool_calls) > 0 - assert chat_completion.choices[0].message.tool_calls[0].id in [ - "functions.get_current_weather:0", - "functions.get_forecast:1", - ] - else: - # Streaming test - output_stream = await k2_client.chat.completions.create( - messages=messages, - model=model_name, - tools=tools, - tool_choice=tool_choice, - stream=True, - ) - - output = [] - async for chunk in output_stream: - if chunk.choices and chunk.choices[0].delta.tool_calls: - output.extend(chunk.choices[0].delta.tool_calls) - for o in output: - assert o.id is None or o.id in [ - "functions.get_current_weather:0", - "functions.get_forecast:1", - ] - - @pytest.mark.asyncio @pytest.mark.parametrize("model_name", [MODEL_NAME]) @pytest.mark.parametrize("arguments", ["{}", ""]) diff --git a/tests/entrypoints/openai/responses/test_parsable_context_unit.py b/tests/entrypoints/openai/responses/test_parsable_context_unit.py index 0aadfbe99d3..2bad3032c46 100644 --- a/tests/entrypoints/openai/responses/test_parsable_context_unit.py +++ b/tests/entrypoints/openai/responses/test_parsable_context_unit.py @@ -183,11 +183,19 @@ def _make_request_output( def _make_context(parser_cls, **overrides): + # ParsableContext no longer lazily builds a parser from ``parser_cls``; + # the caller (here, the serving layer in production) must supply one. + request = overrides.get("request", _make_request()) + response_parser = overrides.pop("response_parser", None) + if response_parser is None and parser_cls is not None: + response_parser = parser_cls(MagicMock(), request.tools) + defaults = dict( tokenizer=MagicMock(), parser_cls=parser_cls, + response_parser=response_parser, response_messages=[], - request=_make_request(), + request=request, available_tools=None, chat_template=None, chat_template_content_format="auto", diff --git a/tests/kernels/attention/test_attention.py b/tests/kernels/attention/test_attention.py index 9ddceef8fb3..0bc2461463c 100644 --- a/tests/kernels/attention/test_attention.py +++ b/tests/kernels/attention/test_attention.py @@ -26,11 +26,11 @@ PARTITION_SIZE_ROCM = 256 DTYPES = [torch.bfloat16] NUM_GEN_SEQS = [7] # Arbitrary values for testing NUM_PREFILL_SEQS = [3] # Arbitrary values for testing -NUM_HEADS = [(40, 40), (64, 8)] # Arbitrary values for testing +NUM_HEADS = [(32, 8), (40, 40), (64, 8)] # Arbitrary values for testing # This should be sync with get_supported_head_sizes() in # vllm.v1.attention.ops.paged_attn.PagedAttention -HEAD_SIZES = [32, 80, 128, 256] +HEAD_SIZES = [32, 64, 80, 128, 256] BLOCK_SIZES = [16, 32] USE_ALIBI = [False, True] @@ -353,8 +353,10 @@ def test_paged_attention( kv_cache_dtype, k_scale, v_scale, + None, + "f16", ), - cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]), + cond=(head_size == 64 and block_size == BLOCK_SIZES[0]), ) else: diff --git a/tests/kernels/attention/test_cache.py b/tests/kernels/attention/test_cache.py index 4cbeb7a0b97..7558da1c600 100644 --- a/tests/kernels/attention/test_cache.py +++ b/tests/kernels/attention/test_cache.py @@ -1015,6 +1015,78 @@ def test_gather_and_maybe_dequant_cache_mla( torch.testing.assert_close(dst, expected) +@pytest.mark.parametrize("kv_lora_rank", [512]) +@pytest.mark.parametrize("qk_rope_head_dim", [64]) +@pytest.mark.parametrize("block_size", [16]) +@pytest.mark.parametrize("num_blocks", [128]) +@pytest.mark.parametrize("dtype", [torch.float32]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("device", CUDA_DEVICES) +@torch.inference_mode() +def test_gather_and_maybe_dequant_cache_mla_with_seq_starts( + kv_lora_rank, + qk_rope_head_dim, + block_size, + num_blocks, + dtype, + kv_cache_dtype, + device, +): + entry_size = kv_lora_rank + qk_rope_head_dim + scale = torch.tensor(0.1, dtype=torch.float32, device=device) + src_cache = _create_mla_cache( + num_blocks, block_size, entry_size, dtype, kv_cache_dtype, device + ) + _fill_mla_cache(src_cache, kv_cache_dtype=kv_cache_dtype) + + seq_starts = torch.tensor([3, 17, 5], dtype=torch.int32, device=device) + seq_lens = torch.tensor([20, 10, 16], dtype=torch.int32, device=device) + batch_size = seq_lens.shape[0] + total_tokens = seq_lens.sum().item() + cu_seq_lens = torch.empty((batch_size + 1), dtype=torch.int32, device=device) + cu_seq_lens[0] = 0 + cu_seq_lens[1:] = seq_lens.cumsum(dim=0) + token_to_seq = torch.repeat_interleave( + torch.arange(batch_size, dtype=torch.int32, device=device), seq_lens + ) + + block_table = torch.empty( + (batch_size, num_blocks), dtype=torch.int32, device=device + ) + for b in range(batch_size): + block_table[b, :] = torch.randperm(num_blocks, device=device) + + if kv_cache_dtype == "fp8": + dequant_src_cache = torch.empty_like(src_cache, dtype=dtype) + ops.convert_fp8(dequant_src_cache, src_cache, scale.item()) + else: + dequant_src_cache = src_cache + + expected_rows = [] + for b in range(batch_size): + start = seq_starts[b].item() + length = seq_lens[b].item() + for offset in range(start, start + length): + block_id = block_table[b, offset // block_size] + slot = offset % block_size + expected_rows.append(dequant_src_cache[block_id, slot]) + expected = torch.stack(expected_rows) + + dst = torch.zeros((total_tokens, entry_size), dtype=dtype, device=device) + ops.gather_and_maybe_dequant_cache( + src_cache, + dst, + block_table, + cu_seq_lens, + token_to_seq, + total_tokens, + kv_cache_dtype, + scale, + seq_starts, + ) + torch.testing.assert_close(dst, expected) + + @pytest.mark.parametrize("kv_lora_rank", [512]) @pytest.mark.parametrize("qk_rope_head_dim", [64]) @pytest.mark.parametrize("block_size", [16]) diff --git a/tests/kernels/helion/test_utils.py b/tests/kernels/helion/test_utils.py index 540cc4f8bc7..f357fbf6473 100644 --- a/tests/kernels/helion/test_utils.py +++ b/tests/kernels/helion/test_utils.py @@ -17,6 +17,7 @@ from vllm.kernels.helion.utils import canonicalize_gpu_name ("NVIDIA H100 SXM5", "nvidia_h100"), ("NVIDIA GeForce RTX 4090", "nvidia_geforce_rtx_4090"), ("AMD Instinct MI300X", "amd_instinct_mi300x"), + ("AMD Instinct MI250X / MI250", "amd_instinct_mi250x_mi250"), ("Tesla V100-SXM2-32GB", "tesla_v100"), ], ) diff --git a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py index 4b800b192b2..bd30bc4f1ce 100644 --- a/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py +++ b/tests/kernels/mamba/cpu/test_cpu_gdn_ops.py @@ -25,6 +25,8 @@ HEAD_DIMS = [ (64, 32), ] CHUNK_SIZE = 64 +CONV_DIM = 128 +CONV_KERNEL = 4 PREFILL_SEQ_LENS = [ [1], [1, 2, 3], @@ -312,3 +314,225 @@ def test_chunk_gated_delta_rule_cpu( atol=1e-2, rtol=1e-2, ) + + +# (total_tokens, split) pairs mimicking where chunked prefill breaks a sequence +# across two scheduler steps: chunk-aligned and non-aligned splits. +TWO_CALL_SPLITS = [ + (2 * CHUNK_SIZE, CHUNK_SIZE), + (2 * CHUNK_SIZE + 17, CHUNK_SIZE), + (2 * CHUNK_SIZE + 17, CHUNK_SIZE + 9), + (4 * CHUNK_SIZE + 17, 2 * CHUNK_SIZE), + (3 * CHUNK_SIZE, CHUNK_SIZE + 1), +] + + +@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) +@pytest.mark.parametrize("num_heads", NUM_HEADS) +@pytest.mark.parametrize("head_dims", HEAD_DIMS) +@torch.inference_mode() +def test_chunk_gated_delta_rule_cpu_two_call_split( + total_tokens: int, + split: int, + num_heads: tuple[int, int], + head_dims: tuple[int, int], +) -> None: + """A prefill split into two calls (the second seeded with the first's + ``final_state`` and a rebased ``cu_seqlens``) must match the single-call + result, mimicking the cross-scheduler-step handoff in + ``cpu_gdn_attention_core``. + """ + q, k, v, a, b, A_log, dt_bias = gdn_inputs( + num_tokens=total_tokens, + num_heads=num_heads, + head_dims=head_dims, + ) + _, num_v_heads = num_heads + head_dim, v_head_dim = head_dims + + g, beta = ref_gdn_gating(A_log, a, b, dt_bias) + g = g.unsqueeze(0) # [1, T, HV] + beta = beta.unsqueeze(0) + + zero_state = torch.zeros(1, num_v_heads, head_dim, v_head_dim, dtype=torch.float32) + + # Reference: whole sequence in one call, no initial state. + out_full, final_full = ops.chunk_gated_delta_rule_cpu( + query=q, + key=k, + value=v, + g=g, + beta=beta, + initial_state=zero_state, + output_final_state=True, + cu_seqlens=torch.tensor([0, total_tokens], dtype=torch.int32), + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + + # Call 1: tokens [0:split], no initial state, capture final state. + out1, state1 = ops.chunk_gated_delta_rule_cpu( + query=q[:, :split], + key=k[:, :split], + value=v[:, :split], + g=g[:, :split], + beta=beta[:, :split], + initial_state=zero_state, + output_final_state=True, + cu_seqlens=torch.tensor([0, split], dtype=torch.int32), + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + # Call 2: tokens [split:T] seeded with call 1's final state and a cu_seqlens + # rebased to start at 0, as cpu_gdn_attention_core continues a prefill chunk. + tail = total_tokens - split + out2, state2 = ops.chunk_gated_delta_rule_cpu( + query=q[:, split:], + key=k[:, split:], + value=v[:, split:], + g=g[:, split:], + beta=beta[:, split:], + initial_state=state1.to(torch.float32), + output_final_state=True, + cu_seqlens=torch.tensor([0, tail], dtype=torch.int32), + head_first=False, + use_qk_l2norm_in_kernel=True, + ) + + out_split = torch.cat([out1, out2], dim=1) + + # State must be near-exact; output allows a looser bound for the bf16 round-trip. + torch.testing.assert_close(state2, final_full, atol=1e-3, rtol=1e-3) + torch.testing.assert_close(out_split, out_full, atol=2e-2, rtol=2e-2) + + +def _conv_inputs(total_tokens: int): + x = tensor_cache(total_tokens * CONV_DIM, torch.bfloat16).view( + total_tokens, CONV_DIM + ) + weight = tensor_cache(CONV_DIM * CONV_KERNEL, torch.bfloat16).view( + CONV_DIM, CONV_KERNEL + ) + bias = tensor_cache(CONV_DIM, torch.bfloat16) + return x, weight, bias + + +@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) +@torch.inference_mode() +def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> None: + """Non-AMX conv-state handoff: a two-call split (the second seeded via + ``has_initial_state=True`` from the conv_states the first wrote back) must + match the single-call result. + """ + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_torch, + ) + + x, weight, bias = _conv_inputs(total_tokens) + state_len = CONV_KERNEL - 1 + # [num_slots, conv_dim, state_len]; slot 0 used here. + conv_states_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype) + conv_states_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype) + + # x is [conv_dim, T] for causal_conv1d_torch. + xt = x.transpose(0, 1).contiguous() + + out_full = causal_conv1d_torch( + x=xt, + weight=weight, + bias=bias, + conv_states=conv_states_full, + query_start_loc=torch.tensor([0, total_tokens], dtype=torch.int32), + cache_indices=torch.tensor([0], dtype=torch.int32), + has_initial_state=torch.tensor([False]), + activation="silu", + ) + + out1 = causal_conv1d_torch( + x=xt[:, :split], + weight=weight, + bias=bias, + conv_states=conv_states_split, + query_start_loc=torch.tensor([0, split], dtype=torch.int32), + cache_indices=torch.tensor([0], dtype=torch.int32), + has_initial_state=torch.tensor([False]), + activation="silu", + ) + out2 = causal_conv1d_torch( + x=xt[:, split:], + weight=weight, + bias=bias, + conv_states=conv_states_split, + query_start_loc=torch.tensor([0, total_tokens - split], dtype=torch.int32), + cache_indices=torch.tensor([0], dtype=torch.int32), + has_initial_state=torch.tensor([True]), + activation="silu", + ) + out_split = torch.cat([out1, out2], dim=1) + + torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not torch.cpu._is_amx_tile_supported(), + reason="causal_conv1d_fwd_cpu requires AMX/AVX512", +) +@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS) +@torch.inference_mode() +def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) -> None: + """AMX prefill conv op must honor ``has_initial_state`` so a two-call split + matches the single-call result. + + Regression test for ``causal_conv1d_fwd_varlen_kernel_impl`` (``conv.cpp``) + ignoring the carried conv state on continued chunks. + """ + state_len = CONV_KERNEL - 1 + x, weight, bias = _conv_inputs(total_tokens) + + def amx(x_seg, conv_states, has_init): + seq = x_seg.shape[0] + return ops.causal_conv1d_fwd_cpu( + x=x_seg.transpose(0, 1), # [dim, seq]; stride(-2)==1 (view of [seq,dim]) + weight=weight, + bias=bias, + conv_states=conv_states, + query_start_loc=torch.tensor([0, seq], dtype=torch.int32), + cache_indices=torch.tensor([0], dtype=torch.int32), + has_initial_state=torch.tensor([has_init]), + silu_activation=True, + is_vnni=False, + ).contiguous() + + # conv_state layout passed by the AMX branch: [num_slots, dim, state_len]. + cs_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype) + out_full = amx(x, cs_full, False) + + cs_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype) + out1 = amx(x[:split], cs_split, False) + out2 = amx(x[split:], cs_split, True) + out_split = torch.cat([out1, out2], dim=1) + + torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2) + + +@torch.inference_mode() +def test_batch_memcpy_cpu_fallback() -> None: + """The ctypes batch_memcpy fallback (used when triton-cpu is absent) must + copy each src into its dst, validating the (src_ptrs, dst_ptrs, sizes) + argument order against ctypes.memmove(dst, src, size). + """ + from vllm.utils.cpu_triton_utils import batch_memcpy_kernel + + # Varied byte sizes, including a non-power-of-two run. + sizes_bytes = [256, 1024, 17 * 4, 4096] + srcs = [torch.rand(n // 4, dtype=torch.float32) for n in sizes_bytes] + dsts = [torch.zeros_like(s) for s in srcs] + + src_ptrs = torch.tensor([s.data_ptr() for s in srcs], dtype=torch.uint64) + dst_ptrs = torch.tensor([d.data_ptr() for d in dsts], dtype=torch.uint64) + sizes = torch.tensor(sizes_bytes, dtype=torch.int32) + + batch_memcpy_kernel[(len(srcs),)](src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=1024) + + for src, dst in zip(srcs, dsts): + torch.testing.assert_close(dst, src) diff --git a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py index eaeca6a8a5d..0f80ca5c55a 100644 --- a/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py +++ b/tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py @@ -5,9 +5,13 @@ import pytest import torch +from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( fused_marlin_moe, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( + TrtLlmMxint4ExpertsMonolithic, +) from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( grouped_topk, ) @@ -77,6 +81,14 @@ __all__ = [ ] +def test_trtllm_mxint4_activation_supports_vllm_gated_silu(): + assert TrtLlmMxint4ExpertsMonolithic._supports_activation(MoEActivation.SILU) + assert TrtLlmMxint4ExpertsMonolithic._supports_activation(MoEActivation.SWIGLUOAI) + assert not TrtLlmMxint4ExpertsMonolithic._supports_activation( + MoEActivation.RELU2_NO_MUL + ) + + def marlin_quantize_moe_weights( weights_bf16: torch.Tensor, group_size: int = 32 ) -> tuple[torch.Tensor, torch.Tensor]: diff --git a/tests/kernels/moe/test_moe_kernel_oracle.py b/tests/kernels/moe/test_moe_kernel_oracle.py new file mode 100644 index 00000000000..fbdf804a3b7 --- /dev/null +++ b/tests/kernels/moe/test_moe_kernel_oracle.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the MoEKernelOracle ABC introduced in PR series for #37753. + +This file contains a single canonical demonstration that +`UnquantizedMoEKernelOracle` methods delegate one-to-one to the +existing module-level functions in `oracle/unquantized.py`. Each method +on `UnquantizedMoEKernelOracle` follows the same `return module_fn(args)` +pattern, so verifying delegation for one method (`make_kernel`) gives +high confidence in the rest. +""" + +from unittest.mock import patch + +from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts +from vllm.model_executor.layers.fused_moe.oracle import UnquantizedMoEKernelOracle +from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, +) + + +class TestUnquantizedDelegation: + """UnquantizedMoEKernelOracle methods must delegate to the existing + module-level functions; behaviour is bit-identical.""" + + def test_make_kernel_delegates(self) -> None: + quant_config = object() + moe_config = object() + experts_cls = TritonExperts + sentinel_kernel = object() + + with patch( + "vllm.model_executor.layers.fused_moe.oracle.unquantized." + "make_unquantized_moe_kernel", + return_value=sentinel_kernel, + ) as mocked: + out = UnquantizedMoEKernelOracle().make_kernel( + quant_config, + moe_config, + UnquantizedMoeBackend.TRITON, + experts_cls, + ) + + mocked.assert_called_once_with( + quant_config, + moe_config, + UnquantizedMoeBackend.TRITON, + experts_cls, + None, # routing_tables default + ) + assert out is sentinel_kernel diff --git a/tests/kernels/quantization/test_triton_scaled_mm.py b/tests/kernels/quantization/test_triton_scaled_mm.py index d857d495f2d..1cef5eb93a5 100644 --- a/tests/kernels/quantization/test_triton_scaled_mm.py +++ b/tests/kernels/quantization/test_triton_scaled_mm.py @@ -60,10 +60,8 @@ def test_rocm_compressed_tensors_w8a8( vllm_runner, example_prompts, model_path, max_tokens, num_logprobs ): dtype = "bfloat16" - # Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179 - with vllm_runner( - model_path, dtype=dtype, attention_backend="TRITON_ATTN" - ) as vllm_model: + + with vllm_runner(model_path, dtype=dtype) as vllm_model: vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs) diff --git a/tests/models/registry.py b/tests/models/registry.py index 61aa4b75055..bd2cba46b67 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -201,8 +201,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { # [Decoder-only] "AfmoeForCausalLM": _HfExamplesInfo("arcee-ai/Trinity-Nano-Preview"), "ApertusForCausalLM": _HfExamplesInfo("swiss-ai/Apertus-8B-Instruct-2509"), - "AquilaModel": _HfExamplesInfo("BAAI/AquilaChat-7B", trust_remote_code=True), - "AquilaForCausalLM": _HfExamplesInfo("BAAI/AquilaChat2-7B", trust_remote_code=True), "ArceeForCausalLM": _HfExamplesInfo("arcee-ai/AFM-4.5B-Base"), "ArcticForCausalLM": _HfExamplesInfo( "Snowflake/snowflake-arctic-instruct", trust_remote_code=True diff --git a/tests/parser/engine/test_seed_oss.py b/tests/parser/engine/test_seed_oss.py new file mode 100644 index 00000000000..ce118faa443 --- /dev/null +++ b/tests/parser/engine/test_seed_oss.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the engine-based seed_oss parser. + +seed_oss is Qwen3 with four overridden wrapper tokens, so the shared grammar +(arg types, multiline values, parallel calls, streaming mechanics, …) is +already covered by ``test_qwen3.py``/``test_qwen3_reasoning.py``. These tests +cover only what is seed_oss-specific: that the ``seed:`` token overrides are +wired through, the reasoning→tool boundary holds with them, the malformed +header from #46314 no longer drops sibling calls, and the registered adapters +resolve. Seed-specific budget-reflect tags inside reasoning are also covered +here because the old dedicated parser tests exercised them. +""" + +import json + +import pytest + +from tests.parser.engine.conftest import make_mock_tokenizer +from tests.parser.engine.streaming_helpers import ( + collect_function_name, + collect_tool_arguments, + simulate_reasoning_streaming, + simulate_tool_streaming, +) +from vllm.parser.engine.registered_adapters import ( + SeedOssParserReasoningAdapter, + SeedOssParserToolAdapter, +) +from vllm.parser.seed_oss import SeedOssParser + +TOOL_CALL_START = "" +TOOL_CALL_END = "" +THINK_START = "" +THINK_END = "" + +_THINK_END_ID = 51 +_TOOL_CALL_ID = 60 + +_SEED_OSS_VOCAB = { + THINK_START: 50, + THINK_END: _THINK_END_ID, + TOOL_CALL_START: _TOOL_CALL_ID, + TOOL_CALL_END: 61, +} + + +@pytest.fixture +def mock_tokenizer(): + return make_mock_tokenizer(_SEED_OSS_VOCAB) + + +@pytest.fixture +def tool_parser(mock_tokenizer): + return SeedOssParser( + mock_tokenizer, chat_template_kwargs={"enable_thinking": False} + ) + + +@pytest.fixture +def parser(mock_tokenizer): + return SeedOssParser(mock_tokenizer) + + +def test_token_overrides_wired(parser): + assert parser.parser_engine_config.name == "seed_oss" + assert parser.reasoning_start_str == THINK_START + assert parser.reasoning_end_str == THINK_END + + +def test_single_tool_call(tool_parser, mock_request): + text = ( + f"{TOOL_CALL_START}\n\n" + "Tokyo\n" + f"\n{TOOL_CALL_END}" + ) + result = tool_parser.extract_tool_calls(text, mock_request) + + assert result.tools_called is True + assert result.tool_calls[0].function.name == "get_weather" + assert json.loads(result.tool_calls[0].function.arguments) == {"city": "Tokyo"} + + +def test_malformed_function_end_does_not_drop_siblings(tool_parser, mock_request): + """Regression for #46314: a malformed ```` with no closing ``>`` + on the header must not discard the other, well-formed calls.""" + text = ( + f"{TOOL_CALL_START}\n\n{TOOL_CALL_END}" + f"{TOOL_CALL_START}\n\n" + "Tokyo\n" + f"\n{TOOL_CALL_END}" + ) + result = tool_parser.extract_tool_calls(text, mock_request) + + weather = next(tc for tc in result.tool_calls if tc.function.name == "get_weather") + assert json.loads(weather.function.arguments) == {"city": "Tokyo"} + + +def test_basic_streaming(tool_parser, mock_request): + chunks = [ + f"{TOOL_CALL_START}\n", + "\n", + "Tokyo", + "\n", + "\n", + f"{TOOL_CALL_END}", + ] + results = simulate_tool_streaming(tool_parser, mock_request, chunks) + + assert collect_function_name(results) == "get_weather" + assert json.loads(collect_tool_arguments(results)) == {"city": "Tokyo"} + + +def test_reasoning_then_tool_call(parser): + text = ( + f"{THINK_START}I need to read the file.{THINK_END}" + f"{TOOL_CALL_START}\n\n" + "/tmp/x\n" + f"\n{TOOL_CALL_END}" + ) + reasoning, _ = parser.extract_reasoning(text, None) + assert reasoning == "I need to read the file." + assert TOOL_CALL_START not in reasoning + + +def test_streaming_think_end_and_tool_call_same_delta(parser): + """```` and ```` arriving in one delta must + not leak the terminal tokens into the reasoning text.""" + reasoning, content = simulate_reasoning_streaming( + parser, + [ + "Let me list the directory.", + f"{THINK_END}{TOOL_CALL_START}", + "", + ], + [(1,), (_THINK_END_ID, _TOOL_CALL_ID), (2,)], + ) + assert reasoning == "Let me list the directory." + assert THINK_END not in reasoning + assert TOOL_CALL_START not in reasoning + assert content is not None + + +def test_end_to_end_through_registered_adapters(mock_tokenizer, mock_request): + reasoning_parser = SeedOssParserReasoningAdapter(mock_tokenizer) + tool_parser = SeedOssParserToolAdapter(mock_tokenizer) + text = ( + f"{THINK_START}Plan the call.{THINK_END}" + f"{TOOL_CALL_START}\n\n" + "Tokyo\n" + f"\n{TOOL_CALL_END}" + ) + reasoning, remaining = reasoning_parser.extract_reasoning(text, mock_request) + assert reasoning == "Plan the call." + + tool_result = tool_parser.extract_tool_calls(remaining, mock_request) + assert tool_result.tool_calls[0].function.name == "get_weather" + assert json.loads(tool_result.tool_calls[0].function.arguments) == {"city": "Tokyo"} + + +def test_budget_reflect_tags_do_not_break_adapter_pipeline( + mock_tokenizer, + mock_request, +): + reasoning_parser = SeedOssParserReasoningAdapter(mock_tokenizer) + tool_parser = SeedOssParserToolAdapter(mock_tokenizer) + text = ( + f"{THINK_START}" + "The user's current thinking budget is 512.\n" + "I need the weather.\n" + "I have used 131 tokens." + "\n" + f"{THINK_END}" + f"{TOOL_CALL_START}\n\n" + "Barcelona\n" + f"\n{TOOL_CALL_END}" + ) + + reasoning, remaining = reasoning_parser.extract_reasoning(text, mock_request) + assert reasoning is not None + assert "current thinking budget is 512" in reasoning + assert "" in reasoning + assert "" in reasoning + + tool_result = tool_parser.extract_tool_calls(remaining, mock_request) + assert tool_result.tool_calls[0].function.name == "get_weather" + assert json.loads(tool_result.tool_calls[0].function.arguments) == { + "city": "Barcelona" + } diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index bee3d5d8b28..7f41b2b9513 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -34,6 +34,7 @@ from vllm.parser.engine.registered_adapters import ( MinimaxM2Parser, NemotronV3Parser, Qwen3Parser, + SeedOssParser, ) # ── Data structures ────────────────────────────────────────────────── @@ -587,6 +588,61 @@ def _build_nemotron_v3(scenario: Scenario, validate: bool = True) -> Sample: ) +# ── Seed-OSS (Qwen3 XML grammar with Seed wrapper tokens) ──────────── + +_SEED_OSS_VOCAB: dict[str, int] = { + "": 50, + "": 51, + "": 60, + "": 61, +} + + +def _seed_oss_tool_segments(tc: ToolCallSpec) -> list[tuple[str, bool]]: + parts = [f"\n"] + for key, value in tc.arguments.items(): + parts.append(f"\n{_qwen3_arg_value(value)}") + parts.append("\n\n") + return [ + ("", True), + ("".join(parts), False), + ("", True), + ] + + +def _seed_oss_segments(scenario: Scenario) -> list[tuple[str, bool]]: + segs: list[tuple[str, bool]] = [] + if scenario.reasoning is not None: + segs.append((scenario.reasoning, False)) + if scenario.content is not None or scenario.tool_calls is not None: + segs.append(("", True)) + if scenario.tool_calls is not None and not scenario.tool_calls: + segs.append(("", True)) + segs.append(("", True)) + if scenario.content is not None: + segs.append((scenario.content, False)) + if scenario.tool_calls: + for tc in scenario.tool_calls: + segs.extend(_seed_oss_tool_segments(tc)) + return segs + + +def _build_seed_oss(scenario: Scenario, validate: bool = True) -> Sample: + sample = _make_sample( + sample_id=f"seed_oss-{scenario.id}", + description=scenario.description, + vocab=_SEED_OSS_VOCAB, + segments=_seed_oss_segments(scenario), + expected_reasoning=scenario.reasoning if scenario.reasoning is not None else "", + expected_content=_qwen3_expected_content(scenario), + expected_tool_calls=_expected_tc(scenario), + tools=_expected_tools(scenario), + ) + if validate: + _validate_sample(sample, SeedOssParser) + return sample + + # ── GLM-4.7 MoE (XML tool format, starts in REASONING) ────────────── _GLM47_MOE_VOCAB: dict[str, int] = { @@ -668,6 +724,7 @@ _BUILDERS: dict[str, Any] = { "gemma4": _build_gemma4, "minimax_m2": _build_minimax_m2, "nemotron_v3": _build_nemotron_v3, + "seed_oss": _build_seed_oss, "glm47_moe": _build_glm47_moe, } diff --git a/tests/parser/test_parse.py b/tests/parser/test_parse.py index 39c5c2e3d5a..2a34ac7eea7 100644 --- a/tests/parser/test_parse.py +++ b/tests/parser/test_parse.py @@ -3,6 +3,7 @@ import json import os +from types import SimpleNamespace import pytest @@ -13,7 +14,9 @@ os.environ[_STRICT_TOOL_CALLING_ENV] = "0" from vllm.entrypoints.openai.chat_completion.protocol import ( # noqa: E402 ChatCompletionRequest, ) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest # noqa: E402 from vllm.parser.abstract_parser import DelegatingParser # noqa: E402 +from vllm.parser.utils import count_history_tool_calls # noqa: E402 from vllm.reasoning.basic_parsers import ( # noqa: E402 BaseThinkingReasoningParser, ) @@ -82,12 +85,55 @@ TOOLS = [ ] -def make_parser(tokenizer, reasoning=False, tool=False): +KIMI_K2_MODEL_CONFIG = SimpleNamespace( + hf_text_config=SimpleNamespace(model_type="kimi_k2"), + hf_overrides=None, +) + +HISTORY_MESSAGES = [ + {"role": "user", "content": "first"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "functions.get_current_weather:0", + "type": "function", + "function": { + "name": "get_current_weather", + "arguments": "{}", + }, + }, + { + "id": "functions.get_forecast:1", + "type": "function", + "function": { + "name": "get_forecast", + "arguments": "{}", + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "functions.get_current_weather:0", + "content": "{}", + }, + { + "role": "tool", + "tool_call_id": "functions.get_forecast:1", + "content": "{}", + }, + {"role": "user", "content": "again"}, +] + + +def make_parser(tokenizer, reasoning=False, tool=False, **kwargs): class TestParser(DelegatingParser): reasoning_parser_cls = ThinkReasoningParser if reasoning else None tool_parser_cls = Hermes2ProToolParser if tool else None - return TestParser(tokenizer) + return TestParser(tokenizer, **kwargs) @pytest.mark.parametrize( @@ -232,6 +278,107 @@ def test_parse_required_tool_choice(tokenizer): assert json.loads(tool_calls[1].arguments) == {"timezone": "UTC"} +def test_parse_required_tool_choice_kimi_k2_ids(tokenizer): + parser = make_parser( + tokenizer, reasoning=False, tool=True, model_config=KIMI_K2_MODEL_CONFIG + ) + functions_json = json.dumps( + [ + {"name": "get_current_weather", "parameters": {"city": "Dallas"}}, + {"name": "get_forecast", "parameters": {"city": "Dallas", "days": 2}}, + ] + ) + request = make_request(tools=TOOLS, tool_choice="required") + _, content, tool_calls = parser.parse( + functions_json, request, enable_auto_tools=True + ) + + assert content is None + assert tool_calls is not None + assert [tc.id for tc in tool_calls] == [ + "functions.get_current_weather:0", + "functions.get_forecast:1", + ] + + +def test_parse_required_tool_choice_kimi_k2_ids_after_history(tokenizer): + parser = make_parser( + tokenizer, reasoning=False, tool=True, model_config=KIMI_K2_MODEL_CONFIG + ) + functions_json = json.dumps( + [{"name": "get_current_weather", "parameters": {"city": "Dallas"}}] + ) + request = make_request( + messages=HISTORY_MESSAGES, + tools=TOOLS, + tool_choice="required", + ) + _, _, tool_calls = parser.parse(functions_json, request, enable_auto_tools=True) + + assert tool_calls is not None + assert tool_calls[0].id == "functions.get_current_weather:2" + + +def test_count_history_tool_calls_responses_request(): + request = ResponsesRequest.model_validate( + { + "model": "test-model", + "input": [ + { + "type": "function_call", + "call_id": "call_0", + "name": "get_current_weather", + "arguments": "{}", + }, + { + "type": "function_call", + "call_id": "call_1", + "name": "get_forecast", + "arguments": "{}", + }, + ], + } + ) + + assert count_history_tool_calls(request) == 2 + + +def test_parse_required_tool_choice_random_ids_deferred(tokenizer): + parser = make_parser(tokenizer, reasoning=False, tool=True) + functions_json = json.dumps( + [{"name": "get_current_weather", "parameters": {"city": "Dallas"}}] + ) + request = make_request( + messages=HISTORY_MESSAGES, + tools=TOOLS, + tool_choice="required", + ) + _, _, tool_calls = parser.parse(functions_json, request, enable_auto_tools=True) + + assert tool_calls is not None + assert tool_calls[0].id is None + + +def test_parse_named_tool_choice_kimi_k2_id(tokenizer): + parser = make_parser( + tokenizer, reasoning=False, tool=True, model_config=KIMI_K2_MODEL_CONFIG + ) + request = make_request( + tools=TOOLS, + tool_choice={ + "type": "function", + "function": {"name": "get_weather"}, + }, + ) + _, content, tool_calls = parser.parse( + TOOL_ARGUMENTS, request, enable_auto_tools=True + ) + + assert content is None + assert tool_calls is not None + assert tool_calls[0].id == "functions.get_weather:0" + + def test_parse_named_tool_choice_content_none(tokenizer): parser = make_parser(tokenizer, reasoning=False, tool=True) request = make_request( diff --git a/tests/parser/test_streaming.py b/tests/parser/test_streaming.py index dbc64e75593..e6ec273b25e 100644 --- a/tests/parser/test_streaming.py +++ b/tests/parser/test_streaming.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +from types import SimpleNamespace import pytest @@ -47,6 +48,36 @@ TOOLS = [ ] +KIMI_K2_MODEL_CONFIG = SimpleNamespace( + hf_text_config=SimpleNamespace(model_type="kimi_k2"), + hf_overrides=None, +) + +HISTORY_MESSAGES = [ + {"role": "user", "content": "first"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "functions.get_current_weather:0", + "type": "function", + "function": { + "name": "get_current_weather", + "arguments": "{}", + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "functions.get_current_weather:0", + "content": "{}", + }, + {"role": "user", "content": "again"}, +] + + @pytest.fixture def request_obj(): return ChatCompletionRequest( @@ -57,12 +88,12 @@ def request_obj(): ) -def make_parser(tokenizer, reasoning=False, tool=False): +def make_parser(tokenizer, reasoning=False, tool=False, **kwargs): class TestParser(DelegatingParser): reasoning_parser_cls = ThinkReasoningParser if reasoning else None tool_parser_cls = Hermes2ProToolParser if tool else None - return TestParser(tokenizer) + return TestParser(tokenizer, **kwargs) def stream_text(parser, tokenizer, text, request, prompt_token_ids=None): @@ -365,3 +396,76 @@ def test_parse_delta_tool_choice_none_with_reasoning(tokenizer, request_obj): assert len(tool_calls) == 0 assert "" in content assert "get_weather" in content + + +def test_parse_delta_required_tool_choice_kimi_k2_ids(tokenizer, request_obj): + parser = make_parser( + tokenizer, reasoning=False, tool=True, model_config=KIMI_K2_MODEL_CONFIG + ) + request = request_obj.model_copy(update={"tool_choice": "required"}) + output = json.dumps( + [ + { + "name": "get_current_weather", + "parameters": {"city": "Dallas"}, + } + ] + ) + + results: list[DeltaMessage | None] = [] + prompt_token_ids: list[int] | None = [] + for i in range(0, len(output), 3): + chunk = output[i : i + 3] + results.append( + parser.parse_delta( + chunk, + [], + request, + prompt_token_ids=prompt_token_ids, + finished=False, + ) + ) + prompt_token_ids = None + + _, content, tool_calls = collect_fields(results) + assert content == "" + assert any(tc.id == "functions.get_current_weather:0" for tc in tool_calls) + assert all(tc.id in (None, "functions.get_current_weather:0") for tc in tool_calls) + + +def test_parse_delta_required_tool_choice_kimi_k2_ids_after_history( + tokenizer, request_obj +): + parser = make_parser( + tokenizer, reasoning=False, tool=True, model_config=KIMI_K2_MODEL_CONFIG + ) + request = request_obj.model_copy( + update={"messages": HISTORY_MESSAGES, "tool_choice": "required"} + ) + output = json.dumps( + [ + { + "name": "get_current_weather", + "parameters": {"city": "Dallas"}, + } + ] + ) + + results: list[DeltaMessage | None] = [] + prompt_token_ids: list[int] | None = [] + for i in range(0, len(output), 3): + chunk = output[i : i + 3] + results.append( + parser.parse_delta( + chunk, + [], + request, + prompt_token_ids=prompt_token_ids, + finished=False, + ) + ) + prompt_token_ids = None + + _, _, tool_calls = collect_fields(results) + assert any(tc.id == "functions.get_current_weather:1" for tc in tool_calls) + assert all(tc.id in (None, "functions.get_current_weather:1") for tc in tool_calls) diff --git a/tests/reasoning/test_seedoss_reasoning_parser.py b/tests/reasoning/test_seedoss_reasoning_parser.py deleted file mode 100644 index 33d56d32965..00000000000 --- a/tests/reasoning/test_seedoss_reasoning_parser.py +++ /dev/null @@ -1,236 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from typing import Any, cast - -import pytest -from transformers import AutoTokenizer - -from tests.reasoning.utils import run_reasoning_extraction -from vllm.reasoning import ReasoningParser, ReasoningParserManager - -parser_name = "seed_oss" -start_token = "" -end_token = "" - -# Use a test model that contains our custom tokens -REASONING_MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" - - -@pytest.fixture(scope="module") -def seedoss_tokenizer(): - tokenizer = AutoTokenizer.from_pretrained(REASONING_MODEL_NAME) - # Add custom SeedOSS tokens if they don't exist - if start_token not in tokenizer.get_vocab(): - tokenizer.add_tokens([start_token, end_token]) - return tokenizer - - -SIMPLE_REASONING: dict[str, Any] = { - "output": "This is a reasoning sectionThis is the rest", - "reasoning": "This is a reasoning section", - "content": "This is the rest", - "is_reasoning_end": True, -} -COMPLETE_REASONING: dict[str, Any] = { - "output": "This is a reasoning section", - "reasoning": "This is a reasoning section", - "content": None, - "is_reasoning_end": True, -} -NO_CONTENT: dict[str, Any] = { - "output": "This is content", - "reasoning": "This is content", - "content": None, - "is_reasoning_end": False, -} -NO_REASONING_STREAMING: dict[str, Any] = { - "output": "This is a reasoning section", - "reasoning": "This is a reasoning section", - "content": None, - "is_reasoning_end": False, -} -MULTIPLE_LINES: dict[str, Any] = { - "output": "This\nThatThis is the rest\nThat", - "reasoning": "This\nThat", - "content": "This is the rest\nThat", - "is_reasoning_end": True, -} -WITH_START_TOKEN: dict[str, Any] = { - "output": ("This is a reasoning sectionThis is the rest"), - "reasoning": "This is a reasoning section", - "content": "This is the rest", - "is_reasoning_end": True, -} -ONLY_END_TOKEN: dict[str, Any] = { - "output": "Some reasoningThis is the rest", - "reasoning": "Some reasoning", - "content": "This is the rest", - "is_reasoning_end": True, -} -NO_TOKENS: dict[str, Any] = { - "output": "This is just content without any reasoning tokens", - "reasoning": "This is just content without any reasoning tokens", - "content": None, - "is_reasoning_end": False, -} - - -def test_seedoss_reasoning_parser_creation(seedoss_tokenizer): - """Test that the SeedOSS reasoning parser can be created and registered.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - assert isinstance(parser, ReasoningParser) - assert parser.start_token == start_token - assert parser.end_token == end_token - - -@pytest.mark.parametrize("streaming", [True, False]) -def test_simple_reasoning(seedoss_tokenizer, streaming): - """Test basic reasoning extraction with both tokens.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - reasoning, content = run_reasoning_extraction( - parser, [cast(str, SIMPLE_REASONING["output"])], streaming=streaming - ) - - assert reasoning == SIMPLE_REASONING["reasoning"] - assert content == SIMPLE_REASONING["content"] - - -@pytest.mark.parametrize("streaming", [True, False]) -def test_complete_reasoning(seedoss_tokenizer, streaming): - """Test reasoning extraction when there's no content after reasoning.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - reasoning, content = run_reasoning_extraction( - parser, [cast(str, COMPLETE_REASONING["output"])], streaming=streaming - ) - - assert reasoning == COMPLETE_REASONING["reasoning"] - assert content == COMPLETE_REASONING["content"] - - -@pytest.mark.parametrize("streaming", [True, False]) -def test_no_content(seedoss_tokenizer, streaming): - """Test when there's no end token - everything is reasoning content.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - reasoning, content = run_reasoning_extraction( - parser, [cast(str, NO_CONTENT["output"])], streaming=streaming - ) - - assert reasoning == NO_CONTENT["reasoning"] - assert content == NO_CONTENT["content"] - - -@pytest.mark.parametrize("streaming", [True, False]) -def test_multiple_lines(seedoss_tokenizer, streaming): - """Test reasoning extraction with multiline content.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - reasoning, content = run_reasoning_extraction( - parser, [cast(str, MULTIPLE_LINES["output"])], streaming=streaming - ) - - assert reasoning == MULTIPLE_LINES["reasoning"] - assert content == MULTIPLE_LINES["content"] - - -@pytest.mark.parametrize("streaming", [True, False]) -def test_with_start_token(seedoss_tokenizer, streaming): - """Test reasoning extraction with both start and end tokens.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - reasoning, content = run_reasoning_extraction( - parser, [cast(str, WITH_START_TOKEN["output"])], streaming=streaming - ) - - assert reasoning == WITH_START_TOKEN["reasoning"] - assert content == WITH_START_TOKEN["content"] - - -@pytest.mark.parametrize("streaming", [True, False]) -def test_only_end_token(seedoss_tokenizer, streaming): - """ - Test reasoning extraction with only end token - (SeedOSS typical behavior). - """ - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - reasoning, content = run_reasoning_extraction( - parser, [cast(str, ONLY_END_TOKEN["output"])], streaming=streaming - ) - - assert reasoning == ONLY_END_TOKEN["reasoning"] - assert content == ONLY_END_TOKEN["content"] - - -@pytest.mark.parametrize("streaming", [True, False]) -def test_no_tokens(seedoss_tokenizer, streaming): - """Test when there are no reasoning tokens at all.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - reasoning, content = run_reasoning_extraction( - parser, [cast(str, NO_TOKENS["output"])], streaming=streaming - ) - - assert reasoning == NO_TOKENS["reasoning"] - assert content == NO_TOKENS["content"] - - -def test_is_reasoning_end(seedoss_tokenizer): - """Test the is_reasoning_end method.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - # Test with end token present - end_token_id = parser.end_token_id - assert parser.is_reasoning_end([1, 2, end_token_id, 4]) is True - - # Test without end token - assert parser.is_reasoning_end([1, 2, 3, 4]) is False - - -def test_extract_content_ids(seedoss_tokenizer): - """Test the extract_content_ids method.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - end_token_id = parser.end_token_id - - # Test with end token in the middle - input_ids = [1, 2, end_token_id, 4, 5] - content_ids = parser.extract_content_ids(input_ids) - assert content_ids == [4, 5] - - # Test with end token at the end - input_ids = [1, 2, 3, end_token_id] - content_ids = parser.extract_content_ids(input_ids) - assert content_ids == [] - - # Test without end token - input_ids = [1, 2, 3, 4] - content_ids = parser.extract_content_ids(input_ids) - assert content_ids == [] - - -def test_streaming_delta_processing(seedoss_tokenizer): - """Test streaming processing with small deltas.""" - parser_cls = ReasoningParserManager.get_reasoning_parser(parser_name) - parser = parser_cls(seedoss_tokenizer) - - # Test streaming with incremental tokens - deltas = ["Some ", "reasoning ", "content", "", "Final ", "answer"] - - reasoning, content = run_reasoning_extraction(parser, deltas, streaming=True) - - assert reasoning == "Some reasoning content" - assert content == "Final answer" diff --git a/tests/rocm/aiter/test_quant_op_schema.py b/tests/rocm/aiter/test_quant_op_schema.py index 9b2fac6e017..d5668f95a6d 100644 --- a/tests/rocm/aiter/test_quant_op_schema.py +++ b/tests/rocm/aiter/test_quant_op_schema.py @@ -2,9 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Schema/aliasing tests for the AITER FP8 quantization custom ops. # -# These use torch.library.opcheck, whose test_schema check catches custom ops -# whose implementation aliases an input that the registered schema declares as -# non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant +# These use the shared opcheck helper, whose test_schema check catches custom +# ops whose implementation aliases an input that the registered schema declares +# as non-aliasing -- the failure mode behind the rocm_aiter_per_tensor_quant # regression (a returned scale that aliased the input scale). # # Skipped if AITER is not installed or the platform is not ROCm. @@ -14,6 +14,8 @@ import importlib.util import pytest import torch +from tests.kernels.utils import opcheck + # this import statement is needed to ensure the ops are registered from vllm._aiter_ops import rocm_aiter_ops from vllm.platforms import current_platform @@ -32,26 +34,14 @@ def _x(M=128, N=4096): return torch.randn((M, N), dtype=torch.float16, device="cuda") -# The in-place per-tensor op takes the fp8 output buffer as an input, which -# opcheck's test_schema cannot exercise ("mul_cuda" is unimplemented for fp8), -# so restrict to the utils that run on fp8 inputs. The aliasing contract for -# this op is instead covered by test_per_tensor_quant_torch_compile below. -_INPLACE_OPCHECK_UTILS = ( - "test_faketensor", - "test_aot_dispatch_dynamic", - "test_autograd_registration", -) - - def test_per_tensor_quant_static_schema(): """Static per-tensor: caller provides scale (the aliasing regression).""" x = _x() out = torch.empty_like(x, dtype=FP8_DTYPE) scale = torch.ones(1, dtype=torch.float32, device="cuda") - torch.library.opcheck( + opcheck( torch.ops.vllm.rocm_aiter_per_tensor_quant, (out, x, scale, False), - test_utils=_INPLACE_OPCHECK_UTILS, ) @@ -60,17 +50,16 @@ def test_per_tensor_quant_dynamic_schema(): x = _x() out = torch.empty_like(x, dtype=FP8_DTYPE) scale = torch.empty(1, dtype=torch.float32, device="cuda") - torch.library.opcheck( + opcheck( torch.ops.vllm.rocm_aiter_per_tensor_quant, (out, x, scale, True), - test_utils=_INPLACE_OPCHECK_UTILS, ) def test_per_token_quant_dynamic_schema(): """Dynamic per-token: op computes scale into a freshly allocated buffer.""" x = _x() - torch.library.opcheck( + opcheck( torch.ops.vllm.rocm_aiter_per_token_quant, (x, FP8_DTYPE, None), ) @@ -79,7 +68,7 @@ def test_per_token_quant_dynamic_schema(): def test_group_fp8_quant_schema(): """Dynamic per-token-group quant.""" x = _x() - torch.library.opcheck( + opcheck( torch.ops.vllm.rocm_aiter_group_fp8_quant, (x, 128), ) @@ -103,43 +92,24 @@ def test_per_tensor_quant_matches_native(dynamic): assert out.shape == x.shape assert out.dtype == FP8_DTYPE assert scale.shape == ref_scale.shape - if not dynamic: - # static scale is passed through unchanged - assert torch.equal(scale, scale_in) - # Compare dequantized values to be robust to 1-ULP fp8 boundary flips. deq = out.to(torch.float32) * scale - ref_deq = ref_out.to(torch.float32) * ref_scale - torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2) + if dynamic: + # Dynamic mode: AITER and native each compute their own scale, so their + # outputs differ and can't be compared. Just check that AITER's output + # dequantizes back to the input, within fp8 rounding error. + torch.testing.assert_close(deq, x.to(torch.float32), rtol=0.07, atol=5e-2) + else: + # Static mode: both use the caller's scale, so the outputs must match. + assert torch.equal(scale, scale_in) + ref_deq = ref_out.to(torch.float32) * ref_scale + torch.testing.assert_close(deq, ref_deq, rtol=2e-2, atol=2e-2) -@pytest.mark.parametrize("dynamic", [True, False]) -def test_per_tensor_quant_torch_compile(monkeypatch, dynamic): - """per_tensor_quant compiles under inductor without an aliasing error. - - Forces the custom-op aliasing check to error (it is otherwise only a - warning outside CI), so a regression that returns an input-aliasing - scale fails here regardless of the CI env var. - """ - aliasing_cfg = pytest.importorskip("torch._functorch.config") - monkeypatch.setattr( - aliasing_cfg, "error_on_custom_op_aliasing", True, raising=False - ) - - x = _x() - scale = None if dynamic else torch.tensor([0.5], dtype=torch.float32, device="cuda") - - def fn(x, s): - return rocm_aiter_ops.per_tensor_quant(x, FP8_DTYPE, s) - - compiled = torch.compile(fn, fullgraph=True, backend="inductor", dynamic=False) - - out_eager, scale_eager = fn(x, scale) - out_compiled, scale_compiled = compiled(x, scale) - - assert out_compiled.shape == out_eager.shape - torch.testing.assert_close( - out_compiled.to(torch.float32) * scale_compiled, - out_eager.to(torch.float32) * scale_eager, - rtol=2e-2, - atol=2e-2, - ) +# A test_per_tensor_quant_torch_compile test previously lived here to validate +# the per-tensor aliasing contract. It existed because opcheck's test_schema +# could not check this op directly: test_schema compares the op's outputs with +# torch.allclose, but on fp8 outputs that comparison runs arithmetic fp8 does not +# support and raises "mul_cuda" is unimplemented for fp8. The fp8-safe opcheck +# helper fixes that by casting to double before the comparison, so the per-tensor +# schema tests above can now run test_schema directly. That makes this test +# redundant, so it has been removed. diff --git a/tests/tool_parsers/test_rust_tool_parser.py b/tests/tool_parsers/test_rust_tool_parser.py index 75468487783..2349d4d292a 100644 --- a/tests/tool_parsers/test_rust_tool_parser.py +++ b/tests/tool_parsers/test_rust_tool_parser.py @@ -171,7 +171,7 @@ def test_rust_tool_parser_extension_typed_api() -> None: parser.parse_into(build_tool_call(), output) output.append(parser.finish()) - output = output.coalesce_calls() + output = output.coalesce() assert parser.preserve_special_tokens() assert output.normal_text == "" diff --git a/tests/tool_parsers/test_seed_oss_tool_parser.py b/tests/tool_parsers/test_seed_oss_tool_parser.py deleted file mode 100644 index 4ff96fb01be..00000000000 --- a/tests/tool_parsers/test_seed_oss_tool_parser.py +++ /dev/null @@ -1,522 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# ruff: noqa: E501 - -import json -from collections.abc import Generator - -import pytest - -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ChatCompletionToolsParam, -) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaMessage, - FunctionCall, - ToolCall, -) -from vllm.tokenizers import TokenizerLike, get_tokenizer -from vllm.tokenizers.detokenizer_utils import detokenize_incrementally -from vllm.tool_parsers.seed_oss_tool_parser import SeedOssToolParser - -# Use a common model that is likely to be available -MODEL = "ByteDance-Seed/Seed-OSS-36B-Instruct" - - -@pytest.fixture(scope="module") -def seed_oss_tokenizer(): - return get_tokenizer(tokenizer_name=MODEL, trust_remote_code=True) - - -@pytest.fixture -def seed_oss_tool_parser(seed_oss_tokenizer, sample_tools): - return SeedOssToolParser(seed_oss_tokenizer, tools=sample_tools) - - -@pytest.fixture -def sample_tools(): - return [ - ChatCompletionToolsParam( - type="function", - function={ - "name": "get_weather", - "description": "Get current temperature for a given location.", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "City and country e.g. Bogotá, Colombia", - }, - "unit": { - "type": "string", - "description": "this is the unit of temperature", - }, - }, - "required": ["location"], - "additionalProperties": False, - }, - "returns": { - "type": "object", - "properties": { - "temperature": { - "type": "number", - "description": "temperature in celsius", - } - }, - "required": ["temperature"], - "additionalProperties": False, - }, - "strict": True, - }, - ), - ] - - -def assert_tool_calls( - actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall] -): - assert len(actual_tool_calls) == len(expected_tool_calls) - - for actual_tool_call, expected_tool_call in zip( - actual_tool_calls, expected_tool_calls - ): - # Seed-OSS tool call will not generate id - assert actual_tool_call.type == "function" - assert actual_tool_call.function == expected_tool_call.function - - assert actual_tool_call.function.name == expected_tool_call.function.name - assert ( - actual_tool_call.function.arguments == expected_tool_call.function.arguments - ) - - -def test_extract_tool_calls_no_tools(seed_oss_tool_parser): - model_output = "This is a test response without any tool calls" - extracted_tool_calls = seed_oss_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - - assert not extracted_tool_calls.tools_called - assert extracted_tool_calls.tool_calls == [] - assert extracted_tool_calls.content == model_output - - -@pytest.mark.parametrize( - ids=[ - "tool_call_0_thinking_budget", - "tool_call_512_thinking_budget", - "tool_call_unlimited_thinking_budget", - ], - argnames=["model_output", "expected_tool_calls", "expected_content"], - argvalues=[ - ( - """\n\n""" - """Barcelona, Spain\n\n""", - [ - ToolCall( - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "location": "Barcelona, Spain", - }, - ), - ), - type="function", - ) - ], - None, - ), - ( - """The user\'s current thinking budget is 512.\nLet me analyze the """ - """question. The user wants to know the weather in Barcelona, Spain. Looking at the functions available, """ - """there\'s a get_weather function that can retrieve the current temperature for a given location. \n\nFirst, """ - """check the parameters required by get_weather: location is mandatory (needs city and country), and unit is """ - """optional. The user provided "Barcelona Spain" as the location, which fits the required format (city, """ - """country). \nI have used 131 tokens, and there are 381 tokens remaining for use.""" - """\n Since the unit isn\'t specified, the function will default to Celsius, which """ - """is fine. \n\nThere\'s no need to ask for more information because the location is clear. So I should call """ - """the get_weather function with location set to "Barcelona, Spain" (adding a comma for clarity, though the """ - """user\'s input has a space, but the function might accept either; to be safe, using the standard format """ - """with a comma).\nI have used 257 tokens, and there are 255 tokens remaining for """ - """use.\n The unit parameter can be omitted since it\'s optional.\n""" - """\n\nBarcelona, Spain\n""" - """\n""", - [ - ToolCall( - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "location": "Barcelona, Spain", - }, - ), - ), - type="function", - ) - ], - """The user\'s current thinking budget is 512.\nLet me analyze the """ - """question. The user wants to know the weather in Barcelona, Spain. Looking at the functions available, """ - """there\'s a get_weather function that can retrieve the current temperature for a given location. \n\nFirst, """ - """check the parameters required by get_weather: location is mandatory (needs city and country), and unit is """ - """optional. The user provided "Barcelona Spain" as the location, which fits the required format (city, """ - """country). \nI have used 131 tokens, and there are 381 tokens remaining for use.""" - """\n Since the unit isn\'t specified, the function will default to Celsius, which """ - """is fine. \n\nThere\'s no need to ask for more information because the location is clear. So I should call """ - """the get_weather function with location set to "Barcelona, Spain" (adding a comma for clarity, though the """ - """user\'s input has a space, but the function might accept either; to be safe, using the standard format """ - """with a comma).\nI have used 257 tokens, and there are 255 tokens remaining for """ - """use.\n The unit parameter can be omitted since it\'s optional.\n""", - ), - ( - """\nGot it, let\'s see. The user asked for the weather in Barcelona, Spain. """ - """First, I need to remember the function I can use: get_weather. The function requires a """ - """location (city and country) which is "Barcelona, Spain" here, and unit is optional. Since """ - """the user didn\'t specify the unit, the default in the function is Celsius, right? Wait, """ - """let me check the function docstring again. Oh, the function says unit is optional, and """ - """returns temperature in Celsius. So I should call get_weather with location "Barcelona, """ - """Spain" and maybe omit unit or set to Celsius. Let me format the function call correctly. """ - """The format is \n\nBarcelona, """ - """Spain\ncelsius\n\n. """ - """Wait, but does the unit parameter accept "celsius"? The docstring says unit is the unit """ - """of temperature, but the return is in Celsius anyway. Maybe even if I don\'t pass unit, """ - """it\'s okay, but to be explicit, maybe pass "celsius". Let me go with that. So the function """ - """call should be as above. Then wait for the result to come back and tell the user the """ - """temperature in Celsius.\n\n""" - """Barcelona, Spain\ncelsius\n\n""", - [ - ToolCall( - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "location": "Barcelona, Spain", - "unit": "celsius", - }, - ), - ), - type="function", - ) - ], - """\nGot it, let\'s see. The user asked for the weather in Barcelona, Spain. """ - """First, I need to remember the function I can use: get_weather. The function requires a """ - """location (city and country) which is "Barcelona, Spain" here, and unit is optional. Since """ - """the user didn\'t specify the unit, the default in the function is Celsius, right? Wait, """ - """let me check the function docstring again. Oh, the function says unit is optional, and """ - """returns temperature in Celsius. So I should call get_weather with location "Barcelona, """ - """Spain" and maybe omit unit or set to Celsius. Let me format the function call correctly. """ - """The format is \n\nBarcelona, """ - """Spain\ncelsius\n\n. """ - """Wait, but does the unit parameter accept "celsius"? The docstring says unit is the unit """ - """of temperature, but the return is in Celsius anyway. Maybe even if I don\'t pass unit, """ - """it\'s okay, but to be explicit, maybe pass "celsius". Let me go with that. So the function """ - """call should be as above. Then wait for the result to come back and tell the user the """ - """temperature in Celsius.""", - ), - ], -) -def test_extract_tool_calls( - seed_oss_tool_parser, - sample_tools, - model_output, - expected_tool_calls, - expected_content, -): - request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) - extracted_tool_calls = seed_oss_tool_parser.extract_tool_calls( - model_output, request=request - ) # type: ignore[arg-type] - assert extracted_tool_calls.tools_called - - assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) - - assert extracted_tool_calls.content == expected_content - - -def test_streaming_tool_calls_no_tools(seed_oss_tool_parser): - model_output = "This is a test response without any tool calls" - - result = seed_oss_tool_parser.extract_tool_calls_streaming( - previous_text="his is a test response", - current_text=model_output, - delta_text=" without any tool calls.", - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - # Should return the delta text as content - assert result is not None - assert hasattr(result, "content") - assert result.content == " without any tool calls." - - -def stream_delta_message_generator( - seed_oss_tool_parser: SeedOssToolParser, - seed_oss_tokenizer: TokenizerLike, - model_output: str, - request: ChatCompletionRequest | None = None, -) -> Generator[DeltaMessage, None, None]: - all_token_ids = seed_oss_tokenizer.encode(model_output, add_special_tokens=False) - - previous_text = "" - previous_tokens = None - prefix_offset = 0 - read_offset = 0 - for i, delta_token in enumerate(all_token_ids): - delta_token_ids = [delta_token] - previous_token_ids = all_token_ids[:i] - current_token_ids = all_token_ids[: i + 1] - - (new_tokens, delta_text, new_prefix_offset, new_read_offset) = ( - detokenize_incrementally( - tokenizer=seed_oss_tokenizer, - all_input_ids=current_token_ids, - prev_tokens=previous_tokens, - prefix_offset=prefix_offset, - read_offset=read_offset, - skip_special_tokens=False, - spaces_between_special_tokens=True, - ) - ) - - current_text = previous_text + delta_text - - delta_message = seed_oss_tool_parser.extract_tool_calls_streaming( - previous_text, - current_text, - delta_text, - previous_token_ids, - current_token_ids, - delta_token_ids, - request=request, - ) - if delta_message: - yield delta_message - - previous_text = current_text - previous_tokens = ( - previous_tokens + new_tokens if previous_tokens else new_tokens - ) - prefix_offset = new_prefix_offset - read_offset = new_read_offset - - -@pytest.mark.parametrize( - ids=[ - "tool_call_0_thinking_budget", - "tool_call_512_thinking_budget", - "tool_call_unlimited_thinking_budget", - ], - argnames=["model_output", "expected_tool_calls", "expected_content"], - argvalues=[ - ( - """\n\n\n""" - """The current thinking budget is 0, so I will directly start answering the question.\n\n""" - """\n\n""" - """Barcelona, Spain\n\n""", - [ - ToolCall( - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "location": "Barcelona, Spain", - }, - ), - ), - type="function", - ) - ], - """\n\n\n""" - """The current thinking budget is 0, so I will directly start answering the question.\n\n""", - ), - ( - """The user\'s current thinking budget is 512.\nLet me analyze the """ - """question. The user wants to know the weather in Barcelona, Spain. Looking at the functions available, """ - """there\'s a get_weather function that can retrieve the current temperature for a given location. \n\nFirst, """ - """check the parameters required by get_weather: location is mandatory (needs city and country), and unit is """ - """optional. The user provided "Barcelona Spain" as the location, which fits the required format (city, """ - """country). \nI have used 131 tokens, and there are 381 tokens remaining for use.""" - """\n Since the unit isn\'t specified, the function will default to Celsius, which """ - """is fine. \n\nThere\'s no need to ask for more information because the location is clear. So I should call """ - """the get_weather function with location set to "Barcelona, Spain" (adding a comma for clarity, though the """ - """user\'s input has a space, but the function might accept either; to be safe, using the standard format """ - """with a comma).\nI have used 257 tokens, and there are 255 tokens remaining for """ - """use.\n The unit parameter can be omitted since it\'s optional.\n""" - """\n\nBarcelona, Spain\n""" - """\n""", - [ - ToolCall( - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "location": "Barcelona, Spain", - }, - ), - ), - type="function", - ) - ], - """The user\'s current thinking budget is 512.\nLet me analyze the """ - """question. The user wants to know the weather in Barcelona, Spain. Looking at the functions available, """ - """there\'s a get_weather function that can retrieve the current temperature for a given location. \n\nFirst, """ - """check the parameters required by get_weather: location is mandatory (needs city and country), and unit is """ - """optional. The user provided "Barcelona Spain" as the location, which fits the required format (city, """ - """country). \nI have used 131 tokens, and there are 381 tokens remaining for use.""" - """\n Since the unit isn\'t specified, the function will default to Celsius, which """ - """is fine. \n\nThere\'s no need to ask for more information because the location is clear. So I should call """ - """the get_weather function with location set to "Barcelona, Spain" (adding a comma for clarity, though the """ - """user\'s input has a space, but the function might accept either; to be safe, using the standard format """ - """with a comma).\nI have used 257 tokens, and there are 255 tokens remaining for """ - """use.\n The unit parameter can be omitted since it\'s optional.\n""", - ), - ( - """\nGot it, let\'s see. The user asked for the weather in Barcelona, Spain. """ - """First, I need to remember the function I can use: get_weather. The function requires a """ - """location (city and country) which is "Barcelona, Spain" here, and unit is optional. Since """ - """the user didn\'t specify the unit, the default in the function is Celsius, right? Wait, """ - """let me check the function docstring again. Oh, the function says unit is optional, and """ - """returns temperature in Celsius. So I should call get_weather with location "Barcelona, """ - """Spain" and maybe omit unit or set to Celsius. Let me format the function call correctly. """ - """The format is \n\nBarcelona, """ - """Spain\ncelsius\n\n. """ - """Wait, but does the unit parameter accept "celsius"? The docstring says unit is the unit """ - """of temperature, but the return is in Celsius anyway. Maybe even if I don\'t pass unit, """ - """it\'s okay, but to be explicit, maybe pass "celsius". Let me go with that. So the function """ - """call should be as above. Then wait for the result to come back and tell the user the """ - """temperature in Celsius.\n\n""" - """Barcelona, Spain\ncelsius\n\n""", - [ - ToolCall( - function=FunctionCall( - name="get_weather", - arguments=json.dumps( - { - "location": "Barcelona, Spain", - "unit": "celsius", - }, - ), - ), - type="function", - ) - ], - """\nGot it, let\'s see. The user asked for the weather in Barcelona, Spain. """ - """First, I need to remember the function I can use: get_weather. The function requires a """ - """location (city and country) which is "Barcelona, Spain" here, and unit is optional. Since """ - """the user didn\'t specify the unit, the default in the function is Celsius, right? Wait, """ - """let me check the function docstring again. Oh, the function says unit is optional, and """ - """returns temperature in Celsius. So I should call get_weather with location "Barcelona, """ - """Spain" and maybe omit unit or set to Celsius. Let me format the function call correctly. """ - """The format is \n\nBarcelona, """ - """Spain\ncelsius\n\n. """ - """Wait, but does the unit parameter accept "celsius"? The docstring says unit is the unit """ - """of temperature, but the return is in Celsius anyway. Maybe even if I don\'t pass unit, """ - """it\'s okay, but to be explicit, maybe pass "celsius". Let me go with that. So the function """ - """call should be as above. Then wait for the result to come back and tell the user the """ - """temperature in Celsius.""", - ), - ], -) -def test_streaming_tool_calls( - seed_oss_tool_parser, - seed_oss_tokenizer, - sample_tools, - model_output, - expected_tool_calls, - expected_content, -): - """Test incremental streaming behavior""" - request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) - - other_content = "" - tool_states = {} # Track state per tool index - - for delta_message in stream_delta_message_generator( - seed_oss_tool_parser, seed_oss_tokenizer, model_output, request - ): - # role should never be streamed from tool parser - assert not delta_message.role - - if delta_message.content: - other_content += delta_message.content - - if delta_message.tool_calls: - for tool_call in delta_message.tool_calls: - idx = tool_call.index - - # Initialize state for new tool - if idx not in tool_states: - tool_states[idx] = { - "id": None, - "name": None, - "arguments": "", - "type": None, - } - - # First chunk should have id, name, and type - if tool_call.id: - tool_states[idx]["id"] = tool_call.id - - if tool_call.type: - assert tool_call.type == "function" - tool_states[idx]["type"] = tool_call.type - - if tool_call.function: - if tool_call.function.name: - # Should only be set once - assert tool_states[idx]["name"] is None - tool_states[idx]["name"] = tool_call.function.name - - if tool_call.function.arguments is not None: - # Accumulate arguments incrementally - tool_states[idx]["arguments"] += tool_call.function.arguments - - # Verify final content - assert other_content == expected_content - - # Verify we got all expected tool calls - assert len(tool_states) == len(expected_tool_calls) - - # Verify each tool call - for idx, expected_tool in enumerate(expected_tool_calls): - state = tool_states[idx] - assert state["id"] is not None - assert state["type"] == "function" - assert state["name"] == expected_tool.function.name - - # Parse accumulated arguments - arguments_str = state["arguments"] - assert arguments_str is not None - actual_args = json.loads(arguments_str) - expected_args = json.loads(expected_tool.function.arguments) - assert actual_args == expected_args - - -def test_streaming_tool_calls_non_ascii( - seed_oss_tool_parser, seed_oss_tokenizer, sample_tools -): - request = ChatCompletionRequest(model=MODEL, messages=[], tools=sample_tools) - model_output = ( - """\n\n\n""" - """The current thinking budget is 0, so I will directly start answering the question.\n\n""" - """\n\n""" - """北京\n\n""" - ) - - args = "".join( - tool_call.function.arguments - for delta_message in stream_delta_message_generator( - seed_oss_tool_parser, seed_oss_tokenizer, model_output, request - ) - if delta_message.tool_calls - for tool_call in delta_message.tool_calls - if tool_call.function and tool_call.function.arguments is not None - ) - - assert "北京" in args - assert "\\u" not in args diff --git a/tests/utils.py b/tests/utils.py index db5905b9275..2acb9716302 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1464,7 +1464,7 @@ def error_on_warning(category: type[Warning] = Warning): yield -def get_physical_device_indices(devices): +def get_physical_device_indices(devices: list[int]): visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") if visible_devices is None: return devices @@ -1475,82 +1475,114 @@ def get_physical_device_indices(devices): @_nvml() +def record_gpu_memory_usage_stats( + *, + devices: list[int], +) -> dict[int, tuple[float, float]]: + output: dict[int, tuple[float, float]] = {} + for device in devices: + if current_platform.is_rocm(): + dev_handle = amdsmi_get_processor_handles()[device] + mem_info = amdsmi_get_gpu_vram_usage(dev_handle) + gb_used = mem_info["vram_used"] / 2**10 + gb_total = mem_info["vram_total"] / 2**10 + else: + dev_handle = nvmlDeviceGetHandleByIndex(device) + mem_info = nvmlDeviceGetMemoryInfo(dev_handle) + gb_used = mem_info.used / 2**30 + gb_total = mem_info.total / 2**30 + output[device] = (gb_used, gb_total) + return output + + def wait_for_gpu_memory_to_clear( *, devices: list[int], - threshold_bytes: int | None = None, - threshold_ratio: float | None = None, + threshold_bytes: int | dict[int, int] | None = None, + threshold_ratio: float | dict[int, float] | None = None, timeout_s: float = 120, ) -> None: assert threshold_bytes is not None or threshold_ratio is not None - if ( - current_platform.is_rocm() - and threshold_ratio is not None - and threshold_ratio < 0.05 - ): + devices = get_physical_device_indices(devices) + if isinstance(threshold_bytes, int): + threshold_bytes = {device: threshold_bytes for device in devices} + elif isinstance(threshold_bytes, dict): + assert threshold_bytes.keys() == set(devices) + if isinstance(threshold_ratio, float): + threshold_ratio = {device: threshold_ratio for device in devices} + elif isinstance(threshold_ratio, dict): + assert threshold_ratio.keys() == set(devices) + if current_platform.is_rocm() and threshold_ratio is not None: # ROCm can keep a small runtime/driver footprint resident even after # all model allocations are gone. On MI300 this has been observed # around 2.5 GiB, which is above a strict 1% idle threshold but nowhere # near the amount of free memory needed by the next vLLM runner. - min_threshold_bytes = 4 * 1024**3 - threshold_bytes = max(threshold_bytes or 0, min_threshold_bytes) + min_threshold_b = 4 * 1024**3 + if threshold_bytes is None: + threshold_bytes = {} + for device, ratio in threshold_ratio.items(): + threshold_bytes[device] = max( + threshold_bytes.get(device, 0), min_threshold_b if ratio < 0.05 else 0 + ) # Use nvml instead of pytorch to reduce measurement error from torch cuda # context. - devices = get_physical_device_indices(devices) start_time = time.time() while True: - output: dict[int, str] = {} - output_raw: dict[int, tuple[float, float]] = {} - for device in devices: - if current_platform.is_rocm(): - dev_handle = amdsmi_get_processor_handles()[device] - mem_info = amdsmi_get_gpu_vram_usage(dev_handle) - gb_used = mem_info["vram_used"] / 2**10 - gb_total = mem_info["vram_total"] / 2**10 - else: - dev_handle = nvmlDeviceGetHandleByIndex(device) - mem_info = nvmlDeviceGetMemoryInfo(dev_handle) - gb_used = mem_info.used / 2**30 - gb_total = mem_info.total / 2**30 - output_raw[device] = (gb_used, gb_total) - output[device] = f"{gb_used:.02f}/{gb_total:.02f}" - + output_raw = record_gpu_memory_usage_stats(devices=devices) + output = { + device: f"{gb_used:.02f}/{gb_total:.02f}" + for device, (gb_used, gb_total) in output_raw.items() + } print("gpu memory used/total (GiB): ", end="") for k, v in output.items(): print(f"{k}={v}; ", end="") print("") if threshold_bytes is not None and threshold_ratio is not None: - threshold_gib = threshold_bytes / 2**30 - threshold = f"max({threshold_gib:.2f} GiB, {threshold_ratio:.3f})" + threshold_gib = { + device: threshold_b / 2**30 + for device, threshold_b in threshold_bytes.items() + } + threshold = "; ".join( + f"{device=}: max({threshold_gib[device]:.2f} GiB, " + f"{threshold_ratio[device]:.3f})" + for device in devices + ) all_free = all( - used <= max(threshold_gib, total * threshold_ratio) - for used, total in output_raw.values() + used <= max(threshold_gib[device], total * threshold_ratio[device]) + for device, (used, total) in output_raw.items() ) elif threshold_bytes is not None: - threshold_gib = threshold_bytes / 2**30 - threshold = f"{threshold_gib} GiB" - all_free = all(used <= threshold_gib for used, _ in output_raw.values()) + threshold_gib = { + device: threshold_b / 2**30 + for device, threshold_b in threshold_bytes.items() + } + threshold = "; ".join( + f"{device=}: {threshold_gib[device]:.2f} GiB" for device in devices + ) + all_free = all( + used <= threshold_gib[device] + for device, (used, _) in output_raw.items() + ) else: assert threshold_ratio is not None - threshold = f"{threshold_ratio:.3f}" + threshold = "; ".join( + f"{device=}: {threshold_ratio[device]:.3f}" for device in devices + ) all_free = all( - used / total <= threshold_ratio for used, total in output_raw.values() + used / total <= threshold_ratio[device] + for device, (used, total) in output_raw.items() ) dur_s = time.time() - start_time if all_free: - print( - f"Done waiting for free GPU memory on devices {devices=} " - f"({threshold=}) {dur_s=:.02f}" - ) + print(f"Done waiting for free GPU memory on ({threshold=}) {dur_s=:.02f}") break if dur_s >= timeout_s: raise ValueError( - f"Memory of devices {devices=} not free after " - f"{dur_s=:.02f} ({threshold=})" + f"Memory of devices not free after {dur_s=:.02f} ({threshold=})" ) time.sleep(5) @@ -1558,7 +1590,7 @@ def wait_for_gpu_memory_to_clear( def wait_for_rocm_memory_to_settle( *, - threshold_ratio: float = 0.1, + threshold_ratio: float | dict[int, float] | None = 0.1, timeout_s: float = 240, ) -> None: """Block until ROCm device VRAM usage drops below ``threshold_ratio``. diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 1ef4f96617e..315c77de392 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -19,6 +19,7 @@ from tests.v1.attention.utils import ( ) from vllm import _custom_ops as ops from vllm.config.vllm import set_current_vllm_config +from vllm.model_executor.layers.attention import mla_attention as mla_attention_module from vllm.model_executor.layers.attention.mla_attention import ( MLAAttention, QueryLenSupport, @@ -30,6 +31,7 @@ from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla +from vllm.v1.attention.backends.mla import flashmla as flashmla_module from vllm.v1.attention.backends.mla.prefill import ( MLAPrefillBackendEnum, get_mla_prefill_backend, @@ -552,6 +554,10 @@ class MockMLAAttentionLayer(MLAAttention): ) else: mqa_q = (mqa_ql_nope, mqa_q_pe) + if self.impl.dcp_world_size > 1: + if isinstance(mqa_q, tuple): + mqa_q = torch.cat(mqa_q, dim=-1) + mqa_q = mla_attention_module.get_dcp_group().all_gather(mqa_q, dim=1) attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) @@ -569,6 +575,215 @@ class MockMLAAttentionLayer(MLAAttention): return output +def test_mock_mla_dcp_fp8_decode_gathers_quantized_query( + monkeypatch, default_vllm_config +): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for FP8 decode query quantization path.") + + device = torch.device(f"{DEVICE_TYPE}:0") + num_tokens = 2 + num_heads = 2 + qk_nope_head_dim = 4 + qk_rope_head_dim = 2 + v_head_dim = 3 + kv_lora_rank = 5 + + class _DummyKVProj: + def __init__(self): + # Shape expected by MockMLAAttentionLayer.__init__ + self.weight = torch.randn( + num_heads * (qk_nope_head_dim + v_head_dim), + kv_lora_rank, + device=device, + dtype=torch.float32, + ) + + class _FakeImpl: + def __init__(self): + self.kv_cache_dtype = "fp8" + self.supports_quant_query_input = True + self.dcp_world_size = 2 + self.forward_q = None + + def forward_mha(self, *args, **kwargs): + return None + + def forward_mqa(self, q, kv_cache, attn_metadata, layer): + self.forward_q = q + assert isinstance(q, torch.Tensor) + bsz, _, _ = q.shape + return ( + torch.zeros( + bsz, + num_heads, + kv_lora_rank, + device=q.device, + dtype=torch.float32, + ), + None, + ) + + class _FakeDCPGroup: + def __init__(self): + self.calls = 0 + self.input_dtype = None + self.input_shape = None + + def all_gather(self, x, dim=1): + self.calls += 1 + self.input_dtype = x.dtype + self.input_shape = tuple(x.shape) + return torch.cat([x, x], dim=dim) + + fake_group = _FakeDCPGroup() + monkeypatch.setattr(mla_attention_module, "get_dcp_group", lambda: fake_group) + + impl = _FakeImpl() + with set_current_vllm_config(default_vllm_config): + layer = MockMLAAttentionLayer( + impl=impl, + num_heads=num_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + kv_lora_rank=kv_lora_rank, + device=device, + kv_b_proj=_DummyKVProj(), + q_scale=1.0, + k_scale=1.0, + ) + + q = torch.randn( + num_tokens, + num_heads, + qk_nope_head_dim + qk_rope_head_dim, + device=device, + dtype=torch.float32, + ) + kv_c = torch.randn(num_tokens, kv_lora_rank, device=device, dtype=torch.float32) + k_pe = torch.randn( + num_tokens, 1, qk_rope_head_dim, device=device, dtype=torch.float32 + ) + kv_cache = torch.empty(0, device=device, dtype=torch.float32) + output = torch.empty( + num_tokens, num_heads * v_head_dim, device=device, dtype=torch.float32 + ) + + class _AttnMeta: + num_decode_tokens = num_tokens + num_decodes = 1 + num_prefills = 0 + slot_mapping = torch.empty(0, dtype=torch.long, device=device) + + layer.forward_impl(q, kv_c, k_pe, kv_cache, _AttnMeta(), output) + + assert fake_group.calls == 1 + assert fake_group.input_dtype == current_platform.fp8_dtype() + assert fake_group.input_shape == ( + num_tokens, + num_heads, + kv_lora_rank + qk_rope_head_dim, + ) + assert isinstance(impl.forward_q, torch.Tensor) + assert tuple(impl.forward_q.shape) == ( + num_tokens, + num_heads * impl.dcp_world_size, + kv_lora_rank + qk_rope_head_dim, + ) + + +@pytest.mark.parametrize("is_fp8_kvcache", [False, True], ids=["bf16", "fp8"]) +def test_flashmla_dcp_decode_metadata_uses_gathered_query_heads( + monkeypatch, is_fp8_kvcache +): + class _FakeSchedulerMetadata: + tile_scheduler_metadata = None + num_splits = None + + base_call: tuple[torch.Tensor, int, int, bool] | None = None + fp8_call: tuple[torch.Tensor, int, int] | None = None + + def fake_get_mla_metadata( + seq_lens_device, + num_q_tokens_per_head_k, + num_heads_k, + is_fp8_kvcache=False, + ): + nonlocal base_call + base_call = ( + seq_lens_device, + num_q_tokens_per_head_k, + num_heads_k, + is_fp8_kvcache, + ) + return _FakeSchedulerMetadata(), None + + def fake_get_mla_metadata_dense_fp8( + seq_lens_device, num_q_tokens_per_head_k, num_heads_k + ): + nonlocal fp8_call + fp8_call = ( + seq_lens_device, + num_q_tokens_per_head_k, + num_heads_k, + ) + return ( + torch.empty((0, 8), dtype=torch.int32), + torch.empty((0,), dtype=torch.int32), + ) + + monkeypatch.setattr(flashmla_module, "get_mla_metadata", fake_get_mla_metadata) + monkeypatch.setattr( + flashmla_module, + "get_mla_metadata_dense_fp8", + fake_get_mla_metadata_dense_fp8, + ) + + builder = object.__new__(flashmla_module.FlashMLAMetadataBuilder) + builder.num_q_heads = 4 + builder.dcp_world_size = 2 + builder.is_fp8_kvcache = is_fp8_kvcache + builder.compilation_config = type( + "_CompilationConfig", + (), + { + "cudagraph_mode": type( + "_CudaGraphMode", + (), + {"has_full_cudagraphs": lambda self: False}, + )() + }, + )() + + seq_lens = torch.tensor([16, 24], dtype=torch.int32) + query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32) + + metadata = builder._build_decode( + block_table_tensor=torch.empty((2, 1), dtype=torch.int32), + seq_lens_device=seq_lens, + max_seq_len=24, + query_start_loc_cpu=query_start_loc, + query_start_loc_device=query_start_loc, + num_decode_tokens=2, + dcp_tot_seq_lens_device=None, + ) + + assert base_call is not None + assert base_call[0] is seq_lens + assert base_call[1:] == (8, 1, is_fp8_kvcache) + if is_fp8_kvcache: + assert metadata.scheduler_metadata.tile_scheduler_metadata is not None + assert metadata.scheduler_metadata.num_splits is not None + assert fp8_call is not None + assert fp8_call[0] is seq_lens + assert fp8_call[1:] == (8, 1) + else: + assert metadata.scheduler_metadata.tile_scheduler_metadata is None + assert metadata.scheduler_metadata.num_splits is None + assert fp8_call is None + + def run_attention_backend( backend: AttentionBackendEnum, kv_cache_spec: MLAAttentionSpec, diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index f4b7ee520ad..647241ce73c 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -7,9 +7,8 @@ from unittest.mock import MagicMock import pytest import torch -from vllm import envs from vllm.v1.core.kv_cache_utils import ( - _get_kv_cache_config_deepseek_v4, + _get_kv_cache_config_packed, get_kv_cache_config_from_groups, ) from vllm.v1.kv_cache_interface import ( @@ -84,15 +83,19 @@ def _make_groups(n_c4, n_c128, n_swa): return [mla_group, swa_group] -def _mock_vllm_config(): +def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None): config = MagicMock() config.cache_config.num_gpu_blocks_override = None + config.kv_transfer_config = None + if kv_connector_extra_config is not None: + config.kv_transfer_config = MagicMock() + config.kv_transfer_config.kv_connector_extra_config = kv_connector_extra_config return config def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024): groups = _make_groups(n_c4, n_c128, n_swa) - return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem) + return _get_kv_cache_config_packed(_mock_vllm_config(), groups, mem) def _page_sizes_by_layer( @@ -135,7 +138,7 @@ class TestInterleavedPacking: def test_strided_views_are_independent(self): groups = _make_groups(n_c4=3, n_c128=2, n_swa=5) page_sizes = _page_sizes_by_layer(groups) - num_blocks, tensors = _get_kv_cache_config_deepseek_v4( + num_blocks, tensors = _get_kv_cache_config_packed( _mock_vllm_config(), groups, 100 * 1024 * 1024 ) backing = torch.zeros(tensors[0].size, dtype=torch.uint8) @@ -156,8 +159,7 @@ class TestInterleavedPacking: for i, v in enumerate(views): assert (v == i + 1).all(), f"View {i} was corrupted" - def test_hma_attention_groups_keep_default_backing(self, monkeypatch): - monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False) + def test_hma_attention_groups_keep_default_backing(self): full = _make_full_spec() sw = _make_sw_spec() page_size = full.page_size_bytes @@ -178,8 +180,7 @@ class TestInterleavedPacking: KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]), ] - def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch): - monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False) + def test_hma_attention_groups_use_packed_backing_with_enable_cross_layers(self): full = _make_full_spec() sw = _make_sw_spec() page_size = full.page_size_bytes @@ -190,7 +191,9 @@ class TestInterleavedPacking: ] config = get_kv_cache_config_from_groups( - _mock_vllm_config(), groups, available_memory=page_size * 2 * 32 + _mock_vllm_config({"enable_cross_layers_blocks": "True"}), + groups, + available_memory=page_size * 2 * 32, ) assert config.num_blocks == 32 diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 004adf4a67b..d0168c9a935 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1302,6 +1302,101 @@ def test_no_spec_tokens_scheduled_for_prefill_chunks(): assert len(output.scheduled_spec_decode_tokens[req.request_id]) == num_spec_tokens +def _model_output(scheduler, output, sampled): + """Feed `sampled` (per-request list) back to the scheduler.""" + req_ids = list(output.num_scheduled_tokens.keys()) + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={r: i for i, r in enumerate(req_ids)}, + sampled_token_ids=sampled, + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ), + ) + + +def test_spec_decode_padding_first_decode_step(): + """A request taking its first decode step (whole prompt already computed via + a prefix-cache hit) is padded with placeholder (-1) spec tokens so it enters + the worker with the same 1 + num_spec_tokens shape as the other speculative + decodes, keeping the batch uniform. + """ + num_spec = 3 + scheduler = create_scheduler( + num_speculative_tokens=num_spec, + enable_prefix_caching=True, + block_size=16, + ) + # Two identical 33-token prompts: 2 full blocks (32 tokens) get cached, so a + # second identical request hits num_computed == num_prompt_tokens - 1. + r1, r2 = create_requests( + num_requests=2, num_tokens=33, same_prompt=True, max_tokens=16 + ) + + # Drive r1 through prefill so its prompt blocks are cached, then give it real + # drafts so it is a running speculative decode (1 + num_spec shape). + scheduler.add_request(r1) + out = scheduler.schedule() + assert out.num_scheduled_tokens[r1.request_id] == 33 + _model_output(scheduler, out, [[100]]) + scheduler.update_draft_token_ids(DraftTokenIds([r1.request_id], [[1, 2, 3]])) + + # r2 arrives; its whole prompt is a prefix-cache hit -> first decode step. + scheduler.add_request(r2) + out = scheduler.schedule() + + # r1 verifies its real drafts. + assert out.scheduled_spec_decode_tokens[r1.request_id] == [1, 2, 3] + # r2 is padded to the 1 + num_spec shape with placeholder (-1) drafts. + assert out.num_scheduled_tokens[r2.request_id] == 1 + num_spec + assert out.scheduled_spec_decode_tokens[r2.request_id] == [-1] * num_spec + + +def test_spec_decode_padding_skipped_with_prefill_in_batch(): + """Padding is skipped when the batch contains a prefill chunk: the batch is + already mixed/non-uniform, so padding a new decode request buys nothing. + """ + num_spec = 3 + scheduler = create_scheduler( + num_speculative_tokens=num_spec, + enable_prefix_caching=True, + block_size=16, + max_num_batched_tokens=64, + ) + # r_warm + r_candidate share a prompt so r_candidate gets a full prefix hit. + r_warm, r_candidate = create_requests( + num_requests=2, num_tokens=33, same_prompt=True, max_tokens=1 + ) + # r_long has a different, long prompt that prefills over multiple chunks. + (r_long,) = create_requests(num_requests=1, num_tokens=100, max_tokens=16) + + # Warm the prefix cache with r_warm's prompt (it finishes; blocks stay cached). + scheduler.add_request(r_warm) + out = scheduler.schedule() + assert out.num_scheduled_tokens[r_warm.request_id] == 33 + _model_output(scheduler, out, [[100]]) + assert r_warm.request_id in scheduler.finished_req_ids + + # Start r_long; after one chunk it remains a prefill chunk in the running queue. + scheduler.add_request(r_long) + out = scheduler.schedule() + _model_output(scheduler, out, [[]]) # still prefilling, no sampled token + assert r_long.is_prefill_chunk + + # r_candidate arrives (prefix-cache hit -> first decode step) alongside the + # in-flight prefill chunk. + scheduler.add_request(r_candidate) + out = scheduler.schedule() + + # The batch has a prefill chunk, so r_candidate is NOT padded. + assert r_long.request_id in out.num_scheduled_tokens + assert out.num_scheduled_tokens[r_candidate.request_id] == 1 + assert r_candidate.request_id not in out.scheduled_spec_decode_tokens + + def test_scheduler_stats_waiting_queues(): """Test that scheduler stats correctly report waiting and skipped_waiting queues.""" # Create scheduler with limited capacity so we can have waiting requests diff --git a/tests/v1/e2e/general/test_context_length.py b/tests/v1/e2e/general/test_context_length.py index c9dc8354fa1..cd0aff79de8 100644 --- a/tests/v1/e2e/general/test_context_length.py +++ b/tests/v1/e2e/general/test_context_length.py @@ -75,8 +75,9 @@ def test_auto_fit_max_model_len_rejects_oversized_input( must see this reduced value and reject prompts that exceed it, rather than accepting them and hanging.""" - # Use a tiny KV cache budget to force auto-fit to a very small - # max_model_len (e.g. ~16 tokens). + # Use a small KV cache budget to force auto-fit to a small + # max_model_len. Pin block_size=16 so the budget is independent + # of the platform's default block size. kv_cache_bytes = 1_000_000 # 1 MB with vllm_runner( @@ -84,6 +85,7 @@ def test_auto_fit_max_model_len_rejects_oversized_input( max_model_len=-1, max_num_seqs=1, enforce_eager=True, + block_size=16, kv_cache_memory_bytes=kv_cache_bytes, load_format="dummy", ) as vllm_model: diff --git a/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py b/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py new file mode 100644 index 00000000000..71484d3b05c --- /dev/null +++ b/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU chunked-prefill / prefix-caching correctness for linear-attention models.""" + +import os + +import pytest + +from tests.models.utils import check_logprobs_close +from vllm import LLM, SamplingParams +from vllm.platforms import current_platform + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +# Bound the KV cache so the run does not scale with host memory; these engines +# only need a few thousand tokens. +os.environ.setdefault("VLLM_CPU_KVCACHE_SPACE", "1") + +MODEL = "Qwen/Qwen3.5-0.8B" +CHUNK_TOKENS = 128 # max_num_batched_tokens for the chunked engine +NUM_LOGPROBS = 5 +SP = SamplingParams(max_tokens=32, temperature=0, logprobs=NUM_LOGPROBS) + + +def _long_prompt(repeat: int) -> str: + return "Solve the following arithmetic step by step. " * repeat + "What is 7*8?" + + +# Prompts long enough to span several CHUNK_TOKENS-sized chunks; a single-chunk +# prompt is bit-identical to full prefill regardless of the bug. +PROMPTS = [_long_prompt(r) for r in (40, 60, 80)] +# Spans several full cache blocks; prefix caching only reuses complete blocks. +PREFIX_PROMPT = "You are a helpful assistant. " * 230 + " Now answer: what is 2+2?" + + +def _make_llm(**overrides) -> LLM: + base = dict( + model=MODEL, + dtype="bfloat16", + max_model_len=2048, + enforce_eager=True, + trust_remote_code=True, + ) + base.update(overrides) + return LLM(**base) + + +def _tuples(outputs) -> list[tuple[list[int], str, object]]: + """(token_ids, text, sample_logprobs) per request, for check_logprobs_close.""" + return [ + (list(o.outputs[0].token_ids), o.outputs[0].text, o.outputs[0].logprobs) + for o in outputs + ] + + +@pytest.fixture(scope="module") +def full_prefill_refs(): + """Reference (ids, text, logprobs) for PROMPTS and PREFIX_PROMPT, full prefill.""" + llm = _make_llm(enable_chunked_prefill=False, enable_prefix_caching=False) + refs = _tuples(llm.generate(PROMPTS, SP)) + prefix_ref = _tuples(llm.generate([PREFIX_PROMPT], SP))[0] + del llm + return refs, prefix_ref + + +def test_chunked_prefill_matches_full_prefill(full_prefill_refs): + """Batched multi-chunk prefill must stay close to per-prompt full prefill. + + Prompts are scheduled together so the scheduler interleaves prefill chunks + across requests (the cross-request path where the accuracy gap was strongest). + """ + refs, _ = full_prefill_refs + llm = _make_llm( + enable_chunked_prefill=True, + max_num_batched_tokens=CHUNK_TOKENS, + enable_prefix_caching=False, + ) + got = _tuples(llm.generate(PROMPTS, SP)) + del llm + + check_logprobs_close( + outputs_0_lst=refs, + outputs_1_lst=got, + name_0="full_prefill", + name_1="chunked_prefill", + ) + + +def test_prefix_cache_hit_matches_cold_cache(full_prefill_refs): + """A prefix-cache hit must stay close to the cold-cache (reference) output. + + The warm run continues prefill from the restored GDN state; the + num_cached_tokens check guards against a vacuous (no-hit) pass. + """ + _, ref = full_prefill_refs + llm = _make_llm(enable_prefix_caching=True) + llm.generate([PREFIX_PROMPT], SP) # prime the cache + warm_out = llm.generate([PREFIX_PROMPT], SP)[0] + warm = _tuples([warm_out])[0] + del llm + + assert warm_out.num_cached_tokens > 0, ( + "expected a prefix-cache hit but num_cached_tokens=0; " + "PREFIX_PROMPT may be shorter than one cache block" + ) + check_logprobs_close( + outputs_0_lst=[ref], + outputs_1_lst=[warm], + name_0="cold_cache", + name_1="warm_cache", + ) diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index 57602289ce6..9ce225c3a49 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -28,6 +28,8 @@ hybrid_ssm_configs=( # GDN (Qwen3.5) "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" "VLLM_SSM_CONV_STATE_LAYOUT=DS PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=Qwen/Qwen3.5-0.8B" + # Mamba1 (Jamba) + "VLLM_SSM_CONV_STATE_LAYOUT=DS GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ai21labs/AI21-Jamba2-3B VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" ) sw_attn_configs=( # NOTE: gemma3 does not work with FlashInfer diff --git a/tests/v1/kv_connector/nixl_integration/test_accuracy.py b/tests/v1/kv_connector/nixl_integration/test_accuracy.py index eead3de1532..bb68b7a5724 100644 --- a/tests/v1/kv_connector/nixl_integration/test_accuracy.py +++ b/tests/v1/kv_connector/nixl_integration/test_accuracy.py @@ -25,6 +25,7 @@ EXPECTED_VALUES = { "ibm-granite/granite-4.0-h-tiny": 0.77, "Qwen/Qwen3.5-0.8B": 0.33, "google/gemma-4-E2B-it": 0.485, + "ai21labs/AI21-Jamba2-3B": 0.74, } SIMPLE_PROMPT = ( diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 32abd05242f..ad5792e6c3c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import ( SlidingWindowSpec, ) from vllm.v1.kv_offload.base import ( + LookupResult, OffloadingManager, OffloadPolicy, ReqContext, @@ -484,7 +485,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # full 3 blocks hit [0, 1, 2] runner.new_request(token_ids=[0] * (block_size * 3 + 1)) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.run( decoded_tokens=[EOS_TOKEN_ID], # Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2 @@ -504,7 +505,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # 3 blocks are hit on GPU [0, 1, 2] # 1 block loaded [3,] runner.new_request(token_ids=[0] * (block_size * 4 + 1)) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.run( decoded_tokens=[EOS_TOKEN_ID], # Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2 @@ -632,7 +633,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # 48 tokens (3 block) from the second group # Total 48 tokens can be loaded runner.new_request(token_ids=[0] * 48) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -648,7 +649,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # extra tokens [0, 36] (blocks [4, 5, 6]) from the first group # extra tokens [0, 32] (block [3, 4]) from the second group runner.new_request(token_ids=[0] * (48 + 37)) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -665,12 +666,12 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool def _make_scheduler_with_lookup( - lookup_results: dict[int, bool | None], + lookup_results: dict[int, LookupResult], ) -> OffloadingConnectorScheduler: """Create an OffloadingConnectorScheduler with a mocked manager.lookup.""" manager = MagicMock(spec=OffloadingManager) manager.lookup.side_effect = lambda key, req_context: lookup_results.get( - int(get_offload_block_hash(key).decode()), False + int(get_offload_block_hash(key).decode()), LookupResult.MISS ) scheduler = object.__new__(OffloadingConnectorScheduler) @@ -683,7 +684,7 @@ _EMPTY_REQ_CTX = ReqContext(req_id="") class TestMaximalPrefixLookup: def test_all_hit(self): - sched = _make_scheduler_with_lookup({1: True, 2: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2 def test_all_miss(self): @@ -691,32 +692,54 @@ class TestMaximalPrefixLookup: assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 def test_partial_prefix(self): - sched = _make_scheduler_with_lookup({1: True, 2: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2 def test_miss_then_hit(self): - sched = _make_scheduler_with_lookup({2: True}) + sched = _make_scheduler_with_lookup({2: LookupResult.HIT}) assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 def test_single_hit(self): - sched = _make_scheduler_with_lookup({1: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT}) assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1 def test_empty(self): sched = _make_scheduler_with_lookup({}) assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0 - def test_none_defers(self): - sched = _make_scheduler_with_lookup({1: None, 2: True}) + def test_retry_defers(self): + sched = _make_scheduler_with_lookup( + {1: LookupResult.RETRY, 2: LookupResult.HIT} + ) + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + assert sched.manager.lookup.call_count == 2 + + def test_retry_after_hit_defers(self): + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.RETRY} + ) assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None - def test_none_after_hit_defers(self): - sched = _make_scheduler_with_lookup({1: True, 2: None}) + def test_hit_pending_defers(self): + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT} + ) assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + assert sched.manager.lookup.call_count == 2 - def test_none_stops_at_miss(self): - """None is treated as hit for iteration, but miss stops the scan.""" - sched = _make_scheduler_with_lookup({1: None, 2: False, 3: True}) + def test_hit_pending_does_not_stop_scan(self): + """HIT_PENDING defers but does not break — scan continues until miss.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT_PENDING, 2: LookupResult.MISS, 3: LookupResult.HIT} + ) + assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None + assert sched.manager.lookup.call_count == 2 + + def test_retry_stops_at_miss(self): + """RETRY is treated as hit for iteration, but miss stops the scan.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.RETRY, 2: LookupResult.MISS, 3: LookupResult.HIT} + ) assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None # lookup should have been called for blocks 1 and 2 (stops at miss) assert sched.manager.lookup.call_count == 2 @@ -724,7 +747,7 @@ class TestMaximalPrefixLookup: class TestSlidingWindowLookup: def test_all_hit_exact_window(self): - sched = _make_scheduler_with_lookup({1: True, 2: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2 def test_all_miss(self): @@ -732,25 +755,27 @@ class TestSlidingWindowLookup: assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 0 def test_window_at_end(self): - sched = _make_scheduler_with_lookup({2: True, 3: True}) + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 2, _EMPTY_REQ_CTX) == 3 def test_window_in_middle(self): - sched = _make_scheduler_with_lookup({2: True, 3: True}) + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) assert ( sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 3 ) def test_no_full_window_falls_back_to_prefix(self): - sched = _make_scheduler_with_lookup({1: True, 2: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) == 2 def test_single_block_window(self): - sched = _make_scheduler_with_lookup({2: True, 3: True}) + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 3 def test_gap_resets_consecutive(self): - sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True}) + sched = _make_scheduler_with_lookup( + {2: LookupResult.HIT, 3: LookupResult.HIT, 4: LookupResult.HIT} + ) # [1, 2, 3, 0, 4] — gap at 0 resets, window of 2 found at [2,3] assert ( sched._sliding_window_lookup(to_keys([1, 2, 3, 0, 4]), 2, _EMPTY_REQ_CTX) @@ -758,7 +783,14 @@ class TestSlidingWindowLookup: ) def test_window_prefers_rightmost(self): - sched = _make_scheduler_with_lookup({1: True, 2: True, 4: True, 5: True}) + sched = _make_scheduler_with_lookup( + { + 1: LookupResult.HIT, + 2: LookupResult.HIT, + 4: LookupResult.HIT, + 5: LookupResult.HIT, + } + ) # two valid windows: [1,2] at positions 0-1 and [4,5] at positions 3-4 # scans right-to-left, finds [4,5] first assert ( @@ -767,7 +799,14 @@ class TestSlidingWindowLookup: ) def test_prefix_fallback_with_gap(self): - sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True, 5: True}) + sched = _make_scheduler_with_lookup( + { + 2: LookupResult.HIT, + 3: LookupResult.HIT, + 4: LookupResult.HIT, + 5: LookupResult.HIT, + } + ) # window of 4 not found contiguously (gap at 1) assert ( sched._sliding_window_lookup(to_keys([2, 1, 3, 4, 5]), 4, _EMPTY_REQ_CTX) @@ -778,20 +817,47 @@ class TestSlidingWindowLookup: sched = _make_scheduler_with_lookup({}) assert sched._sliding_window_lookup([], 1, _EMPTY_REQ_CTX) == 0 - def test_none_defers(self): - sched = _make_scheduler_with_lookup({1: True, 2: None}) + def test_retry_defers(self): + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.RETRY} + ) assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None - def test_none_with_full_window_still_defers(self): - """Even if a real window is found after a None, result is deferred.""" - # Scan right-to-left: 4(True), 3(None) resets, 2(True), 1(True) = window - # but block 3 was None so defer_lookup is set - sched = _make_scheduler_with_lookup({1: True, 2: True, 3: None, 4: True}) + def test_retry_with_full_window_still_defers(self): + """Even if a real window is found after a RETRY, result is deferred.""" + # Scan right-to-left: 4(HIT), 3(RETRY) resets, 2(HIT), 1(HIT) = window + # but block 3 was RETRY so defer_lookup is set + sched = _make_scheduler_with_lookup( + { + 1: LookupResult.HIT, + 2: LookupResult.HIT, + 3: LookupResult.RETRY, + 4: LookupResult.HIT, + } + ) assert ( sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) is None ) + def test_hit_pending_counts_as_hit(self): + """HIT_PENDING counts toward the consecutive-hit streak.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING} + ) + # window=2: both count as hits, but defer_lookup is set + assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None + + def test_hit_pending_does_not_break_streak(self): + """HIT_PENDING in the middle of a window doesn't reset the streak.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING, 3: LookupResult.HIT} + ) + # window=3: right-to-left finds 3(HIT),2(HIT_PENDING),1(HIT) = 3 consecutive + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) is None + ) + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool): @@ -1485,7 +1551,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): # Verify that loads still work correctly for the stored SWA blocks. runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * num_tokens + [1]) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2 runner.run( decoded_tokens=[EOS_TOKEN_ID], @@ -1701,7 +1767,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -1732,7 +1800,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -1762,7 +1832,7 @@ class TestEagle: async_scheduling=False, kv_cache_groups=groups, ) - runner.manager.lookup.return_value = False + runner.manager.lookup.return_value = LookupResult.MISS sched = runner.connector_scheduler req_status = self._make_req_status( sched, num_tokens=8, offload_keys_per_group=[[1, 2]] @@ -1803,7 +1873,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3, 4} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3, 4} + else LookupResult.MISS ) sched = runner.connector_scheduler @@ -1858,7 +1930,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -1891,7 +1965,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler # num_tokens=13 → max_hit=13-1=12, query_max=min(12+4,12)=12 @@ -1940,7 +2016,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -1995,7 +2073,9 @@ class TestEagle: # Group 0 keys [10,11,12]: only 10 hits. # Group 1 keys [1,2,3]: all hit. runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {10, 1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {10, 1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -2046,7 +2126,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -2273,7 +2355,7 @@ class TestEagle: runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * offloaded_block_size * 3 + [1]) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index a232082879d..482a2f25a56 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -46,6 +46,7 @@ from vllm.v1.kv_offload.base import ( CanonicalKVCaches, GPULoadStoreSpec, LoadStoreSpec, + LookupResult, OffloadingManager, OffloadingSpec, OffloadingWorker, @@ -129,9 +130,8 @@ class MockOffloadingSpec(OffloadingSpec): super().__init__(vllm_config, kv_cache_config) self.manager = MagicMock(spec=OffloadingManager) - self.manager.lookup.return_value = 0 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) - self.manager.lookup.return_value = False + self.manager.lookup.return_value = LookupResult.MISS self.manager.on_new_request.return_value = RequestOffloadingContext() self.handler = MockOffloadingWorker() diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 8ef1277bb39..7e291962987 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py @@ -19,7 +19,6 @@ def _make_bare_scheduler() -> MooncakeStoreScheduler: scheduler.lookup_async = False scheduler._block_size = 16 scheduler.load_specs = {} - scheduler._preempted_req_ids = set() scheduler._unfinished_request_ids = {"req-0"} scheduler._unfinished_requests = {} scheduler._request_trackers = {} @@ -35,6 +34,7 @@ def _make_scheduler_output(*, scheduled_spec_tokens: list[int] | None): req_ids=["req-0"], new_block_ids=[([2],)], num_computed_tokens=[44], + resumed_req_ids=set(), ), num_scheduled_tokens={"req-0": 4}, scheduled_spec_decode_tokens=( @@ -52,6 +52,7 @@ def _make_preemption_scheduler_output(): req_ids=[], new_block_ids=[], num_computed_tokens=[], + resumed_req_ids=set(), ), num_scheduled_tokens={}, scheduled_spec_decode_tokens={}, @@ -195,6 +196,7 @@ def _make_pending_load_scheduler_output() -> SimpleNamespace: req_ids=[], new_block_ids=[], num_computed_tokens=[], + resumed_req_ids=set(), ), num_scheduled_tokens={}, scheduled_spec_decode_tokens={}, @@ -253,14 +255,17 @@ def _make_resumed_unfinished_request( def _make_resumed_scheduler_output(*, num_scheduled_tokens: int) -> SimpleNamespace: + # A resumed-from-preemption step: the scheduler lists the request in + # resumed_req_ids and sends the FULL block table (replace semantics). return SimpleNamespace( finished_req_ids=set(), preempted_req_ids=set(), scheduled_new_reqs=[], scheduled_cached_reqs=SimpleNamespace( req_ids=["req-0"], - new_block_ids=[([2],)], + new_block_ids=[([0, 1, 2],)], num_computed_tokens=[0], + resumed_req_ids={"req-0"}, ), num_scheduled_tokens={"req-0": num_scheduled_tokens}, scheduled_spec_decode_tokens={}, @@ -273,7 +278,6 @@ def test_resumed_from_preemption_with_load_skips_save(): # passes load_spec.can_load=True. Skip save in this step; subsequent # cached_reqs steps will save new tokens normally. scheduler = _make_bare_scheduler() - scheduler._preempted_req_ids = {"req-0"} _make_resumed_unfinished_request( scheduler, token_ids=list(range(48)), @@ -303,7 +307,6 @@ def test_resumed_from_preemption_with_load_skips_save(): def test_resumed_from_preemption_without_load_still_saves(): # No load_spec → behavior is unchanged: save proceeds. scheduler = _make_bare_scheduler() - scheduler._preempted_req_ids = {"req-0"} _make_resumed_unfinished_request( scheduler, token_ids=list(range(48)), @@ -324,6 +327,71 @@ def test_resumed_from_preemption_without_load_still_saves(): assert tracker.num_saved_tokens == 48 +def test_running_request_not_in_resumed_req_ids_appends_blocks(): + """Regression: the replace-vs-append choice must follow the scheduler's + cached_reqs.resumed_req_ids, NOT connector-local preemption history. + + A running request that is not resumed this step carries a *delta* + new_block_ids and must be APPENDED to the tracker's existing blocks. + Treating it as resumed would replace allocated_block_ids with just the + delta while token_len stays at the full computed length, so the store + path's block_ids[start // block_size] runs off the end (the + "list index out of range" / token_len >> len(block_ids) bug). + """ + scheduler = _make_bare_scheduler() + _add_unfinished_request( + scheduler, + token_ids=list(range(48)), + block_hashes=[b"h0", b"h1", b"h2"], + prefill_end_tokens=48, + ) + + out = _make_scheduler_output(scheduled_spec_tokens=None) + assert "req-0" not in out.scheduled_cached_reqs.resumed_req_ids + + meta = scheduler.build_connector_meta(out) + + tracker = scheduler._request_trackers["req-0"] + # Delta [2] appended to existing [0, 1] (decode path), not replaced by [2]. + assert tracker.allocated_block_ids == ([0, 1, 2],) + # token_len stays covered by the block table: no store-path under-count. + blocks_held = sum(len(g) for g in tracker.allocated_block_ids) + assert tracker.token_len // scheduler._block_size <= blocks_held + assert len(meta.requests) == 1 + assert meta.requests[0].token_len_chunk == 48 + + +def test_resumed_request_in_resumed_req_ids_replaces_blocks(): + """A request the scheduler marks resumed gets the FULL block table in + new_block_ids and must REPLACE the tracker's blocks (not append), even if + a stale tracker from before preemption is still present.""" + scheduler = _make_bare_scheduler() + _make_resumed_unfinished_request( + scheduler, + token_ids=list(range(48)), + block_hashes=[b"h0", b"h1", b"h2"], + num_computed_tokens=0, + ) + # Stale pre-preemption tracker that must be overwritten, not appended to. + scheduler._request_trackers["req-0"] = RequestTracker( + req_id="req-0", + token_len=99, + allocated_block_ids=([7, 8, 9],), + num_saved_tokens=0, + ) + + scheduler.build_connector_meta( + _make_resumed_scheduler_output(num_scheduled_tokens=48) + ) + + tracker = scheduler._request_trackers["req-0"] + # Replaced with the full table from new_block_ids, not appended to [7,8,9]. + assert tracker.allocated_block_ids == ([0, 1, 2],) + assert tracker.token_len == 48 + blocks_held = sum(len(g) for g in tracker.allocated_block_ids) + assert tracker.token_len // scheduler._block_size <= blocks_held + + # Focused tests for ReqMeta.from_request_tracker — the centralized guard that # enforces "a ReqMeta never carries both a save and a load". diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index dce582946b5..d6ce200f4cf 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -4,6 +4,7 @@ import json import logging import math +import queue import sys import threading import types @@ -760,7 +761,7 @@ def test_store_worker_get_block_ids_with_load_errors_delegates_to_recv_thread(): recv_thread = MagicMock() recv_thread.get_and_clear_block_ids_with_load_errors.return_value = {3, 4} w = _make_bare_worker() - w.kv_recv_thread = recv_thread + w.kv_recv_threads = [recv_thread] assert w.get_block_ids_with_load_errors() == {3, 4} recv_thread.get_and_clear_block_ids_with_load_errors.assert_called_once_with() @@ -1574,7 +1575,9 @@ def _make_bare_worker( worker.put_step = 1 worker.enable_kv_events = False worker.kv_send_thread = None - worker.kv_recv_thread = None + worker.kv_recv_threads = [] + worker.num_recv_threads = 1 + worker.recv_request_queue = queue.Queue() worker.tp_size = 1 worker.num_kv_head = 1 worker.pp_size = 1 diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index d508f3cae2c..c6bcdc2896c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -630,6 +630,19 @@ def _make_mock_worker_for_desc_ids( worker._has_mamba = has_mamba worker._group_spec_types = group_spec_types worker.block_len_per_layer = block_len_per_layer or [100] + worker._conv_decomp = None + if has_mamba: + from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( # noqa: E501 + MambaConvSplitInfo, + ) + + # Mamba2/GDN layout: 3 conv sub-projections -> 4 NIXL regions per layer. + worker._conv_decomp = MambaConvSplitInfo( + conv_rows=3, + local_proj_dims=(1, 1, 1), + conv_dtype_size=2, + ssm_sizes=(0, 0), + ) worker._compute_desc_ids = NixlConnectorWorker._compute_desc_ids.__get__( worker, NixlConnectorWorker ) @@ -976,6 +989,37 @@ def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_rati (256, 256, 768), id="qwen35_27b_tp8", ), + # ai21labs/AI21-Jamba2-Mini (Mamba1) + # mamba d_inner = mamba_expand(2) * hidden_size(4096) = 8192 + # mamba_d_state=16, mamba_d_conv=4 → conv_rows=3. + # Conv state holds only x: a single contiguous sub-projection. + pytest.param( + "mamba1", + 1, + 8192, + 3, + (8192, 16), + (8192,), + id="jamba_mini_tp1", + ), + pytest.param( + "mamba1", + 4, + 2048, + 3, + (2048, 16), + (2048,), + id="jamba_mini_tp4", + ), + pytest.param( + "mamba1", + 8, + 1024, + 3, + (1024, 16), + (1024,), + id="jamba_mini_tp8", + ), ], ) def test_derive_mamba_conv_split( @@ -999,6 +1043,7 @@ def test_derive_mamba_conv_split( from vllm.v1.kv_cache_interface import MambaSpec _TYPE_MAP = { + "mamba1": MambaAttentionBackendEnum.MAMBA1, "mamba2": MambaAttentionBackendEnum.MAMBA2, "gdn_attention": MambaAttentionBackendEnum.GDN_ATTN, } diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index d568357224c..89a0374b462 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -8,6 +8,7 @@ import pytest from vllm.v1.kv_offload.base import ( LoadStoreSpec, + LookupResult, OffloadingEvent, OffloadKey, PrepareStoreOutput, @@ -160,7 +161,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) # block 2 must still be present in the cache - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT def test_filter_reused_manager_reports_stores_skipped_counter(): @@ -242,8 +243,8 @@ def test_cpu_manager(): ) # lookup [1, 2] -> write in-flight, not yet ready - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING # no events so far assert list(cpu_manager.take_events()) == [] @@ -253,9 +254,9 @@ def test_cpu_manager(): verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare store [2, 3, 4, 5] -> evicts [1] prepare_store_output = cpu_manager.prepare_store( @@ -280,12 +281,12 @@ def test_cpu_manager(): cpu_manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) # lookup (now that we have [2, 3, 4, 5]) - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare load [2, 3] prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) @@ -329,8 +330,8 @@ def test_cpu_manager(): cpu_manager.complete_store(to_keys([7, 9]), _EMPTY_REQ_CTX, success=False) # assert [7] is still stored, but [9] is not - assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is LookupResult.MISS verify_events( cpu_manager.take_events(), @@ -412,8 +413,8 @@ class TestARCPolicy: ) # lookup [1, 2] -> write in-flight, not yet ready - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING # no events so far assert list(cpu_manager.take_events()) == [] @@ -423,9 +424,9 @@ class TestARCPolicy: verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS # blocks should be in T1 (recent) assert len(arc_policy.t1) == 2 @@ -629,7 +630,7 @@ class TestARCPolicy: cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False) # block 5 should not be in cache - assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.MISS # block 5 should not be in T1 or T2 assert to_keys([5])[0] not in arc_policy.t1 assert to_keys([5])[0] not in arc_policy.t2 @@ -670,8 +671,8 @@ class TestARCPolicy: cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX) # verify blocks 2, 3 (in T2) are still present - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT # verify events events = list(cpu_manager.take_events()) @@ -691,8 +692,8 @@ def test_filter_reused_manager(): ) # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet - assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare store [1, 2] -> should be filtered prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -700,7 +701,7 @@ def test_filter_reused_manager(): assert prepare_store_output.keys_to_store == [] # Lookup [1] -> 2nd time, eligible now - assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -709,13 +710,13 @@ def test_filter_reused_manager(): # Lookup [3, 4] -> 1st time # (evicts [2] from tracker since max_size is 3 and tracker has [1]) - assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False - assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS + assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.MISS # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4]) assert to_keys([2])[0] not in manager.counts # Lookup [2] again -> (this adds [2] back to the tracker as 1st time) - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS # Verify [2] was re-added with count=1 (not eligible yet) assert manager.counts.get(to_keys([2])[0]) == 1 diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 9e19bd18fec..7245ae1ba7a 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -18,7 +18,12 @@ import numpy as np import pytest import torch -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + make_offload_key, +) from vllm.v1.kv_offload.tiering.base import JobMetadata from vllm.v1.kv_offload.tiering.fs.manager import ( FileSystemTierManager, @@ -166,7 +171,7 @@ def fs_tier(tmp_path): def test_lookup_empty_tier(fs_tier): tier, _ = fs_tier results = lookup_and_wait(tier, [key(1), key(2)]) - assert results == [False, False] + assert results == [LookupResult.MISS, LookupResult.MISS] def test_store_creates_file_and_lookup_succeeds(fs_tier): @@ -176,7 +181,7 @@ def test_store_creates_file_and_lookup_succeeds(fs_tier): results = drain(tier) assert len(results) == 1 assert results[0].success - assert lookup_and_wait(tier, [key(1)]) == [True] + assert lookup_and_wait(tier, [key(1)]) == [LookupResult.HIT] dest = tier.file_mapper.get_file_name(key(1)) assert os.path.exists(dest), f"Expected file at {dest}" @@ -188,14 +193,20 @@ def test_store_then_load_roundtrip(fs_tier): store_results = drain(tier) assert all(r.success for r in store_results) - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] + assert lookup_and_wait(tier, [key(1), key(2)]) == [ + LookupResult.HIT, + LookupResult.HIT, + ] job_l = make_job(2, [key(1), key(2)], [2, 3], is_promotion=True) tier.submit_load(job_l) load_results = drain(tier) assert all(r.success for r in load_results) # Blocks stay on disk after load - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] + assert lookup_and_wait(tier, [key(1), key(2)]) == [ + LookupResult.HIT, + LookupResult.HIT, + ] def test_invalid_path_raises_at_construction(): @@ -231,7 +242,10 @@ def test_multiple_jobs_tracked_independently(fs_tier): results = drain(tier) job_ids = {r.job_id for r in results} assert job_ids == {1, 2} - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] + assert lookup_and_wait(tier, [key(1), key(2)]) == [ + LookupResult.HIT, + LookupResult.HIT, + ] def test_multi_block_job_partial_failure(fs_tier): diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index aae3c60c539..28570926db2 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -17,7 +17,12 @@ from unittest.mock import MagicMock, patch import numpy as np import torch -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + make_offload_key, +) from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager @@ -236,19 +241,19 @@ class TestMockObjTierBasic: self.tier, self.agent = _make_tier(num_blocks=4) def test_lookup_empty_tier(self): - assert lookup_and_wait(self.tier, [key(1)]) == [False] + assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.MISS] def test_store_and_lookup(self): self.tier.submit_store(make_job(1, [key(1)], [0])) results = drain(self.tier) assert len(results) == 1 assert results[0].success - assert lookup_and_wait(self.tier, [key(1)]) == [True] + assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.HIT] def test_lookup_unrelated_key_returns_false(self): self.tier.submit_store(make_job(1, [key(1)], [0])) drain(self.tier) - assert lookup_and_wait(self.tier, [key(999)]) == [False] + assert lookup_and_wait(self.tier, [key(999)]) == [LookupResult.MISS] def test_store_then_load_roundtrip(self): self.tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) @@ -327,13 +332,17 @@ class TestMockObjTierMultiBlock: results = drain(tier) assert len(results) == 1 assert results[0].success - assert lookup_and_wait(tier, keys) == [True] * 8 + assert lookup_and_wait(tier, keys) == [LookupResult.HIT] * 8 def test_partial_block_lookup(self): tier, _ = _make_tier(num_blocks=4) tier.submit_store(make_job(1, [key(0), key(1)], [0, 1])) drain(tier) - assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [True, True, False] + assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [ + LookupResult.HIT, + LookupResult.HIT, + LookupResult.MISS, + ] class TestMockObjTierFailures: @@ -342,7 +351,7 @@ class TestMockObjTierFailures: agent.query_memory = lambda *a, **k: (_ for _ in ()).throw( RuntimeError("backend error") ) - assert lookup_and_wait(tier, [key(1)]) == [False] + assert lookup_and_wait(tier, [key(1)]) == [LookupResult.MISS] def test_submit_store_register_memory_failure_reported_in_get_finished(self): tier, agent = _make_tier(num_blocks=4) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index fca84532445..f06b91aa208 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -21,6 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) from vllm.v1.kv_offload.base import ( + LookupResult, OffloadingCounterMetadata, OffloadKey, OffloadPolicy, @@ -60,15 +61,15 @@ def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]: def count_hits(manager, keys: list[OffloadKey]) -> int | None: """Count consecutive lookup hits from the start of keys. - Returns the count of leading True results, or None if any lookup - returns None (retry-later signal). + Returns the count of leading HIT results, or None if any lookup + returns HIT_PENDING or RETRY. """ count = 0 for key in keys: result = manager.lookup(key, _CTX) - if result is None: + if result in (LookupResult.HIT_PENDING, LookupResult.RETRY): return None - if not result: + if result is not LookupResult.HIT: break count += 1 return count @@ -185,18 +186,18 @@ class TestExampleSecondaryTierManager: # Initially empty blocks = to_keys(range(3)) - assert tier.lookup(blocks[0], _CTX) is False + assert tier.lookup(blocks[0], _CTX) is LookupResult.MISS # Store blocks (simulate with direct insertion for testing) tier.blocks[blocks[0]] = True tier.blocks[blocks[1]] = True # Lookup should find first two blocks - assert tier.lookup(blocks[0], _CTX) is True - assert tier.lookup(blocks[1], _CTX) is True + assert tier.lookup(blocks[0], _CTX) is LookupResult.HIT + assert tier.lookup(blocks[1], _CTX) is LookupResult.HIT # Third block not present - assert tier.lookup(blocks[2], _CTX) is False + assert tier.lookup(blocks[2], _CTX) is LookupResult.MISS class TestTieringOffloadingManager: @@ -283,8 +284,12 @@ class TestTieringOffloadingManager: assert self.secondary_tier2.get_num_blocks() == 3 # Verify blocks are present - assert all(self.secondary_tier1.lookup(b, _CTX) for b in blocks) - assert all(self.secondary_tier2.lookup(b, _CTX) for b in blocks) + assert all( + self.secondary_tier1.lookup(b, _CTX) is LookupResult.HIT for b in blocks + ) + assert all( + self.secondary_tier2.lookup(b, _CTX) is LookupResult.HIT for b in blocks + ) def test_ref_cnt_protection_during_cascade(self, manager_setup): """Test that ref_cnt protects blocks during cascade.""" @@ -355,7 +360,7 @@ class TestTieringOffloadingManager: # Lookup each block to initiate promotion for all of them for block in blocks: result = self.manager.lookup(block, _CTX) - assert result is None # Retry later (promotion initiated) + assert result is LookupResult.RETRY # promotion initiated # End of step 1: flushes deferred submit_load() calls self._simulate_on_schedule_end() @@ -470,11 +475,11 @@ class TestTieringOffloadingManager: ctx_a = ReqContext(req_id="req_a") ctx_b = ReqContext(req_id="req_b") - # All lookups return None: secondary hit triggers promotion (in-flight) - assert self.manager.lookup(blocks[0], ctx_a) is None - assert self.manager.lookup(blocks[1], ctx_a) is None - assert self.manager.lookup(blocks[2], ctx_b) is None - assert self.manager.lookup(blocks[3], ctx_b) is None + # All lookups return RETRY: secondary hit triggers promotion + assert self.manager.lookup(blocks[0], ctx_a) is LookupResult.RETRY + assert self.manager.lookup(blocks[1], ctx_a) is LookupResult.RETRY + assert self.manager.lookup(blocks[2], ctx_b) is LookupResult.RETRY + assert self.manager.lookup(blocks[3], ctx_b) is LookupResult.RETRY # submit_load must not fire during lookup - only at end of step self.secondary_tier1.submit_load.assert_not_called() @@ -511,9 +516,10 @@ class TestTieringOffloadingManager: result_a = self.manager.lookup(shared_block, ctx_a) result_b = self.manager.lookup(shared_block, ctx_b) - # Both see None (in-flight), but promotion is only queued once - assert result_a is None - assert result_b is None + # First lookup triggers promotion (RETRY), second finds block + # already in primary with write in-flight (HIT_PENDING). + assert result_a is LookupResult.RETRY + assert result_b is LookupResult.HIT_PENDING self._simulate_on_schedule_end() @@ -796,7 +802,10 @@ class TestTieringOffloadingManager: # the lookup that staged it). promo_block = to_keys([99])[0] self.secondary_tier1.blocks[promo_block] = True - assert self.manager.lookup(promo_block, ReqContext(req_id="pending")) is None + assert ( + self.manager.lookup(promo_block, ReqContext(req_id="pending")) + is LookupResult.RETRY + ) assert self.manager._pending_load_submissions # Request-level tier registration. @@ -829,7 +838,7 @@ class TestTieringOffloadingManager: assert self.primary_tier._num_allocated_blocks == 0 assert self.primary_tier._free_list == [] for block in blocks: - assert self.primary_tier.lookup(block, _CTX) is False + assert self.primary_tier.lookup(block, _CTX) is LookupResult.MISS # Pending submission was dropped, not submitted. self.secondary_tier1.submit_load.assert_not_called() diff --git a/tests/v1/logits_processors/test_correctness.py b/tests/v1/logits_processors/test_correctness.py index c93593865e0..a38d8a6cf71 100644 --- a/tests/v1/logits_processors/test_correctness.py +++ b/tests/v1/logits_processors/test_correctness.py @@ -1226,3 +1226,34 @@ def test_thinking_budget_invalid_budget_rejected(invalid_budget): with pytest.raises(VLLMValidationError, match="thinking_token_budget"): SamplingParams(thinking_token_budget=invalid_budget) + + +def test_thinking_budget_long_thinking_section_end_marker_found_at_correct_index(): + """Test thinking budget enforced for a long thinking run, + then a natural end marker.""" + h = ThinkingBudgetStateHolder( + MockReasoningConfig(), 8, 0, torch.device("cpu"), False + ) + h.sync_batch( + BatchUpdate( + batch_size=1, + removed=(), + added=[(0, SamplingParams(thinking_token_budget=10_000), None, [])], + moved=(), + ) + ) + start = MockReasoningConfig.reasoning_start_token_ids + end = MockReasoningConfig.reasoning_end_token_ids + + out: list[int] = list(start) + h.update_state([out], None, None) + for tok in range(500): # 500 filler thinking tokens, one decode step each + out.append(tok) + h.update_state([out], None, None) + assert h._state[0]["end_thinking"] == -1 # not present yet + expected_end_idx = len(out) # marker appended next + out.extend(end) + h.update_state([out], None, None) + + assert h._state[0]["start_thinking"] == 0 + assert h._state[0]["end_thinking"] == expected_end_idx diff --git a/tests/v1/sample/test_logprobs.py b/tests/v1/sample/test_logprobs.py index 863c1e7a8e5..5ed0a476279 100644 --- a/tests/v1/sample/test_logprobs.py +++ b/tests/v1/sample/test_logprobs.py @@ -1263,3 +1263,37 @@ def test_prompt_logprobs_with_chunking_and_preemption(): assert preemptions > 0, "Test did not trigger any preemptions" print(f"Test passed with {preemptions} preemptions") + + +@large_gpu_mark(min_gb=24) +def test_token_logprobs_large_batch_int64_row_offset(): + """Regression: logprob kernel row offset (row * vocab_size) must use int64. + + The rejection-sampler logprobs path runs the logprob kernels over the + spec-expanded logits batch, so batch_size * vocab_size can exceed 2**31 + (e.g. DFlash drafts K tokens per request). With int32 offset arithmetic the + per-row pointer wraps to a negative address and the kernel hits a CUDA + illegal memory access. Run over a batch where batch_size * vocab_size > 2**31 + and check the highest-offset row matches a reference log-softmax. + """ + if not current_platform.is_cuda(): + pytest.skip("int32 row-offset overflow is a CUDA kernel issue") + from vllm.v1.worker.gpu.sample.logprob import compute_token_logprobs + + device = torch.device("cuda") + vocab_size = 131072 + batch_size = 2**31 // vocab_size + 64 # batch_size * vocab_size > 2**31 + # logits (the large input) plus small logprob/rank outputs; ~1 GB headroom. + required_bytes = batch_size * vocab_size * 4 + (1 << 30) + if torch.cuda.mem_get_info()[0] < required_bytes: + pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory") + + logits = torch.randn(batch_size, vocab_size, device=device, dtype=torch.float32) + token_ids = torch.full((batch_size, 1), 7, device=device, dtype=torch.int64) + logprobs = compute_token_logprobs(logits, token_ids) + torch.accelerator.synchronize() # surface any async illegal memory access + last = batch_size - 1 + ref = torch.log_softmax(logits[last].float(), dim=-1)[7] + assert torch.allclose(logprobs[last, 0], ref, atol=1e-2), ( + f"logprob {logprobs[last, 0].item()} != ref {ref.item()}" + ) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 047e2b754ef..8d906e83f2d 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -4,6 +4,7 @@ import pytest import torch from torch import Generator +from tests.utils import large_gpu_mark from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON from vllm.utils.torch_utils import set_random_seed @@ -404,6 +405,45 @@ class TestTritonTopkTopp: self._compare_results(logits, k, p) + @large_gpu_mark(min_gb=24) + def test_large_batch_int64_row_offset(self): + """Regression: per-row offset (row * vocab_size) must not overflow int32. + + Speculative decoding expands the logits batch (e.g. DFlash drafts K + tokens per request), so batch_size * vocab_size can exceed 2**31. With + int32 offset arithmetic the per-row pointer wraps to a negative address + and the kernel hits a CUDA illegal memory access. Use a batch where + batch_size * vocab_size > 2**31 and give the highest-offset row the same + logits as row 0: an overflow there would read a different row and change + the kept set. + """ + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton + + if not current_platform.is_cuda(): + pytest.skip("int32 row-offset overflow is a CUDA kernel issue") + vocab_size = 131072 + batch_size = 2**31 // vocab_size + 64 # batch_size * vocab_size > 2**31 + # logits is modified in place; the only extra device memory is the + # per-SM scratch buffer (~num_sm * vocab), so allow ~1 GB of headroom. + required_bytes = batch_size * vocab_size * 4 + (1 << 30) + if torch.cuda.mem_get_info()[0] < required_bytes: + pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory") + + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + logits[batch_size - 1] = logits[0] + k = torch.full((batch_size,), 5, dtype=torch.int32) + result = apply_top_k_top_p_triton(logits, k, None) + torch.accelerator.synchronize() # surface any async illegal memory access + kept_first = (result[0] > float("-inf")).nonzero(as_tuple=True)[0] + kept_last = (result[batch_size - 1] > float("-inf")).nonzero(as_tuple=True)[0] + assert kept_first.numel() == 5, f"row 0 kept {kept_first.numel()}, expected 5" + assert torch.equal(kept_first, kept_last), ( + "highest-offset row produced a different top-k mask than the " + "identical row 0 (int32 row-offset overflow)" + ) + @pytest.mark.parametrize( "mode", ["topk_only", "topp_only", "topk_and_topp"], diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 1ec986eada6..09586f5e6b4 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -1655,6 +1655,7 @@ def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None expected_cp = dcp_world_size * pcp_world_size assert sched.cp_world_size == expected_cp assert sched.block_size == BLOCK_SIZE * expected_cp + assert sched.fa_block_size == BLOCK_SIZE * expected_cp # --------------------------------------------------------------------------- @@ -1734,6 +1735,86 @@ def test_cp_eager_store_and_load_roundtrip( assert len(meta2.load_cpu_blocks) == num_blocks +# --------------------------------------------------------------------------- +# Test 18: CP store and load use effective block size +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "dcp_world_size, pcp_world_size", + [ + (2, 1), + (1, 2), + (2, 2), + ], +) +def test_cp_effective_block_size_store_and_load( + dcp_world_size: int, pcp_world_size: int +) -> None: + """Verify ready_blocks_g (store) and n_take_g (load) use the effective + (physical * cp) block size, not the per-rank physical size.""" + fix = _make_cp_scheduler( + dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size + ) + sched = fix.scheduler + gpu_pool = fix.gpu_block_pool + cp = dcp_world_size * pcp_world_size + vbs = BLOCK_SIZE * cp + + # Store: allocate 2 blocks, confirm only 1. Without the fix, + # ready_blocks_g = vbs / BLOCK_SIZE = 2, storing both blocks. + req = _make_cp_request(num_blocks=2, virtual_block_size=vbs) + gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs) + kv = KVCacheBlocks(blocks=(gpu_blocks,)) + req.num_computed_tokens = vbs + sched.update_state_after_alloc(req, kv, num_external_tokens=0) + m1 = sched.build_connector_meta( + make_scheduler_output( + {req.request_id: vbs}, + new_reqs={req.request_id: kv.get_block_ids()}, + ) + ) + assert len(m1.store_gpu_blocks) == 1 + assert len(m1.store_cpu_blocks) == 1 + simulate_store_completion(sched, m1.store_event) + + # Load: store 2 blocks from a second request, accept only 1 as external. + # Without the fix, n_take_g = vbs / BLOCK_SIZE = 2, loading both. + req2 = _make_cp_request(num_blocks=2, virtual_block_size=vbs) + kv2 = KVCacheBlocks(blocks=(_allocate_cp_gpu_blocks(gpu_pool, req2, 2, vbs),)) + req2.num_computed_tokens = 2 * vbs + sched.update_state_after_alloc(req2, kv2, num_external_tokens=0) + m2 = sched.build_connector_meta( + make_scheduler_output( + {req2.request_id: 2 * vbs}, + new_reqs={req2.request_id: kv2.get_block_ids()}, + ) + ) + simulate_store_completion(sched, m2.store_event) + + req3 = Request( + request_id="req-cp-partial-load", + prompt_token_ids=req2.prompt_token_ids, + sampling_params=req2.sampling_params, + pooling_params=None, + mm_features=None, + block_hasher=req2._block_hasher, + ) + hit, _ = sched.get_num_new_matched_tokens(req3, num_computed_tokens=0) + assert hit == 2 * vbs + + kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(2),)) + sched.update_state_after_alloc(req3, kv3, num_external_tokens=vbs) + m3 = sched.build_connector_meta( + make_scheduler_output( + {req3.request_id: vbs}, + new_reqs={req3.request_id: kv3.get_block_ids()}, + ) + ) + assert m3.load_event >= 0 + assert len(m3.load_gpu_blocks) == 1 + assert len(m3.load_cpu_blocks) == 1 + assert m3.load_gpu_blocks == [kv3.get_block_ids()[0][0]] + + # --------------------------------------------------------------------------- # Test 17: CP lazy target blocks are scaled correctly # --------------------------------------------------------------------------- diff --git a/tools/build_rust.py b/tools/build_rust.py index e5c5d0bb2e4..b5951bfe576 100644 --- a/tools/build_rust.py +++ b/tools/build_rust.py @@ -27,7 +27,7 @@ def rust_extensions(*, optional: bool = False) -> list[RustExtension]: ), RustExtension( target="vllm._rust_tool_parser", - path="rust/src/tool-parser/python/Cargo.toml", + path="rust/src/parser/python/Cargo.toml", features=["pyo3/abi3-py38"], binding=Binding.PyO3, optional=optional, diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index d1470029216..ee706037abb 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -98,11 +98,13 @@ FI_ALLREDUCE_FUSION_MAX_SIZE_MB: dict[int, dict[int, float]] = { 2: 64, # 64MB 4: 32, # 32MB 8: 1, # 1MB + 16: 64, # 64MB (mnnvl multi-node) }, 103: { 2: 64, # 64MB 4: 64, # 64MB 8: 2, # 2MB + 16: 64, # 64MB (mnnvl multi-node) }, } diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index bc38ec6a8a8..4a392a7e3bd 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -186,7 +186,7 @@ class PassConfig: """ MiB = 1024 * 1024 - FI_SUPPORTED_WORLD_SIZES = [2, 4, 8] + FI_SUPPORTED_WORLD_SIZES = [2, 4, 8, 16] if world_size not in FI_SUPPORTED_WORLD_SIZES: return None max_size_mb = self.fi_allreduce_fusion_max_size_mb diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py index c54eaf7555d..3c9d759dd7f 100644 --- a/vllm/distributed/device_communicators/quick_all_reduce.py +++ b/vllm/distributed/device_communicators/quick_all_reduce.py @@ -28,11 +28,13 @@ from vllm.distributed.utils import is_weak_contiguous # noqa: E402, F401 class QuickReduceRegime(Enum): + # Keep integer ids aligned with csrc/quickreduce/quick_reduce.h FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 - NONE = 4 + INT3 = 4 + NONE = 5 KB = 1024 @@ -43,14 +45,20 @@ class QuickAllReduce: _SUPPORTED_WORLD_SIZES = [2, 4, 8] _SUPPORTED_DTYPES = [torch.float16, torch.bfloat16] # The following data is based on kernel tests. - # In this order [FP, INT8, INT6, INT4]. + # In this order [FP, INT8, INT6, INT4, INT3]. _QR_MIN_SIZE = { - (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB], - (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB], - (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB], - (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB], - (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB], - (torch.bfloat16, 8): [16 * MB, 2048 * MB, 2048 * MB, 2048 * MB], + (torch.float16, 2): [1 * MB, 2 * MB, 2 * MB, 1 * MB, 1 * MB], + (torch.float16, 4): [1 * MB, 16 * MB, 4 * MB, 2 * MB, 2 * MB], + (torch.float16, 8): [16 * MB, 4 * MB, 4 * MB, 2 * MB, 2 * MB], + (torch.bfloat16, 2): [2 * MB, 8 * MB, 8 * MB, 8 * MB, 8 * MB], + (torch.bfloat16, 4): [8 * MB, 64 * MB, 64 * MB, 16 * MB, 16 * MB], + (torch.bfloat16, 8): [ + 16 * MB, + 2048 * MB, + 2048 * MB, + 2048 * MB, + 2048 * MB, + ], } def __init__(self, group: ProcessGroup, device: int | str | torch.device) -> None: @@ -59,8 +67,10 @@ class QuickAllReduce: available for CUDA and ROCm MI300 series. Custom quick allreduce leverages quantization for further - acceleration on ROCm. It currently supports Q8, Q6, and Q4 - quantization formats and FP(float16, bfloat16). + acceleration on ROCm. It currently supports Q8, Q6, Q4, and Q3 + quantization formats and FP(float16, bfloat16). Q3 (INT3) is + restricted to TP2 (world_size == 2) due to poor performance on + larger world sizes. Quick allreduce is designed as a complement to custom allreduce. Its initialization requires even stricter conditions. @@ -178,6 +188,23 @@ class QuickAllReduce: ) return self.qr_quant_level = QuickReduceRegime[regime_str] + + # INT3 is only enabled for TP2 (world_size == 2). + # Kernel benchmarks show INT3 all-reduce on TP4/TP8 has poor + # performance (the extra ranks make the 3-bit codec's pack/unpack + # overhead outweigh the reduced communication volume), so INT3 is + # restricted to 2-GPU tensor parallelism. For TP4/TP8 use a wider + # codec (e.g. INT4) or NONE instead. + if self.qr_quant_level == QuickReduceRegime.INT3 and self.world_size != 2: + logger.warning( + "Custom quick allreduce is disabled: INT3 quantization is " + "only supported for TP2 (world_size == 2), but world_size " + "is %d. INT3 on TP4/TP8 is disabled due to poor kernel " + "performance. Use INT4/NONE for this world size.", + self.world_size, + ) + return + self.qr_quantization_min_size = self._get_qr_quantization_min_size() vllm_config = get_current_vllm_config_or_none() if ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index 620fa2f5ba1..58dfd5e428e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py @@ -67,7 +67,6 @@ class MooncakeStoreScheduler: # Per-request state self.load_specs: dict[str, LoadSpec] = {} # to be loaded self._request_trackers: dict[str, RequestTracker] = {} # scheduled new requests - self._preempted_req_ids: set[str] = set() # preempted requests self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {} self._unfinished_request_ids: set[str] = set() @@ -175,10 +174,8 @@ class MooncakeStoreScheduler: self._request_trackers.pop(finished_req_id, None) self._unfinished_requests.pop(finished_req_id, None) self._unfinished_request_ids.discard(finished_req_id) - self._preempted_req_ids.discard(finished_req_id) preempted_ids = scheduler_output.preempted_req_ids or set() - self._preempted_req_ids.update(preempted_ids) for req_id in preempted_ids: self.load_specs.pop(req_id, None) if request_tracker := self._request_trackers.get(req_id): @@ -243,13 +240,12 @@ class MooncakeStoreScheduler: continue req_meta = None - if req_id in self._preempted_req_ids: + if req_id in cached_reqs.resumed_req_ids: # Resumed after preemption if isinstance(new_block_ids, tuple): new_block_ids = tuple(b.copy() for b in new_block_ids) else: new_block_ids = (new_block_ids.copy(),) - self._preempted_req_ids.discard(req_id) load_spec = self.load_specs.pop(req_id, None) request_tuple = self._unfinished_requests.get(req_id) request_real = request_tuple[0] # type: ignore[index] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 127b8e4d4b1..aea2d602e72 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -359,6 +359,7 @@ class KVTransferThread(threading.Thread): ready_event: threading.Event, name: str, record_operation: Callable[..., None] | None = None, + request_queue: queue.Queue[Any] | None = None, ): super().__init__(daemon=True, name=name) self.store = store @@ -368,7 +369,7 @@ class KVTransferThread(threading.Thread): self.token_databases = token_databases self._record_operation_cb = record_operation self.done_task_lock = threading.Lock() - self.request_queue: queue.Queue[Any] = queue.Queue() + self.request_queue: queue.Queue[Any] = request_queue or queue.Queue() self.finished_requests: set[str] = set() self.kv_event_lock = threading.Lock() self.kv_events: list[BlockStored] = [] @@ -389,6 +390,7 @@ class KVTransferThread(threading.Thread): def run(self): self.ready_event.set() while True: + request_data = None try: request_data = self.request_queue.get() if request_data is None: @@ -396,8 +398,9 @@ class KVTransferThread(threading.Thread): self.request_queue.task_done() continue self._handle_request(request_data) - except Exception as e: - logger.error("Error in %s: %s", self.name, e) + except Exception: + req_id = getattr(request_data, "req_id", "") + logger.exception("Error in %s (req=%s)", self.name, req_id) def _handle_request(self, req_meta: Any): pass @@ -747,6 +750,7 @@ class KVCacheStoreRecvingThread(KVTransferThread): ready_event: threading.Event, disk_offload_buffer_budget_bytes: int | None = None, record_operation: Callable[..., None] | None = None, + request_queue: queue.Queue[Any] | None = None, ): super().__init__( store, @@ -756,6 +760,7 @@ class KVCacheStoreRecvingThread(KVTransferThread): ready_event, name="KVCacheStoreRecvingThread", record_operation=record_operation, + request_queue=request_queue, ) # _invalid_block_ids can be access by both the Worker and RecvingThread self._invalid_block_ids_lock = threading.Lock() @@ -1099,7 +1104,10 @@ class MooncakeStoreWorker: self.enable_kv_events = True self.kv_send_thread: KVCacheStoreSendingThread | None = None - self.kv_recv_thread: KVCacheStoreRecvingThread | None = None + # Pool of load-receive threads + self.kv_recv_threads: list[KVCacheStoreRecvingThread] = [] + self.num_recv_threads = max(1, envs.VLLM_MOONCAKE_LOAD_RECV_THREADS) + self.recv_request_queue: queue.Queue[ReqMeta] = queue.Queue() self.finished_store_req: set[str] = set() self._kv_connector_stats_lock = threading.Lock() self.kv_connector_stats = MooncakeStoreConnectorStats() @@ -1258,19 +1266,30 @@ class MooncakeStoreWorker: ) self.kv_send_thread.start() - ready_event_recving = threading.Event() - self.kv_recv_thread = KVCacheStoreRecvingThread( - self.store, - self.coord, - self.token_dbs, - self.block_size, - self.tp_rank, - ready_event_recving, - disk_offload_buffer_budget_bytes=self.disk_offload_buffer_budget_bytes, - record_operation=self._record_kv_connector_operation, + self.kv_recv_threads = [] + ready_events_recving = [] + for i in range(self.num_recv_threads): + ready_event_recving = threading.Event() + recv_thread = KVCacheStoreRecvingThread( + self.store, + self.coord, + self.token_dbs, + self.block_size, + self.tp_rank, + ready_event_recving, + disk_offload_buffer_budget_bytes=self.disk_offload_buffer_budget_bytes, + record_operation=self._record_kv_connector_operation, + request_queue=self.recv_request_queue, + ) + recv_thread.name = f"KVCacheStoreRecvingThread-{i}" + recv_thread.start() + self.kv_recv_threads.append(recv_thread) + ready_events_recving.append(ready_event_recving) + for ready_event_recving in ready_events_recving: + ready_event_recving.wait() + logger.info( + "Started %d Mooncake KV-load receive thread(s)", self.num_recv_threads ) - self.kv_recv_thread.start() - ready_event_recving.wait() def start_load_kv( self, @@ -1304,9 +1323,7 @@ class MooncakeStoreWorker: continue load_spec.token_len = load_spec.kvpool_cached_tokens - - assert self.kv_recv_thread is not None - self.kv_recv_thread.add_request(request) + self.recv_request_queue.put(request) assert self.load_async, "load_async must be True for better performance." # Issue stores with CUDA event synchronization @@ -1333,11 +1350,10 @@ class MooncakeStoreWorker: else set() ) - done_recving = ( - self.kv_recv_thread.get_and_clear_finished_requests() - if self.load_async and self.kv_recv_thread is not None - else set() - ) + done_recving: set[str] = set() + if self.load_async: + for recv_thread in self.kv_recv_threads: + done_recving |= recv_thread.get_and_clear_finished_requests() logger.debug( "Completed send: %d, recv: %d, tp_rank: %d", @@ -1348,9 +1364,10 @@ class MooncakeStoreWorker: return done_sending, done_recving def get_block_ids_with_load_errors(self) -> set[int]: - if self.kv_recv_thread is None: - return set() - return self.kv_recv_thread.get_and_clear_block_ids_with_load_errors() + block_ids: set[int] = set() + for recv_thread in self.kv_recv_threads: + block_ids |= recv_thread.get_and_clear_block_ids_with_load_errors() + return block_ids def _record_kv_connector_operation( self, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 060fa5e3228..66bee55f286 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -96,7 +96,13 @@ class NixlBaseConnectorWorker: ) -> np.ndarray: """Compute NIXL descriptor IDs for given block IDs.""" num_fa_regions = self.num_regions - num_ssm_regions = len(self.block_len_per_layer) * 4 if self._has_mamba else 0 + num_ssm_regions = 0 + if self._has_mamba: + assert self._conv_decomp is not None + # NIXL regions per SSM layer = conv sub-projections + 1 SSM temporal + # (Mamba2/GDN: 3+1=4; Mamba1: 1+1=2). + ssm_regions_per_layer = len(self._conv_decomp.local_conv_offsets) + 1 + num_ssm_regions = len(self.block_len_per_layer) * ssm_regions_per_layer num_blocks = dst_num_blocks if block_size_ratio is not None: @@ -279,8 +285,8 @@ class NixlBaseConnectorWorker: # ---- Model state (derived from model config) ---- mamba_ssm_size = (0, 0) # Conv state sub-projection decomposition (None when no Mamba). - # The 3-read transfer requires DS (dim, state_len) conv layout so - # that x/B/C sub-projections are contiguous in memory. + # The transfer requires DS (dim, state_len) conv layout so that + # conv sub-projections are contiguous in memory. self._conv_decomp: MambaConvSplitInfo | None = None self._has_mamba = any( isinstance(g.kv_cache_spec, MambaSpec) @@ -1186,8 +1192,8 @@ class NixlBaseConnectorWorker: base_addresses: list[int], block_size_ratio: int, ) -> list[tuple[int, int, int]]: - """Build 4 desc regions (x, B, C, ssm) per layer for local mamba - blocks, enabling the 3-read transfer with DS conv layout.""" + """Build desc regions (conv sub-projections + ssm) per layer for + local mamba blocks with DS conv layout.""" assert block_size_ratio == 1, ( "Mamba 3-read transfer with block_size_ratio != 1 is not tested. " f"Got block_size_ratio={block_size_ratio}." @@ -1227,9 +1233,9 @@ class NixlBaseConnectorWorker: tp_ratio: int, transfer_info: EngineTransferInfo, ) -> list[tuple[int, int, int]]: - """Build 4 remote desc regions (proj0, proj1, proj2, ssm) per layer - for the 3-read transfer. For hetero-TP, each D rank reads only its - sub-projection slice from the P rank.""" + """Build remote desc regions (conv sub-projections + ssm) per layer. + For hetero-TP, each D rank reads only its sub-projection slice from + the P rank.""" assert self._conv_decomp is not None effective_ratio = max(tp_ratio, 1) # Mamba conv state is always TP-sharded, even when attention KV diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index a35970a6160..7aa2b563ba7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -35,6 +35,7 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, + LookupResult, OffloadingManager, OffloadingSpec, OffloadKey, @@ -393,15 +394,18 @@ class OffloadingConnectorScheduler: hit_count = 0 defer_lookup = False for key in keys: - result = self.manager.lookup(key, req_context) - if result is None: - defer_lookup = True - # continue lookup to allow manager to kick-off async lookups - # for all blocks (until a miss is detected) - result = True - if not result: - break - hit_count += 1 + match self.manager.lookup(key, req_context): + case LookupResult.HIT: + hit_count += 1 + case LookupResult.HIT_PENDING: + defer_lookup = True + hit_count += 1 + case LookupResult.RETRY: + # Don't break: keep scanning to let manager kick off + # async lookups (until a miss is detected). + defer_lookup = True + case LookupResult.MISS: + break return hit_count if not defer_lookup else None def _sliding_window_lookup( @@ -416,18 +420,25 @@ class OffloadingConnectorScheduler: defer_lookup = False consecutive_hits = 0 for idx in range(len(keys) - 1, -1, -1): - result = self.manager.lookup(keys[idx], req_context) - if result is None: - defer_lookup = True - # continue lookup to allow manager to kick-off async lookups - # for all blocks (until a hit is detected) - result = False - if not result: - consecutive_hits = 0 - else: - consecutive_hits += 1 - if consecutive_hits == sliding_window_size: - return idx + sliding_window_size if not defer_lookup else None + match self.manager.lookup(keys[idx], req_context): + case LookupResult.HIT: + consecutive_hits += 1 + case LookupResult.HIT_PENDING: + # Block is in cache, just not readable yet — counts + # as hit for the consecutive streak. Don't break: + # keep scanning to let manager kick off async lookups. + defer_lookup = True + consecutive_hits += 1 + case LookupResult.RETRY: + # Block location uncertain — does not count as hit. + # Don't break: keep scanning to let manager kick off + # async lookups. + defer_lookup = True + consecutive_hits = 0 + case LookupResult.MISS: + consecutive_hits = 0 + if consecutive_hits == sliding_window_size: + return idx + sliding_window_size if not defer_lookup else None return consecutive_hits if not defer_lookup else None def _touch(self, req_status: RequestOffloadState): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 1e0435d371e..254e0dec09f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -124,13 +124,14 @@ class OffloadingConnectorWorker: raise NotImplementedError packed_kv_cache_tensor = next( - (t for t in kv_cache_config.kv_cache_tensors if t.block_stride), None + ( + t + for t in kv_cache_config.kv_cache_tensors + if t.block_stride and t.shared_by + ), + None, ) - is_dsv4 = all( - isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) - for group in kv_cache_config.kv_cache_groups - ) - if packed_kv_cache_tensor is not None and not is_dsv4: + if packed_kv_cache_tensor is not None: (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]] block_stride = tensor.stride(0) packed_tensor = tensor.as_strided( @@ -153,7 +154,7 @@ class OffloadingConnectorWorker: block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list) for kv_cache_tensor in kv_cache_config.kv_cache_tensors: # Filter to layers that were actually processed above. - # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for + # Packed KV allocation emits KVCacheTensor entries for # every (tuple_idx, page_size) slot; slots where no group has a # layer at that index produce an empty shared_by (reserved memory # with no corresponding model layer). diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py index 2a5510656bc..00dc05bfc4a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Mamba conv-state sub-projection decomposition for the 3-read transfer. +"""Mamba conv-state sub-projection decomposition for NIXL transfer. With DS conv state layout (dim, state_len), sub-projections are -contiguous in memory. Each D rank reads its slices via 3 separate +contiguous in memory. Each D rank reads its slices via separate RDMA transfers — no P-side permutation needed. Supported model types: + - Mamba1: conv = [x], temporal = (intermediate_size, state_size) - Mamba2: conv = [x, B, C], temporal = (num_heads, head_dim) - GDN (Gated Delta Net): conv = [Q, K, V] (dim(Q)==dim(K)), temporal = (num_v_heads, v_dim, k_dim) @@ -24,18 +25,21 @@ from vllm.v1.kv_cache_interface import MambaSpec @dataclass(frozen=True) class MambaConvSplitInfo: - """Per-rank byte sizes of the 3 conv sub-projections. + """Per-rank byte sizes of the conv sub-projections. Used by both P and D sides for NIXL descriptor registration. All fields are LOCAL to this engine's TP (already divided by TP size). DS memory layout within one page (contiguous): + Mamba1: |---- x ----| (single sub-projection, no decomposition) Mamba2: |-- x --|- B -|- C -| (B == C) GDN: |- Q -|- K -|-- V --| (dim(Q)==dim(K), V may differ) """ conv_rows: int # conv_kernel - 1 (typically 3) - local_proj_dims: tuple[int, int, int] # per-rank column counts per sub-proj + # Per-rank column counts per sub-projection: + # 1 entry for Mamba1, 3 for Mamba2/GDN. + local_proj_dims: tuple[int, ...] conv_dtype_size: int # bytes per element (e.g. 2 for float16) ssm_sizes: tuple[int, int] # (conv_state_bytes, ssm_state_bytes) @@ -45,10 +49,10 @@ class MambaConvSplitInfo: return sum(self.local_proj_dims) @property - def proj_bytes(self) -> tuple[int, int, int]: - """Byte sizes of the 3 sub-projections for one rank.""" + def proj_bytes(self) -> tuple[int, ...]: + """Byte sizes of the sub-projections for one rank.""" row_bytes = self.conv_rows * self.conv_dtype_size - return tuple(d * row_bytes for d in self.local_proj_dims) # type: ignore[return-value] + return tuple(d * row_bytes for d in self.local_proj_dims) @property def local_conv_offsets(self) -> list[tuple[int, int]]: @@ -57,8 +61,12 @@ class MambaConvSplitInfo: Used by both P and D for local descriptor registration. """ - conv0, conv1, conv2 = self.proj_bytes - return [(0, conv0), (conv0, conv1), (conv0 + conv1, conv2)] + offsets: list[tuple[int, int]] = [] + offset = 0 + for size in self.proj_bytes: + offsets.append((offset, size)) + offset += size + return offsets def remote_conv_offsets( self, local_rank_offset: int, tp_ratio: int @@ -76,28 +84,23 @@ class MambaConvSplitInfo: P page. Local dims are scaled down by |tp_ratio| to get P-sized offsets. """ - conv0, conv1, conv2 = self.proj_bytes + offsets: list[tuple[int, int]] = [] if tp_ratio >= 1: - remote_conv0 = conv0 * tp_ratio - remote_conv1 = conv1 * tp_ratio - return [ - (local_rank_offset * conv0, conv0), - (remote_conv0 + local_rank_offset * conv1, conv1), - (remote_conv0 + remote_conv1 + local_rank_offset * conv2, conv2), - ] + remote_base = 0 + for size in self.proj_bytes: + offsets.append((remote_base + local_rank_offset * size, size)) + remote_base += size * tp_ratio else: # NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages # are smaller than D's. Local dims are D-sized, but we need # P-sized offsets. Scale down by |tp_ratio|. abs_ratio = -tp_ratio - remote_conv0 = conv0 // abs_ratio - remote_conv1 = conv1 // abs_ratio - remote_conv2 = conv2 // abs_ratio - return [ - (0, remote_conv0), - (remote_conv0, remote_conv1), - (remote_conv0 + remote_conv1, remote_conv2), - ] + remote_base = 0 + for size in self.proj_bytes: + remote_size = size // abs_ratio + offsets.append((remote_base, remote_size)) + remote_base += remote_size + return offsets def derive_mamba_conv_split( @@ -120,12 +123,13 @@ def derive_mamba_conv_split( conv_dtype_size, and ssm_sizes (conv_state_bytes, ssm_state_bytes). """ _supported = ( + MambaAttentionBackendEnum.MAMBA1, MambaAttentionBackendEnum.MAMBA2, MambaAttentionBackendEnum.GDN_ATTN, ) if mamba_spec.mamba_type not in _supported: raise NotImplementedError( - f"3-read conv transfer only supports Mamba2 and GDN models, " + f"Conv transfer only supports Mamba1, Mamba2 and GDN models, " f"got mamba_type={mamba_spec.mamba_type!r}." ) @@ -149,7 +153,18 @@ def derive_mamba_conv_split( conv_state_bytes = torch.Size(mamba_spec.shapes[0]).numel() * conv_dtype_size ssm_state_bytes = torch.Size(mamba_spec.shapes[1]).numel() * ssm_dtype_size - if mamba_spec.mamba_type == MambaAttentionBackendEnum.MAMBA2: + local_proj_dims: tuple[int, ...] + if mamba_spec.mamba_type == MambaAttentionBackendEnum.MAMBA1: + # Mamba1 conv state holds only x (no B/C), so it's a single + # contiguous TP shard with no sub-projection decomposition. + temporal_shape = mamba_spec.shapes[1] + assert temporal_shape[0] == local_conv_dim, ( + f"Mamba1 temporal state dim ({temporal_shape[0]}) doesn't match " + f"conv dim ({local_conv_dim}); both should be " + f"intermediate_size/TP." + ) + local_proj_dims = (local_conv_dim,) + elif mamba_spec.mamba_type == MambaAttentionBackendEnum.MAMBA2: # NOTE (ZhanqiuHu): intermediate_size (= global x dim) is not stored # in MambaSpec, so we reconstruct it from the SSM temporal state shape: # shapes[1] = (local_num_heads, head_dim), already divided by TP. diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 0b41c4d7fa6..284b2511dba 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -17,8 +17,6 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ( ChatTemplateContentFormatOption, ConversationMessage, - get_history_tool_calls_cnt, - get_tool_call_id_type, make_tool_call_id, ) from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -170,8 +168,6 @@ class OpenAIServingChat(OpenAIServing): if mc.generation_config not in ("auto", "vllm") else getattr(mc, "override_generation_config", {}).get("max_new_tokens") ) - self.tool_call_id_type = get_tool_call_id_type(self.model_config) - # NOTE(woosuk): While OpenAI's chat completion API supports browsing # for some models, currently vLLM doesn't support it. Please use the # Responses API instead. @@ -261,6 +257,7 @@ class OpenAIServingChat(OpenAIServing): tokenizer, request.tools, chat_template_kwargs=chat_template_kwargs, + model_config=self.model_config, ) result = await self.render_chat_request(request) if isinstance(result, ErrorResponse): @@ -433,11 +430,6 @@ class OpenAIServingChat(OpenAIServing): else: tool_choice_function_name = None - if self.tool_call_id_type == "kimi_k2": - history_tool_call_cnt = get_history_tool_calls_cnt(conversation) - else: - history_tool_call_cnt = 0 - previous_texts = [""] * num_choices try: @@ -451,14 +443,10 @@ class OpenAIServingChat(OpenAIServing): tokenizer, request.tools, chat_template_kwargs=chat_template_kwargs, + model_config=self.model_config, ) for _ in range(num_choices) ] - for p in parsers: - if p is not None: - # NOTE: HarmonyParser ignores _stream_state (uses its own FSM). - p._stream_state.tool_call_id_type = self.tool_call_id_type - p._stream_state.history_tool_call_cnt = history_tool_call_cnt else: parsers = [None] * num_choices except Exception as e: @@ -842,10 +830,6 @@ class OpenAIServingChat(OpenAIServing): ) choices: list[ChatCompletionResponseChoice] = [] - if self.tool_call_id_type == "kimi_k2": - history_tool_call_cnt = get_history_tool_calls_cnt(conversation) - else: - history_tool_call_cnt = 0 role = self.get_chat_request_role(request) tool_parser_cls = ( @@ -885,54 +869,26 @@ class OpenAIServingChat(OpenAIServing): tool_calls = [] auto_tools_called = False + is_named_tool_choice = ( + request.tool_choice is not None + and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam + ) + is_required_tool_choice = request.tool_choice == "required" if (not self.enable_auto_tools or not tool_parser_cls) and ( - not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) - and request.tool_choice != "required" + not is_named_tool_choice and not is_required_tool_choice ): message = ChatMessage(role=role, reasoning=reasoning, content=content) - elif ( - request.tool_choice - and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam - ): - tool_call_items = [] - tool_calls = tool_calls or [] - for tc in tool_calls: - if not tc.id: - tc.id = make_tool_call_id( - id_type=self.tool_call_id_type, - func_name=tc.name, - idx=history_tool_call_cnt, - ) - tool_call_items.append(ToolCall(id=tc.id, function=tc)) - history_tool_call_cnt += 1 + elif is_named_tool_choice or is_required_tool_choice: message = ChatMessage( role=role, reasoning=reasoning, content=content or "", - tool_calls=tool_call_items, - ) - - elif request.tool_choice and request.tool_choice == "required": - tool_call_items = [] - tool_calls = tool_calls or [] - for tool_call in tool_calls: - if not tool_call.id: - tool_call.id = make_tool_call_id( - id_type=self.tool_call_id_type, - func_name=tool_call.name, - idx=history_tool_call_cnt, - ) - tool_call_items.append( - ToolCall(id=tool_call.id, function=tool_call) - ) - history_tool_call_cnt += 1 - message = ChatMessage( - role=role, - content=content or "", - tool_calls=tool_call_items, - reasoning=reasoning, + tool_calls=[ + ToolCall(id=tc.id or make_tool_call_id(), function=tc) + for tc in (tool_calls or []) + ], ) # if the request doesn't use tool choice @@ -949,21 +905,14 @@ class OpenAIServingChat(OpenAIServing): ): auto_tools_called = tool_calls is not None and len(tool_calls) > 0 if tool_calls: - tool_call_items = [] - for tc in tool_calls: - if not tc.id: - tc.id = make_tool_call_id( - id_type=self.tool_call_id_type, - func_name=tc.name, - idx=history_tool_call_cnt, - ) - tool_call_items.append(ToolCall(id=tc.id, function=tc)) - history_tool_call_cnt += 1 message = ChatMessage( role=role, reasoning=reasoning, content=content, - tool_calls=tool_call_items, + tool_calls=[ + ToolCall(id=tc.id or make_tool_call_id(), function=tc) + for tc in tool_calls + ], ) else: diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index 6b987f449d9..3c9a31a141e 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -301,7 +301,6 @@ class ParsableContext(ConversationContext): chat_template_content_format: ChatTemplateContentFormatOption, response_parser: Parser | None = None, enable_auto_tools: bool = False, - tool_call_id_type: str = "random", ): self.num_prompt_tokens = 0 self.num_output_tokens = 0 @@ -314,20 +313,8 @@ class ParsableContext(ConversationContext): self.num_init_messages = len(response_messages) self.finish_reason: str | None = None self.enable_auto_tools = enable_auto_tools - self.tool_call_id_type = tool_call_id_type self.response_parser = response_parser - if self.response_parser is None and parser_cls is not None: - chat_template_kwargs = request.build_chat_params( - default_template=chat_template, - default_template_content_format=chat_template_content_format, - ).chat_template_kwargs - self.response_parser = parser_cls( - tokenizer, - tools=request.tools, - chat_template_kwargs=chat_template_kwargs, - ) - self.parser_cls = parser_cls self.request = request @@ -365,7 +352,6 @@ class ParsableContext(ConversationContext): reasoning=reasoning, content=content, tool_calls=tool_calls, - tool_call_id_type=self.tool_call_id_type, ) ) elif completion.text: diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 62af1953dd0..f2e1f8e5d80 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -30,7 +30,6 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.chat_utils import ( ChatCompletionMessageParam, ChatTemplateContentFormatOption, - get_tool_call_id_type, ) from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.entrypoints.openai.engine.protocol import ( @@ -222,9 +221,6 @@ class OpenAIServingResponses(OpenAIServing): "For gpt-oss, we ignore --enable-auto-tool-choice " "and always enable tool use." ) - - self.tool_call_id_type = get_tool_call_id_type(self.model_config) - self.enable_auto_tools = enable_auto_tools # HACK(woosuk): This is a hack. We should use a better store. # FIXME: If enable_store=True, this may cause a memory leak since we @@ -272,6 +268,7 @@ class OpenAIServingResponses(OpenAIServing): tokenizer, request.tools, chat_template_kwargs=chat_template_kwargs, + model_config=self.model_config, ) def _validate_generator_input( @@ -493,7 +490,6 @@ class OpenAIServingResponses(OpenAIServing): chat_template=self.chat_template, chat_template_content_format=self.chat_template_content_format, enable_auto_tools=self.enable_auto_tools, - tool_call_id_type=self.tool_call_id_type, ) else: context = SimpleContext( @@ -1073,7 +1069,6 @@ class OpenAIServingResponses(OpenAIServing): content=content, tool_calls=tool_calls, logprobs=logprobs, - tool_call_id_type=self.tool_call_id_type, ) # Fallback when no parser is configured diff --git a/vllm/entrypoints/openai/responses/utils.py b/vllm/entrypoints/openai/responses/utils.py index 15b6fa88abc..a2f35dca235 100644 --- a/vllm/entrypoints/openai/responses/utils.py +++ b/vllm/entrypoints/openai/responses/utils.py @@ -45,7 +45,6 @@ def build_response_output_items( content: str | None, tool_calls: list[FunctionCall] | None, logprobs: list[Logprob] | None = None, - tool_call_id_type: str = "random", ) -> list[ResponseOutputItem]: outputs: list[ResponseOutputItem] = [] @@ -86,12 +85,7 @@ def build_response_output_items( ResponseFunctionToolCall( id=f"fc_{random_uuid()}", call_id=tool_call.id - if tool_call.id - else make_tool_call_id( - id_type=tool_call_id_type, - func_name=tool_call.name, - idx=idx, - ), + or make_tool_call_id(func_name=tool_call.name, idx=idx), type="function_call", status="completed", name=tool_call.name, diff --git a/vllm/entrypoints/pooling/embed/serving.py b/vllm/entrypoints/pooling/embed/serving.py index 5d9616f00c0..fd8140982ec 100644 --- a/vllm/entrypoints/pooling/embed/serving.py +++ b/vllm/entrypoints/pooling/embed/serving.py @@ -188,7 +188,12 @@ class ServingEmbedding(PoolingServing): ] total_tokens = get_pooling_usage(ctx.final_res_batch).prompt_tokens - image_tokens = total_tokens if request.images is not None else 0 + has_image_input = request.images is not None or any( + content.type == "image_url" + for input_item in request.inputs or [] + for content in input_item.content + ) + image_tokens = total_tokens if has_image_input else 0 texts_echo = request.texts embedding_types = request.embedding_types or ["float"] diff --git a/vllm/envs.py b/vllm/envs.py index 9cfd4792e14..08314a8c88d 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -200,6 +200,7 @@ if TYPE_CHECKING: VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600 VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998 VLLM_MOONCAKE_STORE_TIER_LOG: bool = False + VLLM_MOONCAKE_LOAD_RECV_THREADS: int = 1 VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO: float = 0.9 MOONCAKE_PREFERRED_SEGMENT: str | None = None MOONCAKE_REQUESTER_LOCAL_HOSTNAME: str | None = None @@ -210,11 +211,10 @@ if TYPE_CHECKING: VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5 VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None - VLLM_USE_PACKED_HMA_KV_CACHE: bool = False VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None VLLM_COMPUTE_NANS_IN_LOGITS: bool = False VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[ - "FP", "INT8", "INT6", "INT4", "NONE" + "FP", "INT8", "INT6", "INT4", "INT3", "NONE" ] = "NONE" VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None @@ -1224,12 +1224,12 @@ environment_variables: dict[str, Callable[[], Any]] = { os.getenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "False").lower() in ("true", "1") ), # Custom quick allreduce kernel for MI3* cards - # Choice of quantization level: FP, INT8, INT6, INT4 or NONE + # Choice of quantization level: FP, INT8, INT6, INT4, INT3 or NONE # Recommended for large models to get allreduce "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION": env_with_choices( "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE", - ["FP", "INT8", "INT6", "INT4", "NONE"], + ["FP", "INT8", "INT6", "INT4", "INT3", "NONE"], ), # Custom quick allreduce kernel for MI3* cards # Due to the lack of the bfloat16 asm instruction, bfloat16 @@ -1532,6 +1532,14 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_MOONCAKE_STORE_TIER_LOG": lambda: ( os.getenv("VLLM_MOONCAKE_STORE_TIER_LOG", "False").lower() in ("true", "1") ), + # Number of parallel KV-load receive threads per worker rank. Lets the + # per-request control overhead (Python prep + master key lookup) of one + # request overlap with the RDMA transfer of another, keeping the transfer + # engine's queue pairs busy. Helps when that overhead is significant or + # per-request batches are too small to saturate the link on their own. + "VLLM_MOONCAKE_LOAD_RECV_THREADS": lambda: int( + os.getenv("VLLM_MOONCAKE_LOAD_RECV_THREADS", "1") + ), # Fraction of the owner's DirectIO staging buffer to fill per GET batch. "VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO": lambda: float( os.getenv("VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO", "0.9") @@ -1614,11 +1622,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_KV_CACHE_LAYOUT": env_with_choices( "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"] ), - # Opt into packed per-block KV cache allocation for multi-group - # attention-only HMA models (e.g. gpt-oss, Gemma 3/4). - "VLLM_USE_PACKED_HMA_KV_CACHE": lambda: bool( - int(os.getenv("VLLM_USE_PACKED_HMA_KV_CACHE", "0")) - ), # SSM conv state layout used for Mamba models. # - SD: (state_len, dim) — dim contiguous (default) # - DS: (dim, state_len) — TP-sharded dim on dim1, diff --git a/vllm/kernels/helion/utils.py b/vllm/kernels/helion/utils.py index 460fcc85065..f4ace6cff3b 100644 --- a/vllm/kernels/helion/utils.py +++ b/vllm/kernels/helion/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for Helion kernel management.""" +import regex as re import torch from vllm.logger import init_logger @@ -62,7 +63,7 @@ def canonicalize_gpu_name(name: str) -> str: """ Canonicalize GPU name for use as a platform identifier. - Converts to lowercase, replaces spaces and hyphens with underscores, + Converts to lowercase, replaces separators with underscores, and maps known variant names to their canonical form via _GPU_NAME_ALIASES. e.g., "NVIDIA H100 80GB HBM3" -> "nvidia_h100" "NVIDIA A100-SXM4-80GB" -> "nvidia_a100" @@ -70,9 +71,7 @@ def canonicalize_gpu_name(name: str) -> str: """ if not name or not name.strip(): raise ValueError("GPU name cannot be empty") - name = name.lower() - name = name.replace(" ", "_") - name = name.replace("-", "_") + name = re.sub(r"[\s/-]+", "_", name.lower()) if name in _GPU_NAME_ALIASES: return _GPU_NAME_ALIASES[name] return name diff --git a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py index 13012015069..c2627668cf6 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py @@ -168,11 +168,13 @@ class CPUWNA16LinearKernel(MPLinearKernel): if zp.output_dim == 0: zp.data = zp.t().contiguous() + supports_amx = torch.cpu._is_amx_tile_supported() + supports_riscv = current_platform.get_cpu_architecture() == CpuArchEnum.RISCV layer.use_w4a8 = ( envs.VLLM_CPU_INT4_W4A8 and not self.config.has_g_idx and self.config.act_type == torch.bfloat16 - and torch.cpu._is_amx_tile_supported() + and (supports_amx or supports_riscv) ) # layer.use_w4a8 = False # AWQ format will be converted to GPTQ format in `AutoAWQMarlinLinearMethod` diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 051468ed14c..8d9a674319d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -800,9 +800,9 @@ class MLAAttention(nn.Module, AttentionLayerBase): else: mqa_q = (mqa_ql_nope, mqa_q_pe) if self.impl.dcp_world_size > 1: - assert not fp8_attention, "DCP not support fp8 kvcache now." - # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) - mqa_q = torch.cat(mqa_q, dim=-1) + if isinstance(mqa_q, tuple): + # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) + mqa_q = torch.cat(mqa_q, dim=-1) # mqa_q do allgather in head dim. mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) @@ -1263,6 +1263,7 @@ class MLACommonPrefillMetadata: padded_local_chunk_seq_lens: list[list[int]] | None = None local_context_lens_allranks: list[list[int]] | None = None padded_local_cu_seq_lens: torch.Tensor | None = None + padded_local_token_to_seq: torch.Tensor | None = None cu_seq_lens_lst: list[list[int]] | None = None chunk_size: int | None = None prefill_tokens_with_context: int | None = None @@ -1531,9 +1532,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size - # Don't try to access the runner on AMD - if self.aot_schedule: - self.page_size = self.kv_cache_spec.block_size + self.page_size = self.kv_cache_spec.block_size self.chunked_prefill_workspace_size = ( self.determine_chunked_prefill_workspace_size(vllm_config) @@ -1683,12 +1682,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): self.chunked_prefill_workspace_size // num_prefills_with_context_cpu ) - if self.aot_schedule: - # align max_context_chunk to page_size by rounding down, - # currently the `gather_and_maybe_dequant_cache` kernel - # cannot handle `context_chunk_starts` that are not aligned - # to page_size - max_context_chunk = round_down(max_context_chunk, self.page_size) + # align max_context_chunk to page_size by rounding down, + # currently the `gather_and_maybe_dequant_cache` kernel + # cannot handle `context_chunk_starts` that are not aligned + # to page_size + max_context_chunk = round_down(max_context_chunk, self.page_size) assert max_context_chunk > 0 num_chunks = cdiv(max_context_len_cpu, max_context_chunk) @@ -1787,6 +1785,21 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): out=padded_local_cu_chunk_seq_lens_cpu[:, 1:], dtype=torch.int32, ) + max_padded_local_tokens_over_chunk = ( + padded_local_cu_chunk_seq_lens_cpu[:, -1].max().item() + ) + padded_local_token_to_seq_tensor_cpu = torch.zeros( + [num_chunks, max_padded_local_tokens_over_chunk], + dtype=torch.int32, + ) + for i in range(num_chunks): + chunk_token_to_seq_tensor = torch.repeat_interleave( + range_idx, padded_local_chunk_seq_lens[i] + ) + chunk_len = chunk_token_to_seq_tensor.shape[0] + padded_local_token_to_seq_tensor_cpu[i, :chunk_len] = ( + chunk_token_to_seq_tensor + ) prefill_tokens_with_context = None if num_prefills_with_context_cpu > 0: @@ -1811,6 +1824,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): padded_local_cu_seq_lens=padded_local_cu_chunk_seq_lens_cpu.to( device, non_blocking=True ), + padded_local_token_to_seq=padded_local_token_to_seq_tensor_cpu.to( + device, non_blocking=True + ), cu_seq_lens_lst=cu_seq_lens_cpu.tolist(), chunk_size=padded_local_max_context_chunk_across_ranks, prefill_tokens_with_context=prefill_tokens_with_context, @@ -2187,7 +2203,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): k_scale: torch.Tensor, dcp_world_size: int, ): - assert k_scale is None, "DCP not support scaled kvcache now." assert attn_metadata.prefill is not None prefill_metadata = attn_metadata.prefill assert prefill_metadata.prefill_backend is not None @@ -2195,9 +2210,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): assert prefill_metadata.chunked_context.padded_local_chunk_seq_lens is not None assert prefill_metadata.chunked_context.local_context_lens_allranks is not None assert prefill_metadata.chunked_context.padded_local_cu_seq_lens is not None + assert prefill_metadata.chunked_context.padded_local_token_to_seq is not None assert prefill_metadata.chunked_context.cu_seq_lens_lst is not None assert prefill_metadata.chunked_context.chunk_size is not None + use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype() output = None merge_output = None iters = len(prefill_metadata.chunked_context.seq_tot) @@ -2205,16 +2222,37 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): for i in range(iters): toks = prefill_metadata.chunked_context.seq_tot[i] - ops.cp_gather_cache( - src_cache=kv_c_and_k_pe_cache, - dst=workspace, - block_table=prefill_metadata.block_table, - cu_seq_lens=prefill_metadata.chunked_context.padded_local_cu_seq_lens[ - i - ], - batch_size=attn_metadata.num_prefills, - seq_starts=prefill_metadata.chunked_context.starts[i], + if toks == 0: + continue + padded_local_cu_seq_lens = ( + prefill_metadata.chunked_context.padded_local_cu_seq_lens[i] ) + if is_quantized_kv_cache(self.kv_cache_dtype) and ( + self.kv_cache_dtype != "fp8_ds_mla" + ): + assert k_scale is not None + ops.gather_and_maybe_dequant_cache( + src_cache=kv_c_and_k_pe_cache, + dst=workspace, + block_table=prefill_metadata.block_table, + cu_seq_lens=padded_local_cu_seq_lens, + token_to_seq=prefill_metadata.chunked_context.padded_local_token_to_seq[ + i + ], + num_tokens=toks, + kv_cache_dtype=self.kv_cache_dtype, + scale=k_scale, + seq_starts=prefill_metadata.chunked_context.starts[i], + ) + else: + ops.cp_gather_cache( + src_cache=kv_c_and_k_pe_cache, + dst=workspace, + block_table=prefill_metadata.block_table, + cu_seq_lens=padded_local_cu_seq_lens, + batch_size=attn_metadata.num_prefills, + seq_starts=prefill_metadata.chunked_context.starts[i], + ) # workspace # |------- N tokens --------|--------- N*dcp_size tokens ----------| # |<- use for local_gather ->|<--------- use for allgather -------->| @@ -2252,9 +2290,22 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): toks=toks, ) + kv_b_proj_w_dtype = ( + self.kv_b_proj.weight.dtype + if hasattr(self.kv_b_proj, "weight") + else self.kv_b_proj.params_dtype + ) + if ( + use_fp8_prefill or kv_b_proj_w_dtype != current_platform.fp8_dtype() + ) and kv_b_proj_w_dtype != torch.uint8: + kv_c_normed = kv_c_normed.to(kv_b_proj_w_dtype) + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim ) + if use_fp8_prefill: + kv_nope = kv_nope.to(prefill_metadata.q_data_type) + k_pe = k_pe.to(prefill_metadata.q_data_type) k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) k = self._concat_k_nope_k_pe(k_nope, k_pe) @@ -2346,7 +2397,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]): q, kv_c_and_k_pe_cache, attn_metadata, - k_scale=None, + k_scale=k_scale, dcp_world_size=self.dcp_world_size, ) ) diff --git a/vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json index 3357dc223f7..6e47aa02383 100644 --- a/vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json +++ b/vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json @@ -2,35 +2,35 @@ "triton_version": "3.5.0", "1": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 4 }, "2": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 4 }, "4": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 64, + "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 4 + "num_stages": 3 }, "8": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 64, + "GROUP_SIZE_M": 16, "num_warps": 4, - "num_stages": 4 + "num_stages": 3 }, "16": { "BLOCK_SIZE_M": 16, diff --git a/vllm/model_executor/layers/fused_moe/configs/E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json index 9c07695ba91..04aea87b055 100644 --- a/vllm/model_executor/layers/fused_moe/configs/E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json +++ b/vllm/model_executor/layers/fused_moe/configs/E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json @@ -1,17 +1,17 @@ { "1": { - "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 3 + "num_stages": 5 }, "2": { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 32, + "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 4 }, @@ -25,11 +25,11 @@ }, "8": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 16, + "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 5 + "num_stages": 3 }, "16": { "BLOCK_SIZE_M": 16, @@ -143,4 +143,4 @@ "num_warps": 4, "num_stages": 4 } -} \ No newline at end of file +} diff --git a/vllm/model_executor/layers/fused_moe/configs/E=160,N=640,device_name=NVIDIA_H100,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=160,N=640,device_name=NVIDIA_H100,dtype=fp8_w8a8,block_shape=[128,128].json index aa7610cd75e..b2016140b63 100644 --- a/vllm/model_executor/layers/fused_moe/configs/E=160,N=640,device_name=NVIDIA_H100,dtype=fp8_w8a8,block_shape=[128,128].json +++ b/vllm/model_executor/layers/fused_moe/configs/E=160,N=640,device_name=NVIDIA_H100,dtype=fp8_w8a8,block_shape=[128,128].json @@ -1,7 +1,7 @@ { "1": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, @@ -9,11 +9,11 @@ }, "2": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 3 + "num_stages": 4 }, "4": { "BLOCK_SIZE_M": 16, @@ -24,12 +24,12 @@ "num_stages": 3 }, "8": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 16, + "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 5 + "num_stages": 3 }, "16": { "BLOCK_SIZE_M": 16, diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json index 532c16e8992..4f8ec02b9ff 100644 --- a/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json @@ -1,35 +1,35 @@ { "1": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 16, + "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 4 }, "2": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 16, + "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 4 + "num_stages": 5 }, "4": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 64, + "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 4 }, "8": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 16, + "GROUP_SIZE_M": 32, "num_warps": 4, - "num_stages": 4 + "num_stages": 3 }, "16": { "BLOCK_SIZE_M": 16, diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json index 381eb5d826a..5c8b76d873c 100644 --- a/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json @@ -2,19 +2,19 @@ "triton_version": "3.5.0", "1": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 4 + "num_stages": 5 }, "2": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 4 + "num_stages": 5 }, "4": { "BLOCK_SIZE_M": 16, @@ -34,11 +34,11 @@ }, "16": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 5 + "num_stages": 3 }, "24": { "BLOCK_SIZE_M": 16, diff --git a/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json index 689e553e1c2..690e6190032 100644 --- a/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json +++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json @@ -1,21 +1,21 @@ { "triton_version": "3.6.0", "1": { + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "2": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 3 }, - "2": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 16, - "num_warps": 8, - "num_stages": 4 - }, "4": { "BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json index ac53df14ce8..1907fda73ed 100644 --- a/vllm/model_executor/layers/fused_moe/configs/E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json @@ -1,23 +1,23 @@ { "1": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 3 }, "2": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 4 + "num_stages": 5 }, "4": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, @@ -25,19 +25,19 @@ }, "8": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, + "GROUP_SIZE_M": 32, "num_warps": 4, - "num_stages": 4 + "num_stages": 3 }, "16": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, + "GROUP_SIZE_M": 16, "num_warps": 4, - "num_stages": 4 + "num_stages": 3 }, "24": { "BLOCK_SIZE_M": 16, diff --git a/vllm/model_executor/layers/fused_moe/configs/E=512,N=256,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json b/vllm/model_executor/layers/fused_moe/configs/E=512,N=256,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json index b0bf1bf5178..c532364e811 100644 --- a/vllm/model_executor/layers/fused_moe/configs/E=512,N=256,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json +++ b/vllm/model_executor/layers/fused_moe/configs/E=512,N=256,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json @@ -2,27 +2,27 @@ "triton_version": "3.4.0", "1": { "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 4 + "num_stages": 5 }, "2": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, "num_stages": 4 }, "4": { - "BLOCK_SIZE_M": 16, - "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_M": 8, + "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128, "GROUP_SIZE_M": 1, "num_warps": 4, - "num_stages": 3 + "num_stages": 4 }, "8": { "BLOCK_SIZE_M": 16, @@ -36,9 +36,9 @@ "BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, + "GROUP_SIZE_M": 32, "num_warps": 4, - "num_stages": 4 + "num_stages": 3 }, "24": { "BLOCK_SIZE_M": 16, diff --git a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py index 31ef144e237..4b0a0b8ecad 100644 --- a/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py @@ -4,7 +4,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -220,6 +219,241 @@ def _patch_make_bitmatrix_metadata() -> None: _bm.make_bitmatrix_metadata = _make_bitmatrix_metadata_pow2_safe +def _patch_legacy_routing_for_nonpow2_topk() -> None: + """Monkey-patch the legacy (v3.5.1) triton_kernels routing path to support + non-power-of-2 top_k (e.g. DeepSeek-V4 top_k=6). + + The bundled ``_routing_compute_indx`` does ``tl.arange(0, N_EXPTS_ACT * + BLOCK_M)``, which fails to compile when ``N_EXPTS_ACT`` (top_k) is not a + power of 2 (6 * 32 = 192). This installs a pow2-safe variant that pads the + ``tl.arange`` to the next power of 2, strides by the real per-block size, + and masks the padded tail so it neither loads the next block's gates nor + writes any output. For power-of-2 top_k it is identical to the original. + + A matching ``sort_tokens`` is installed that threads the padded size into + the patched kernel. Only needed on the legacy path; the v3.6+ SparseMatrix + path is handled by ``_patch_make_bitmatrix_metadata``. + """ + import triton + import triton.language as tl + + # Import via the `triton_kernels` alias (set up by has_triton_kernels) so + # we patch the SAME module object that `make_routing_data` consumes. The + # `vllm.third_party.triton_kernels.routing` path is a *different* module + # object under the import alias, so patching it would have no effect. + try: + import triton_kernels.routing as _routing + from triton_kernels.routing_details import _routing_compute as _rc + except ImportError: + return + + _keyed_add = _rc._keyed_add + _expt_data_compute = _rc._expt_data_compute + + @triton.jit + def _routing_compute_indx_pow2( + pid_m, + GatherIndx, + ScatterIndx, + GateScal, + ExptScal, + ExptIndx, + PartialOffs, + stride_pm, + stride_pn, + TokensStart, + n_tokens, + BLOCK_M: tl.constexpr, + N_EXPTS_ACT: tl.constexpr, + BLOCK_SIZE_PADDED: tl.constexpr, + ): + if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr(): + n_tokens = tl.load(n_tokens) + n_gates = n_tokens * N_EXPTS_ACT + BLOCK_SIZE: tl.constexpr = N_EXPTS_ACT * BLOCK_M + tl.static_assert(BLOCK_SIZE_PADDED <= 32768) + local_offs = tl.arange(0, BLOCK_SIZE_PADDED) + offs = pid_m * BLOCK_SIZE + local_offs + expert = tl.load( + ExptIndx + offs, + mask=(local_offs < BLOCK_SIZE) & (offs < n_gates), + other=-1, + ).to(tl.uint32) + kv_pairs = ((expert << 16) | local_offs).to(tl.uint32) + kv_pairs = tl.sort(kv_pairs, 0) + expert = kv_pairs >> 16 + offs = pid_m * BLOCK_SIZE + (kv_pairs & 0xFFFF) + mask = expert != 0xFFFF + gate_scal = tl.load(ExptScal + offs, mask=mask) + x = kv_pairs & 0xFFFF0000 | 0x00000001 + run_lengths = tl.associative_scan(x, 0, _keyed_add) + exclusive_run_lengths = (run_lengths - 1) & 0xFFFF + gates = tl.load(PartialOffs + pid_m * stride_pm + expert * stride_pn, mask=mask) + gates += tl.load(TokensStart + expert, mask=mask) + gates += exclusive_run_lengths + tl.store(ScatterIndx + offs, gates, mask=mask) + tl.store(GatherIndx + gates, offs, mask=mask) + tl.store(GateScal + gates, gate_scal, mask=mask) + + @triton.jit + def _combined_routing_compute_pow2( + GatherIndx, + ScatterIndx, + GateScal, + ExptScal, + ExptIndx, + PartialOffs, + stride_pm, + stride_pn, + TokensStart, + n_tokens, + BLOCK_M: tl.constexpr, + N_EXPTS_ACT: tl.constexpr, + Hist, + MDTileStarts, + tile_starts_stridem, + MDTileInfo, + tile_info_stridem, + first_tile_dim_log2, + SIZES: tl.constexpr, + BLOCK: tl.constexpr, + blocks2a, + BLOCK_SIZE_PADDED: tl.constexpr, + ): + pid = tl.program_id(0) + if pid < blocks2a: + _expt_data_compute( + Hist, + MDTileStarts, + tile_starts_stridem, + MDTileInfo, + tile_info_stridem, + first_tile_dim_log2, + SIZES, + BLOCK, + ) + else: + pid -= blocks2a + _routing_compute_indx_pow2( + pid, + GatherIndx, + ScatterIndx, + GateScal, + ExptScal, + ExptIndx, + PartialOffs, + stride_pm, + stride_pn, + TokensStart, + n_tokens, + BLOCK_M, + N_EXPTS_ACT, + BLOCK_SIZE_PADDED, + ) + + def _sort_tokens_pow2(expt_scal, expt_indx, n_expts_tot, bitmatrix): + import torch + + HIST_BLOCK_M = 32 + INDX_OFFS_BLOCK_M = 512 + MEMSET_BLOCK = 1024 + cdiv = triton.cdiv + device = expt_scal.device + dtype = expt_scal.dtype + n_tokens_raw, _ = bitmatrix.shape + n_tokens_pad, n_expts_act = expt_scal.shape + n_gates_pad = n_tokens_pad * n_expts_act + # pad per-block gate count (HIST_BLOCK_M * top_k) up to a pow2. + block_size_padded = triton.next_power_of_2(HIST_BLOCK_M * n_expts_act) + + hist, partial_hist = bitmatrix.sum(partials_block_size=HIST_BLOCK_M) + hist = hist[:n_expts_tot] + expt_offs = torch.empty(n_expts_tot, dtype=torch.int32, device=device) + combined_indx = torch.empty(n_gates_pad * 2, dtype=torch.int32, device=device) + topk_indx = combined_indx[:n_gates_pad] + gate_indx = combined_indx[n_gates_pad:] + gate_scal = torch.empty(n_gates_pad, dtype=dtype, device=device) + + ( + token_offs_combined, + token_offs_raw, + token_offs_pad, + block_pid_map, + blocks1a, + blocks2a, + MEMSET_BLOCK_A, + HIST2_BLOCK_M, + block_m_log2_start, + block_m_num, + ) = _routing._compute_expt_data_internal(hist, n_expts_tot, n_gates_pad) + + blocks1b = cdiv(n_gates_pad * 2, MEMSET_BLOCK) + n_expts_tot + 1 + blocks2b = cdiv(n_tokens_pad, HIST_BLOCK_M) + + _rc._combined_routing_memset[(blocks1a + blocks1b,)]( + combined_indx, + n_gates_pad * 2, + -1, + MEMSET_BLOCK, + hist, + expt_offs, + hist.shape[0], + n_expts_tot, + partial_hist, + partial_hist.shape[0], + partial_hist.stride(0), + partial_hist.stride(1), + token_offs_combined, + token_offs_combined.stride(0), + blocks1a, + block_pid_map, + block_m_log2_start, + SIZES=block_m_num, + BLOCK_A=MEMSET_BLOCK_A, + BLOCK_N=512, + BLOCK_M=INDX_OFFS_BLOCK_M, + ) + + indx_offs = partial_hist + _combined_routing_compute_pow2[(blocks2a + blocks2b,)]( + topk_indx, + gate_indx, + gate_scal, + expt_scal, + expt_indx, + indx_offs, + indx_offs.stride(0), + indx_offs.stride(1), + expt_offs, + n_tokens_raw, + HIST_BLOCK_M, + n_expts_act, + hist, + token_offs_pad, + token_offs_pad.stride(0), + block_pid_map, + block_pid_map.stride(0), + block_m_log2_start, + block_m_num, + HIST2_BLOCK_M, + blocks2a, + block_size_padded, + ) + return ( + hist, + topk_indx, + gate_indx, + gate_scal, + token_offs_raw, + token_offs_pad, + block_pid_map, + ) + + # `routing_from_bitmatrix` looks up `sort_tokens` via the routing module + # global, so replacing it here redirects the legacy path to the pow2 kernel. + _routing.sort_tokens = _sort_tokens_pow2 + + # Two API generations of triton_kernels are supported: # - v3.5.1 (the version bundled with vLLM): exposes `routing()` and # `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix` @@ -260,6 +494,9 @@ if has_triton_kernels(): use_legacy_triton_kernels = True if not use_legacy_triton_kernels: _patch_make_bitmatrix_metadata() + else: + # Legacy routing fails to compile for non-pow2 top_k (DeepSeek-V4). + _patch_legacy_routing_for_nonpow2_topk() except (AttributeError, ImportError) as e: logger.error( "Failed to import Triton kernels. Please make sure your triton " @@ -577,6 +814,85 @@ def make_routing_data( return routing_data, gather_indx, scatter_indx +@triton.jit +def _masked_topk_sum_kernel( + inp_ptr, # (M, topk, K) contiguous + topk_ids_ptr, # (M, topk) int: -1 marks an invalid / non-local slot + out_ptr, # (M, K), same dtype as inp + K, + topk: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_m = tl.program_id(0).to(tl.int64) + k = tl.program_id(1) * BLOCK_K + tl.arange(0, BLOCK_K) + k_mask = k < K + base = pid_m * topk + acc = tl.zeros((BLOCK_K,), dtype=tl.float32) + for j in tl.static_range(topk): + eid = tl.load(topk_ids_ptr + base + j) + # NOTE: This is NaN-safe because the invalid slots are skipped. + if eid >= 0: + x = tl.load(inp_ptr + (base + j) * K + k, mask=k_mask) + acc += x.to(tl.float32) + tl.store(out_ptr + pid_m * K + k, acc.to(out_ptr.dtype.element_ty), mask=k_mask) + + +def masked_moe_sum( + intermediate: torch.Tensor, # (M, topk, K) + topk_ids: torch.Tensor, # (M, topk) int, -1 = invalid / non-local slot + output: torch.Tensor, # (M, K) +) -> None: + M, topk, K = intermediate.shape + BLOCK_K = 1024 + grid = (M, triton.cdiv(K, BLOCK_K)) + _masked_topk_sum_kernel[grid]( + intermediate, topk_ids, output, K, topk=topk, BLOCK_K=BLOCK_K + ) + + +@triton.jit +def _remap_topk_to_local_kernel( + topk_ids_ptr, # [n] global expert IDs (-1 = invalid) + expert_map_ptr, # [num_experts] global->local (-1 for non-local) + out_ptr, # [n] int64 local expert IDs (-1 for invalid/non-local) + n_elements, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + tid = tl.load(topk_ids_ptr + offs, mask=mask, other=-1) + # Gather expert_map[tid] for valid (tid >= 0); clamp the index so invalid + # rows don't read OOB, then select -1 for them. Matches + # torch.where(tid >= 0, expert_map[clamp(tid, 0)], -1) -- preserving -1 (a + # plain expert_map[-1] would wrap to a valid local id and misroute). + valid = tid >= 0 + idx = tl.where(valid, tid, 0) + local = tl.load(expert_map_ptr + idx, mask=mask, other=-1) + out = tl.where(valid, local.to(tl.int64), -1) + tl.store(out_ptr + offs, out, mask=mask) + + +def remap_topk_to_local( + topk_ids: torch.Tensor, expert_map: torch.Tensor +) -> torch.Tensor: + """Fused global->local expert-id mapping over a topk_ids tensor, preserving -1. + + Replaces ``torch.where(topk_ids >= 0, expert_map[topk_ids.clamp(min=0)], -1)`` + with one kernel. Returns a NEW int64 tensor -- the caller keeps the original + ``topk_ids`` as ``global_topk_ids``, so this must not write in place. + + (Distinct from ``deep_gemm_utils.apply_expert_map``, which is a scalar + ``@triton.jit`` device helper called from within other kernels.) + """ + out = torch.empty_like(topk_ids, dtype=torch.int64) + n = topk_ids.numel() + BLOCK = 1024 + grid = (triton.cdiv(n, BLOCK),) + _remap_topk_to_local_kernel[grid](topk_ids, expert_map, out, n, BLOCK=BLOCK) + return out + + class BaseOAITritonExperts(mk.FusedMoEExpertsModular): @property def expects_unquantized_inputs(self) -> bool: @@ -708,7 +1024,9 @@ class OAITritonExperts(BaseOAITritonExperts): self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG if expert_map is not None: - topk_ids = expert_map[topk_ids] + # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch): + # make_routing_data treats -1 as the skip sentinel. + topk_ids = remap_topk_to_local(topk_ids, expert_map) local_num_experts = w1.shape[0] if global_num_experts == -1: @@ -780,9 +1098,6 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): output = (M, K) return (workspace1, workspace2, output) - def moe_sum(self, input: torch.Tensor, output: torch.Tensor): - ops.moe_sum(input, output) - def activation( self, activation: MoEActivation, @@ -853,7 +1168,9 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): global_topk_ids = topk_ids if expert_map is not None: - topk_ids = expert_map[topk_ids] + # Preserve -1 (invalid / non-local slots, e.g. from EP dispatch): + # make_routing_data treats -1 as the skip sentinel. + topk_ids = remap_topk_to_local(topk_ids, expert_map) local_num_experts = w1.shape[0] if global_num_experts == -1: @@ -976,7 +1293,9 @@ class UnfusedOAITritonExperts(LoRAExpertsMixin, BaseOAITritonExperts): top_k_num=topk, ) - self.moe_sum(intermediate_cache3.view(-1, topk, K), output) + # matmul_ogs leaves invalid (-1 / non-local EP) slots unwritten. + # Reduce over topk skipping those slots. + masked_moe_sum(intermediate_cache3.view(-1, topk, K), topk_ids, output) class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py index a412a6936d3..c6e5e70a14a 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -68,8 +68,10 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - # FlashInfer MxInt4 uses a fused SwiGLU activation. - return activation == MoEActivation.SWIGLUOAI + # FlashInfer MxInt4 names the standard gated SiLU path "SwiGLU". + # In vLLM MoE configs that maps to SILU/silu_and_mul; SWIGLUOAI is + # kept as an alias for consistency with other FlashInfer backends. + return activation in (MoEActivation.SILU, MoEActivation.SWIGLUOAI) @staticmethod def _supports_parallel_config( diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 49957c8f5e3..269b6e3da0b 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, ) from vllm.model_executor.layers.fused_moe.utils import ( + enable_swap_ab, moe_kernel_quantize_input, ) from vllm.platforms import current_platform @@ -343,6 +344,7 @@ def fused_moe_kernel( use_int8_w8a16: tl.constexpr, per_channel_quant: tl.constexpr, HAS_BIAS: tl.constexpr, + SWAP_AB: tl.constexpr, ): """ Implements the fused computation for a Mixture of Experts (MOE) using @@ -432,15 +434,25 @@ def fused_moe_kernel( offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + ( - offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak - ) + if SWAP_AB: + a_ptrs = a_ptr + ( + offs_k[:, None] * stride_ak + offs_token[None, :] // top_k * stride_am + ) + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk) + ) + else: + a_ptrs = a_ptr + ( + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + ) + b_ptrs = ( + b_ptr + + off_experts * stride_be + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + ) - b_ptrs = ( - b_ptr - + off_experts * stride_be - + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - ) if use_int8_w8a16: b_scale_ptrs = ( b_scale_ptr + off_experts * stride_bse + offs_bn[None, :] * stride_bsn @@ -477,16 +489,25 @@ def fused_moe_kernel( # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block # of fp32 values for higher accuracy. # `accumulator` will be converted back to fp16 after the loop. - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + if SWAP_AB: + accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32) + else: + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): # Load the next block of A and B, generate a mask by checking the # K dimension. + if SWAP_AB: + a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :] + b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K + else: + a_mask = token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K) + b_mask = offs_k[:, None] < K - k * BLOCK_SIZE_K a = tl.load( a_ptrs, - mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + mask=a_mask, other=0.0, ) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=b_mask, other=0.0) # We accumulate along the K dimension. if use_int8_w8a16: accumulator = tl.dot(a, b.to(compute_type), acc=accumulator) @@ -498,12 +519,17 @@ def fused_moe_kernel( a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0 ) b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) - - accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] + if SWAP_AB: + accumulator += tl.dot(b, a) * b_scale[:, None] * a_scale[None, :] + else: + accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] else: if use_fp8_w8a8: # acc used to enable fp8_fast_accum - accumulator = tl.dot(a, b, acc=accumulator) + if SWAP_AB: + accumulator = tl.dot(b, a, acc=accumulator) + else: + accumulator = tl.dot(a, b, acc=accumulator) else: accumulator += tl.dot(a, b) else: @@ -512,6 +538,9 @@ def fused_moe_kernel( a_ptrs += BLOCK_SIZE_K * stride_ak b_ptrs += BLOCK_SIZE_K * stride_bk + if SWAP_AB: + accumulator = tl.trans(accumulator, (1, 0)) + # Dequantization for supported quantization schemes: # - int8_w8a16 # - fp8_w8a8 @@ -729,6 +758,11 @@ def invoke_fused_moe_triton_kernel( assert topk_weights is None or topk_weights.stride(1) == 1 assert sorted_token_ids is None or sorted_token_ids.stride(0) == 1 + if use_fp8_w8a8: + SWAP_AB = enable_swap_ab(config["BLOCK_SIZE_M"], config["BLOCK_SIZE_N"]) + else: + SWAP_AB = False + if use_fp8_w8a8 or use_int8_w8a8: assert B_scale is not None assert block_shape is None or triton.cdiv( @@ -810,6 +844,7 @@ def invoke_fused_moe_triton_kernel( naive_block_assignment=(sorted_token_ids is None), HAS_BIAS=HAS_BIAS, BLOCK_SIZE_K=BLOCK_SIZE_K, + SWAP_AB=SWAP_AB, **config, ) @@ -1222,19 +1257,38 @@ def get_default_config( num_stages_rocm = 2 if dtype == "fp8_w8a8" and block_shape is not None: - # Block-wise quant: tile sizes are constrained by block_shape. - # Use a small M tile for decode-like batches where tokens are - # spread thin across experts. Larger batches benefit from - # GROUP_SIZE_M > 1 because the per-block scales add memory - # traffic that benefits from L2 tile reuse. + # Block-wise quant. Use a small M tile for decode-like batches where + # tokens are spread thin across experts. Larger batches benefit from + # GROUP_SIZE_M > 1 because the per-block scales add memory traffic + # that benefits from L2 tile reuse. + # + # BLOCK_SIZE_N need not equal block_shape[0]: the kernel indexes block + # scales per N element (offs_bn // group_n), so any N tile dividing the + # quant block is valid. At decode a 128-wide N tile leaves the gate-up + # GEMM SM-bound; a 64-wide tile exposes ~2x the thread blocks, and the + # swap-AB kernel keeps it efficient on Hopper down to the smallest + # batches, so prefer N=64 through low batch sizes. CUDA only (validated + # on NVIDIA); ROCm keeps its prior tile/pipeline sizes. + if current_platform.is_rocm(): + block_n = block_shape[0] + num_stages = num_stages_rocm + elif M <= 8 and block_shape[0] % 64 == 0: + block_n = 64 + # The smallest batches are memory-latency bound, so a deeper + # pipeline hides the weight loads; by M=8 it turns occupancy/SMEM + # bound and the extra stages hurt. + num_stages = 4 if M <= 4 else 3 + else: + block_n = block_shape[0] + num_stages = 3 config = { "BLOCK_SIZE_M": 16 if M <= 64 else 64, - "BLOCK_SIZE_N": block_shape[0], + "BLOCK_SIZE_N": block_n, "BLOCK_SIZE_K": block_shape[1], "GROUP_SIZE_M": 1 if M <= 16 else 32, "SPLIT_K": 1, "num_warps": 4, - "num_stages": 3 if not current_platform.is_rocm() else num_stages_rocm, + "num_stages": num_stages, } elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None: # moe wna16 kernels diff --git a/vllm/model_executor/layers/fused_moe/oracle/__init__.py b/vllm/model_executor/layers/fused_moe/oracle/__init__.py index 208f01a7cb5..f1942819cd5 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/__init__.py +++ b/vllm/model_executor/layers/fused_moe/oracle/__init__.py @@ -1,2 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.model_executor.layers.fused_moe.oracle.base import MoEKernelOracle +from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoEKernelOracle, +) + +__all__ = [ + "MoEKernelOracle", + "UnquantizedMoEKernelOracle", +] diff --git a/vllm/model_executor/layers/fused_moe/oracle/base.py b/vllm/model_executor/layers/fused_moe/oracle/base.py new file mode 100644 index 00000000000..6f7b8a98558 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/oracle/base.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Abstract base class for MoE kernel oracles. + +Each MoE oracle (unquantized / fp8 / nvfp4 / mxfp4 / mxfp8 / int8 / +int_wna16) is responsible for selecting the right MoE kernel backend for a +given (model, hardware, deployment-config) tuple. The current +implementation expresses this responsibility as module-level functions +that follow an informal convention. + +This module declares the abstract contract; concrete oracles inherit from +`MoEKernelOracle` and provide the platform-specific behaviour. + +This is the first PR in the series suggested by @robertgshaw2-redhat in +PR #37776 (see issue #37753). It intentionally only introduces the ABC; +follow-up PRs migrate each oracle to inherit from it. The single concrete +subclass shipped here (`UnquantizedMoEKernelOracle`) delegates to the +existing module-level functions to keep behaviour bit-identical with +pre-class code. +""" + +from abc import ABC, abstractmethod +from enum import Enum +from typing import TYPE_CHECKING, Generic, TypeVar + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config.kernel import MoEBackend +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, +) + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + +BackendT = TypeVar("BackendT", bound=Enum) + + +class MoEKernelOracle(ABC, Generic[BackendT]): + """Abstract base for MoE kernel-selection oracles. + + Concrete oracles MUST implement: `backend_enum_cls`, + `get_priority_backends`, `backend_to_kernel_cls`, `map_backend`, + `select_backend`, `make_kernel`. + + Concrete oracles MAY override: `convert_to_kernel_format`, + `make_quant_config`. The base class provides default implementations + that are appropriate for oracles which do not need them + (e.g. `make_quant_config` raises on the unquantized oracle). + """ + + @abstractmethod + def backend_enum_cls(self) -> type[BackendT]: + """Return the concrete `Enum` class enumerating this oracle's + backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`).""" + + @abstractmethod + def get_priority_backends(self, moe_config: FusedMoEConfig) -> list[BackendT]: + """Return platform-appropriate backends in priority order for + this `moe_config`.""" + + @abstractmethod + def backend_to_kernel_cls(self, backend: BackendT) -> type[mk.FusedMoEExperts]: + """Map a backend enum value to its concrete `FusedMoEExperts` + subclass.""" + + @abstractmethod + def map_backend(self, runner_backend: MoEBackend) -> BackendT: + """Map a user-facing `MoEBackend` (from the runner config) to + this oracle's enum.""" + + @abstractmethod + def select_backend( + self, + moe_config: FusedMoEConfig, + weight_key: "QuantKey | None" = None, + activation_key: "QuantKey | None" = None, + ) -> tuple[BackendT, type[mk.FusedMoEExperts] | None]: + """Primary entry point: choose the best supported backend for + the given `moe_config`. + + `weight_key` / `activation_key` carry the quantization scheme of + the weights and activations and are consumed by quantized oracles + (fp8, nvfp4, int8, ...) to disambiguate backends. The unquantized + oracle ignores them. Subclasses with additional selection inputs + (e.g. int_wna16 needs `weight_bits`, fp8 needs + `allow_vllm_cutlass`) widen the signature in their override; a + per-oracle config object is the longer-term target tracked in + the #37753 follow-up PRs. + """ + + @abstractmethod + def make_kernel( + self, + quant_config: FusedMoEQuantConfig, + moe_config: FusedMoEConfig, + backend: BackendT, + experts_cls: type[mk.FusedMoEExperts], + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEKernel: + """Construct the `FusedMoEKernel` (Prepare/Finalize + Experts + combinator) for the chosen backend.""" + + def convert_to_kernel_format( + self, + backend: BackendT, + moe_config: FusedMoEConfig, + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Shuffle weights into the layout expected by `backend`. + + Default implementation returns the inputs unchanged. Oracles + whose backends need weight permutation should override this + (e.g. `UnquantizedMoEKernelOracle` handles AITER and FlashInfer + layouts). + + `moe_config` carries MoE-layer state (e.g. `is_act_and_mul`) + that the conversion needs without coupling the oracle to a + `Module` reference. Quantized oracles whose conversion + additionally needs scales / zero-points / block shapes will + override with a wider signature (and ultimately a per-oracle + config object — tracked in the #37753 follow-up PRs). + """ + return w13_weight, w2_weight + + def make_quant_config(self, *args, **kwargs) -> FusedMoEQuantConfig: + """Build a `FusedMoEQuantConfig` for this oracle. + + Quantized oracles (fp8, nvfp4, mxfp4, ...) override this with + the appropriate signature for their quantization scheme. + Unquantized oracles inherit the default, which raises because + there is no quantization-specific config to build. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement make_quant_config; " + "this oracle has no quantization-specific config to build." + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index a7dcd801376..a8ed9c1d7fe 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -2,9 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum +from typing import TYPE_CHECKING import torch -from torch.nn import Module import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEQuantConfig, ) +from vllm.model_executor.layers.fused_moe.oracle.base import MoEKernelOracle from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( align_moe_weights_for_fi, convert_moe_weights_to_flashinfer_trtllm_block_layout, @@ -25,6 +26,9 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( ) from vllm.platforms import current_platform +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + logger = init_logger(__name__) @@ -256,7 +260,7 @@ def select_unquantized_moe_backend( def convert_to_unquantized_kernel_format( unquantized_backend: UnquantizedMoeBackend, - layer: Module, + moe_config: FusedMoEConfig, w13_weight: torch.Tensor, w2_weight: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -264,13 +268,13 @@ def convert_to_unquantized_kernel_format( w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(w13_weight, w2_weight) elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS: - if layer.moe_config.is_act_and_mul: + if moe_config.is_act_and_mul: # Swap halves to arrange as [w3; w1] (kernel expectation) # Non-gated MoE: w13 is a single projection, no need to swap. w13_weight = swap_w13_to_w31(w13_weight) elif unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: - is_act_and_mul = layer.moe_config.is_act_and_mul + is_act_and_mul = moe_config.is_act_and_mul if not is_act_and_mul: # Kernel requires intermediate_size_per_partition % 128 == 0 (BlockMajorK # weight layout uses block_k=128). Pad along the intermediate dim when @@ -278,7 +282,7 @@ def convert_to_unquantized_kernel_format( w13_weight, w2_weight, padded_intermediate = align_moe_weights_for_fi( w13_weight, w2_weight, is_act_and_mul, min_alignment=128 ) - layer.moe_config.intermediate_size_per_partition = padded_intermediate + moe_config.intermediate_size_per_partition = padded_intermediate _cache_permute_indices: dict[torch.Size, torch.Tensor] = {} w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout( @@ -333,3 +337,71 @@ def make_unquantized_moe_kernel( ) return kernel + + +# --------------------------------------------------------------------------- +# Class-based view (first PR of the #37753 series; see oracle/base.py). +# Methods delegate to the module-level functions above so behaviour is +# bit-identical with pre-class code. +# --------------------------------------------------------------------------- + + +class UnquantizedMoEKernelOracle(MoEKernelOracle[UnquantizedMoeBackend]): + """Class-based view of the unquantized MoE kernel oracle. + + Each method delegates to its module-level counterpart so that + instantiating and calling this class is bit-identical to calling + the standalone functions. Follow-up PRs may move logic from the + module-level functions into these methods. + """ + + def backend_enum_cls(self) -> type[UnquantizedMoeBackend]: + return UnquantizedMoeBackend + + def get_priority_backends( + self, moe_config: FusedMoEConfig + ) -> list[UnquantizedMoeBackend]: + return _get_priority_backends(moe_config) + + def backend_to_kernel_cls( + self, backend: UnquantizedMoeBackend + ) -> type[mk.FusedMoEExperts]: + return backend_to_kernel_cls(backend) + + def map_backend(self, runner_backend: MoEBackend) -> UnquantizedMoeBackend: + return map_unquantized_backend(runner_backend) + + def select_backend( + self, + moe_config: FusedMoEConfig, + weight_key: "QuantKey | None" = None, + activation_key: "QuantKey | None" = None, + ) -> tuple[UnquantizedMoeBackend, type[mk.FusedMoEExperts] | None]: + assert weight_key is None and activation_key is None, ( + "Weights and activations will never be quantized for " + "UnquantizedMoEKernelOracle" + ) + return select_unquantized_moe_backend(moe_config) + + def convert_to_kernel_format( + self, + backend: UnquantizedMoeBackend, + moe_config: FusedMoEConfig, + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return convert_to_unquantized_kernel_format( + backend, moe_config, w13_weight, w2_weight + ) + + def make_kernel( + self, + quant_config: FusedMoEQuantConfig, + moe_config: FusedMoEConfig, + backend: UnquantizedMoeBackend, + experts_cls: type[mk.FusedMoEExperts], + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEKernel: + return make_unquantized_moe_kernel( + quant_config, moe_config, backend, experts_cls, routing_tables + ) diff --git a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py index bd4393be5e7..7a2c670a8ce 100644 --- a/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py +++ b/vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py @@ -161,7 +161,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): # Shuffle weights to runtime format. w13_new, w2_new = convert_to_unquantized_kernel_format( self.unquantized_backend, - layer=layer, + moe_config=layer.moe_config, w13_weight=w13, w2_weight=w2, ) diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index f356ce6f4ff..fce74346d62 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools from math import prod import torch @@ -446,3 +447,12 @@ def swiglu_limit_func( up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit) output.copy_(F.silu(gate) * up) + + +@functools.lru_cache +def enable_swap_ab(BLOCK_SIZE_M: int, BLOCK_SIZE_N: int) -> bool: + return ( + current_platform.is_device_capability(90) + and BLOCK_SIZE_M < 64 + and BLOCK_SIZE_N >= 64 + ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 229112739a4..d52386d5d1a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -37,7 +37,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso CompressedTensorsMoEMethod, ) from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( - WNA16_SUPPORTED_BITS, CompressedTensorsScheme, CompressedTensorsW4A4Fp4, CompressedTensorsW4A4Mxfp4, @@ -680,14 +679,7 @@ class CompressedTensorsConfig(QuantizationConfig): and output_quant.num_bits == 8 and not output_quant.dynamic ) - # Static int8-activation layers, plus sub-byte weight-only layers (e.g. - # 2-bit lm_head) that marlin-backed WNA16 cannot serve. Standard 4/8-bit - # weight-only (no activations) falls through to WNA16. - is_subbyte_weight_only = weight_quant.num_bits not in WNA16_SUPPORTED_BITS - needs_wNa8o8 = is_intN_weight and ( - (is_static_int8_in and is_static_int8_out) or is_subbyte_weight_only - ) - return needs_wNa8o8 + return is_intN_weight and (is_static_int8_in or is_static_int8_out) def _get_scheme_from_parts( self, @@ -740,10 +732,8 @@ class CompressedTensorsConfig(QuantizationConfig): quant_format=format, ) - if ( - self._is_wNa16_group_channel(weight_quant, input_quant) - and (format == CompressionFormat.pack_quantized.value) - and (weight_quant.num_bits in WNA16_SUPPORTED_BITS) + if self._is_wNa16_group_channel(weight_quant, input_quant) and ( + format == CompressionFormat.pack_quantized.value ): return CompressedTensorsWNA16( num_bits=weight_quant.num_bits, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py index d81db4a052f..2826bf7b471 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py @@ -11,7 +11,7 @@ from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8 from .compressed_tensors_w8a8_mxfp8 import CompressedTensorsW8A8Mxfp8 from .compressed_tensors_w8a16_fp8 import CompressedTensorsW8A16Fp8 from .compressed_tensors_wNa8o8 import CompressedTensorsWNA8O8Int -from .compressed_tensors_wNa16 import WNA16_SUPPORTED_BITS, CompressedTensorsWNA16 +from .compressed_tensors_wNa16 import CompressedTensorsWNA16 __all__ = [ "CompressedTensorsScheme", @@ -20,7 +20,6 @@ __all__ = [ "CompressedTensorsW8A16Fp8", "CompressedTensorsW8A8Int8", "CompressedTensorsW8A8Fp8", - "WNA16_SUPPORTED_BITS", "CompressedTensorsW4A4Mxfp4", "CompressedTensorsW4A4Fp4", "CompressedTensorsW4A8Int", diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py index 1883d4ae322..f69c11f3d5e 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math from collections.abc import Callable +from fractions import Fraction import torch from compressed_tensors.quantization import ActivationOrdering @@ -32,7 +34,15 @@ from vllm.scalar_type import scalar_types logger = init_logger(__name__) __all__ = ["CompressedTensorsWNA16"] -WNA16_SUPPORTED_TYPES_MAP = {4: scalar_types.uint4b8, 8: scalar_types.uint8b128} +WNA16_SUPPORTED_TYPES_MAP = { + 2: scalar_types.uint2b2, + 3: scalar_types.uint3b4, + 4: scalar_types.uint4b8, + 5: scalar_types.uint5b16, + 6: scalar_types.uint6b32, + 7: scalar_types.uint7b64, + 8: scalar_types.uint8b128, +} WNA16_ZP_SUPPORTED_TYPES_MAP = {4: scalar_types.uint4, 8: scalar_types.uint8} WNA16_SUPPORTED_BITS = list(WNA16_SUPPORTED_TYPES_MAP.keys()) @@ -49,7 +59,8 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): actorder: ActivationOrdering | None = None, layer_name: str | None = None, ): - self.pack_factor = 32 // num_bits + self.num_bits = num_bits + self.pack_factor = Fraction(32, num_bits) self.strategy = strategy self.symmetric = symmetric self.group_size = -1 if group_size is None else group_size @@ -58,15 +69,22 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): if self.group_size == -1 and self.strategy != "channel": raise ValueError( - "Marlin kernels require group quantization or " - "channelwise quantization, but found no group " + "Pack-quantized format requires group quantization " + "or channelwise quantization, but found no group " "size and strategy is not channelwise." ) if num_bits not in WNA16_SUPPORTED_TYPES_MAP: raise ValueError( f"Unsupported num_bits = {num_bits}. " - f"Supported num_bits = {WNA16_SUPPORTED_TYPES_MAP.keys()}" + f"Supported num_bits = {list(WNA16_SUPPORTED_TYPES_MAP)}" + ) + + if not self.symmetric and num_bits not in WNA16_ZP_SUPPORTED_TYPES_MAP: + raise ValueError( + f"Asymmetric quantization not supported for " + f"num_bits = {num_bits}. Supported: " + f"{list(WNA16_ZP_SUPPORTED_TYPES_MAP)}" ) self.quant_type = ( @@ -92,6 +110,12 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): **kwargs, ): output_size_per_partition = sum(output_partition_sizes) + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.output_partition_sizes = output_partition_sizes + layer.params_dtype = params_dtype + if not hasattr(layer, "has_bias"): + layer.has_bias = False mp_linear_kernel_config = MPLinearLayerConfig( full_weight_shape=(input_size, output_size), @@ -130,6 +154,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): assert input_size_per_partition % group_size == 0 scales_and_zp_size = input_size_per_partition // group_size + packed_input_dim = math.ceil(input_size_per_partition * self.num_bits / 32) weight = PackedvLLMParameter( input_dim=1, output_dim=0, @@ -138,7 +163,7 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): packed_dim=1, data=torch.empty( output_size_per_partition, - input_size_per_partition // self.pack_factor, + packed_input_dim, dtype=torch.int32, ), ) @@ -152,10 +177,11 @@ class CompressedTensorsWNA16(CompressedTensorsScheme): ), } + packed_output_dim = math.ceil(output_size_per_partition * self.num_bits / 32) zeros_args = { "weight_loader": weight_loader, "data": torch.zeros( - output_size_per_partition // self.pack_factor, + packed_output_dim, scales_and_zp_size, dtype=torch.int32, ), diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 820260f795c..36c0a357878 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -131,7 +131,7 @@ class DFlashQwen3Attention(nn.Module): with the context K/V from the target model's hidden states. This forward op computes attention for the query tokens only. See also: precompute_and_store_context_kv""" - qkv = F.linear(hidden_states, self.qkv_proj.weight, self.qkv_proj.bias) + qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # Per-head RMSNorm diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index 4d8ff951c09..5ec0b82dabd 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/vllm/model_executor/models/qwen3_next_mtp.py @@ -64,13 +64,22 @@ class Qwen3NextMultiTokenPredictor(nn.Module): config.hidden_size, ) + # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is + # missing from the checkpoint quant exclude list (its `ignore` glob + # does not cover `mtp.fc`). Force unquantized to match the weights, + # mirroring the Qwen3.5 MTP handling (PR #38832). + fc_quant = ( + None + if (quant_config and quant_config.get_name() == "modelopt_fp4") + else quant_config + ) self.fc = ColumnParallelLinear( self.config.hidden_size * 2, self.config.hidden_size, gather_output=True, bias=False, return_bias=False, - quant_config=quant_config, + quant_config=fc_quant, prefix=f"{prefix}.fc", ) @@ -242,7 +251,7 @@ class Qwen3NextMTP(nn.Module, QwenNextMixtureOfExperts): "k_proj", "v_proj", ], - "gate_up_proj": ["up_proj", "down_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], } def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 0cac651e474..0a90d9f9c28 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -72,8 +72,6 @@ _TEXT_GENERATION_MODELS = { # [Decoder-only] "AfmoeForCausalLM": ("afmoe", "AfmoeForCausalLM"), "ApertusForCausalLM": ("apertus", "ApertusForCausalLM"), - "AquilaModel": ("llama", "LlamaForCausalLM"), - "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2 "ArceeForCausalLM": ("arcee", "ArceeForCausalLM"), "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"), "AXK1ForCausalLM": ("AXK1", "AXK1ForCausalLM"), @@ -731,6 +729,8 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "MiniMaxVL01ForConditionalGeneration": "0.23.0", "BaiChuanForCausalLM": "0.23.0", "BaichuanForCausalLM": "0.23.0", + "AquilaModel": "0.24.0", + "AquilaForCausalLM": "0.24.0", } _OOT_SUPPORTED_MODELS = { diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py new file mode 100644 index 00000000000..2b076a38b89 --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/index_topk.py @@ -0,0 +1,939 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton kernels for MiniMax M3 lightning-indexer block scoring + top-k. + +Index queries score each 128-token block of index keys (max over the block), +then the top-k blocks (plus forced init/local blocks) are selected per query +token. Adapted to vLLM's paged KV cache: the KV page size is forced to equal the +sparse block size (128), so one sparse block maps to exactly one page. + +Index-K cache layout (vLLM): ``(num_blocks, 128, idx_head_dim)`` (single head). + +Only the paths MiniMax M3 uses are implemented: score_type="max", index value +disabled (score-only indexer), single shared index head. The selected block ids +feed the block-sparse attention kernels in ``sparse_attn``. +""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import round_up + +# One sparse block == one KV page. +SPARSE_BLOCK_SIZE = 128 + + +# --------------------------------------------------------------------------- +# Bitonic top-k helpers (layout-agnostic). +# --------------------------------------------------------------------------- +@triton.jit +def _compare_and_swap(x, ids, flip, i: tl.constexpr, n_dims: tl.constexpr): + n_outer: tl.constexpr = x.numel >> n_dims + shape: tl.constexpr = [n_outer * 2**i, 2, 2 ** (n_dims - i - 1)] + y = tl.reshape(x, shape) + mask = tl.arange(0, 2)[None, :, None] + left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype) + right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype) + left = tl.reshape(left, x.shape) + right = tl.reshape(right, x.shape) + y_idx = tl.reshape(ids, shape) + left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape) + right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape) + left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype) + right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype) + idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) + ileft = left.to(idtype, bitcast=True) + iright = right.to(idtype, bitcast=True) + ix = x.to(idtype, bitcast=True) + cond = (left > right) != flip + ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix)) + new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids)) + return ret.to(x.dtype, bitcast=True), new_ids + + +@triton.jit +def _bitonic_merge( + x, ids, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr +): + n_outer: tl.constexpr = x.numel >> n_dims + tl.static_assert(stage <= n_dims) + if order == 2: + shape: tl.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2**stage] + flip = tl.reshape( + tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape + ) + else: + flip = order + for i in tl.static_range(stage): + x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims) + return x, ids + + +# --------------------------------------------------------------------------- +# Index block-score kernel (paged). score[h, token, block] = max over the +# 128-token block of (idx_q . index_k), causal-masked. BLOCK_SIZE_K == 128 so +# each K-tile is exactly one page (BLOCKS_PER_K_BLOCK == 1). +# --------------------------------------------------------------------------- +# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens +# might lose pointer alignment, which trigger Triton recompiles. we don't actually +# need pointer alignment for those tensors anyway because we do scalar load. +@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"]) +def _index_block_score_kernel( + q_ptr, # idx_q: [total_q, num_idx_heads, head_dim] + ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim] + score_ptr, # [num_idx_heads, total_q, max_block] + block_table_ptr, # [num_reqs, max_blocks] + cu_seqlens, # [batch+1] query start offsets + seq_lens, # [batch] total K length + prefix_lens, # [batch] context length before this chunk's queries + num_idx_heads, + head_dim: tl.constexpr, + stride_q_n, + stride_q_h, + stride_q_d, + stride_ik_blk, + stride_ik_pos, + stride_ik_d, + stride_s_h, + stride_s_n, + stride_s_k, + stride_bt_b, + BLOCK_SIZE_Q: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) +): + pid_q = tl.program_id(0) + pid_bh = tl.program_id(1) + pid_b = pid_bh // num_idx_heads + pid_h = pid_bh % num_idx_heads + + seq_start = tl.load(cu_seqlens + pid_b) + q_len = tl.load(cu_seqlens + pid_b + 1) - seq_start + seq_len = tl.load(seq_lens + pid_b) + prefix_len = tl.load(prefix_lens + pid_b) + if BLOCK_SIZE_Q * pid_q >= q_len: + return + + q_ptrs = tl.make_block_ptr( + base=q_ptr + seq_start * stride_q_n + pid_h * stride_q_h, + shape=(q_len, head_dim), + strides=(stride_q_n, stride_q_d), + offsets=(pid_q * BLOCK_SIZE_Q, 0), + block_shape=(BLOCK_SIZE_Q, head_dim), + order=(1, 0), + ) + q = tl.load(q_ptrs, boundary_check=(0,), padding_option="zero") + q_start = prefix_len + pid_q * BLOCK_SIZE_Q + + off_q = tl.arange(0, BLOCK_SIZE_Q) + pid_q * BLOCK_SIZE_Q + prefix_len + off_k = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, head_dim) + # Block table row for this request. + bt_row = block_table_ptr + pid_b * stride_bt_b + # Causal window: only blocks up to the last query token's position. + hi = min(seq_len, prefix_len + (pid_q + 1) * BLOCK_SIZE_Q) + for i in tl.range(0, hi, BLOCK_SIZE_K): + blk = i // BLOCK_SIZE_K + page = tl.load(bt_row + blk).to(tl.int64) + pos = i + off_k + # index-K for this page: [BLOCK_SIZE_D, BLOCK_SIZE_K] (transposed) + # we don't need masked load for K, because KV cache ensures + # allocation is multiple of BLOCK_SIZE_K. + # for tokens beyond seqlen, they will be masked in qk later. + k = tl.load( + ik_cache_ptr + + page * stride_ik_blk + + off_k[None, :] * stride_ik_pos + + off_d[:, None] * stride_ik_d, + ) + qk = tl.dot(q, k) + # apply causal mask as needed + if q_start < i + BLOCK_SIZE_K: + qk = tl.where(off_q[:, None] >= pos[None, :], qk, float("-inf")) + # one sparse block per K-tile -> max over the 128 positions + score = tl.max(qk, axis=1) # [BLOCK_SIZE_Q] + s_ptrs = ( + score_ptr + + pid_h * stride_s_h + + (seq_start + pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) + * stride_s_n + + blk * stride_s_k + ) + q_store_mask = (pid_q * BLOCK_SIZE_Q + tl.arange(0, BLOCK_SIZE_Q)) < q_len + tl.store(s_ptrs, score, mask=q_store_mask) + + +# --------------------------------------------------------------------------- +# Top-k selection over per-token block scores (layout-agnostic). block_size_q +# is 1 for M3, so top-k is computed per query token. +# --------------------------------------------------------------------------- +# since prefill metadata is sliced from mixed batch metadata, prefix_lens +# might lose pointer alignment, which trigger Triton recompiles. we don't actually +# need pointer alignment for those tensors anyway because we do scalar load. +@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])}) +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE_K": 2048}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 1024}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 512}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2), + ], + key=["BLOCK_SIZE_T"], +) +@triton.jit(do_not_specialize_on_alignment=["prefix_lens"]) +def _topk_index_kernel( + s_ptr, # [num_heads, total_q, max_block] + ti_ptr, # [num_heads, total_q, topk] + sample_interval: tl.constexpr, # block_size_q (1 for M3) + block_size: tl.constexpr, # sparse block size (128) + cu_seqlens, + cu_seqblocks_q, + prefix_lens, + topk, + init_blocks: tl.constexpr, + local_blocks: tl.constexpr, + stride_s_h, + stride_s_n, + stride_s_k, + stride_ti_h, + stride_ti_n, + stride_ti_t, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + MASK_INIT: tl.constexpr, + MASK_LOCAL: tl.constexpr, +): + tl.static_assert(BLOCK_SIZE_K > BLOCK_SIZE_T) + pid_q = tl.program_id(0) + pid_b = tl.program_id(1) + pid_h = tl.program_id(2) + seq_start = tl.load(cu_seqlens + pid_b) + block_start = tl.load(cu_seqblocks_q + pid_b) + block_num = tl.load(cu_seqblocks_q + pid_b + 1) - block_start + prefix_len = tl.load(prefix_lens + pid_b) + if pid_q >= block_num: + return + off_k = tl.arange(0, BLOCK_SIZE_K) + off_t = tl.arange(0, BLOCK_SIZE_T) + s_ptrs = ( + s_ptr + + (seq_start + pid_q * sample_interval) * stride_s_n + + pid_h * stride_s_h + + off_k * stride_s_k + ) + topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32) + topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32) + left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2 + valid_blocks = (prefix_len + pid_q * sample_interval + block_size) // block_size + for i in tl.range(0, valid_blocks, BLOCK_SIZE_K): + causal_mask = i + off_k < valid_blocks + local_mask = i + off_k >= max(0, valid_blocks - local_blocks) + init_mask = i + off_k < init_blocks + score = tl.load(s_ptrs, mask=causal_mask, other=-1e30).to(tl.float32) + score = tl.where(score != score, -1e30, score) + s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K + if MASK_INIT: + score = tl.where(causal_mask & init_mask, score - 1e29, score) + else: + score = tl.where(causal_mask & init_mask, 1e30, score) + if MASK_LOCAL: + score = tl.where(causal_mask & local_mask, score - 1e28, score) + else: + score = tl.where(causal_mask & local_mask, 1e29, score) + topk_score, last_topk_score = score, topk_score + topk_idx, last_topk_idx = (tl.where(causal_mask, i + off_k + 1, 0), topk_idx) + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) + for j in tl.static_range(1, n_dims): + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), j, 2, n_dims + ) + if i != 0: + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims + ) + topk_score_new = last_topk_score * left_half_mask + topk_score * ( + 1 - left_half_mask + ) + topk_idx_new = last_topk_idx * left_half_mask + topk_idx * ( + 1 - left_half_mask + ) + topk_score, topk_idx = _bitonic_merge( + topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims + ) + else: + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims + ) + topk_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 + topk_idx = tl.sum( + topk_mask[:, None] + * tl.reshape(topk_idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), + axis=0, + ) + ti_ptrs = ( + ti_ptr + + (block_start + pid_q) * stride_ti_n + + pid_h * stride_ti_h + + off_t * stride_ti_t + ) + store_mask = off_t < topk + valid_mask = off_t < valid_blocks + topk_idx = tl.where(store_mask & valid_mask, topk_idx, -1) + tl.store(ti_ptrs, topk_idx.to(ti_ptrs.dtype.element_ty), mask=store_mask) + + +# --------------------------------------------------------------------------- +# Decode index-score kernel (split-K over seq blocks). Decode batches are +# flattened request-major, with a runtime query length used to map each query +# token back to its request metadata. Chunk counts depend only on shape +# constants so the grid is fixed within a cuda graph. The score scale is omitted +# because decode only consumes block ordering. +# --------------------------------------------------------------------------- +@triton.jit(do_not_specialize=["num_kv_chunks", "decode_query_len"]) +def _decode_index_score_kernel( + q_ptr, # idx_q: [total_q, num_idx_heads, head_dim] + ik_cache_ptr, # index-K cache: [num_blocks, 128, head_dim] + score_ptr, # [num_idx_heads, total_q, max_block] + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + num_idx_heads: tl.constexpr, + head_dim: tl.constexpr, + init_blocks, + local_blocks, + decode_query_len, + stride_q_n, + stride_q_h, + stride_q_d, + stride_ik_blk, + stride_ik_pos, + stride_ik_d, + stride_s_h, + stride_s_n, + stride_s_k, + stride_bt_b, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + BLOCK_SIZE_Q: tl.constexpr, + num_kv_chunks, + USE_PDL: tl.constexpr, +): + BLOCK_SIZE_HQ: tl.constexpr = num_idx_heads * BLOCK_SIZE_Q + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) + hq_offsets = tl.arange(0, BLOCK_SIZE_HQ) + h_offsets = hq_offsets // BLOCK_SIZE_Q + q_offsets = hq_offsets % BLOCK_SIZE_Q + q_mask = q_offsets < decode_query_len + q_ids = pid_r * decode_query_len + q_offsets + + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + pid_r) + query_pos = seq_len - decode_query_len + q_offsets + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + num_blocks_q = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + kv_len_max = tl.max(tl.where(q_mask, kv_len, 0), axis=0) + num_blocks = (kv_len_max + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + + # block-aligned fixed-count split: grid independent of seq_len (cuda graph). + chunk_size_blocks = (num_blocks + num_kv_chunks - 1) // num_kv_chunks + chunk_start_block = pid_c * chunk_size_blocks + chunk_end_block = tl.minimum(chunk_start_block + chunk_size_blocks, num_blocks) + if chunk_start_block >= chunk_end_block: + return + off_k = tl.arange(0, BLOCK_SIZE_K) # positions within a 128-block + off_d = tl.arange(0, head_dim) + bt_row = block_table_ptr + pid_r * stride_bt_b + # Force-select init (1e30) and local (1e29, higher priority) blocks. + local_start = tl.maximum(0, num_blocks_q - local_blocks) + # Query vectors for all index heads in a small spec-decode block. + q = tl.load( + q_ptr + + q_ids[None, :] * stride_q_n + + h_offsets[None, :] * stride_q_h + + off_d[:, None] * stride_q_d, + mask=q_mask[None, :], + other=0.0, + ) # [D,HQ] + for blk in tl.range(chunk_start_block, chunk_end_block): + page = tl.load(bt_row + blk).to(tl.int64) + pos = blk * BLOCK_SIZE_K + off_k + pos_mask = pos[:, None] < kv_len[None, :] + # we don't need masked load for K, because KV cache ensures + # allocation is multiple of BLOCK_SIZE_K. + # for tokens beyond seqlen, they will be masked in qk later. + k = tl.load( + ik_cache_ptr + + page * stride_ik_blk + + off_k[:, None] * stride_ik_pos + + off_d * stride_ik_d, + ) # [N,D] + if BLOCK_SIZE_HQ == 1: + # Degenerate GEMV (q is [D,1]): vectorized fp32 multiply + reduce + # instead of an MFMA tile. Numerically equivalent to tl.dot. + q_vec = tl.sum(q, axis=1).to(tl.float32) # [D] + kq = tl.sum(k.to(tl.float32) * q_vec[None, :], axis=1)[:, None] # [N,1] + else: + # fp32 accumulation is required for the fp8 (e4m3) index cache: q/k + # are loaded in their stored dtype (bf16 or e4m3) and the MMA + # accumulates in fp32 so the per-block max score is exact for the + # fp8 indexer too. + kq = tl.dot(k, q, out_dtype=tl.float32) # [N,HQ] + kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf")) + score = tl.max(kq, axis=0) # [HQ] + is_visible_block = blk < num_blocks_q + is_init = (blk < init_blocks) & is_visible_block + is_local = (blk >= local_start) & is_visible_block + score = tl.where(is_local, 1e29, tl.where(is_init, 1e30, score)) + tl.store( + score_ptr + h_offsets * stride_s_h + q_ids * stride_s_n + blk * stride_s_k, + score, + mask=q_mask, + ) + + +# --------------------------------------------------------------------------- +# Decode top-k (split-K): per-chunk partial top-k + merge. Forced init/local +# blocks are already encoded in the scores. +# --------------------------------------------------------------------------- +@triton.heuristics({"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"])}) +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_SIZE_K": 128}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_SIZE_K": 64}, num_warps=2, num_stages=2), + ], + key=["topk"], +) +@triton.jit(do_not_specialize=["chunk_blocks", "decode_query_len"]) +def _topk_index_partial_kernel( + s_ptr, # score: [num_idx_heads, total_q, max_block] + ts_partial_ptr, # partial scores out: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T] + ti_partial_ptr, # partial idx out (1-indexed global, 0=invalid): same shape + seq_lens, # [num_reqs] + block_size: tl.constexpr, # sparse block size (128) + topk: tl.constexpr, + chunk_blocks, # how many score-blocks each chunk owns + decode_query_len, + stride_s_h, + stride_s_b, + stride_s_k, + stride_ts_c, + stride_ts_h, + stride_ts_b, + stride_ts_t, + stride_ti_c, + stride_ti_h, + stride_ti_b, + stride_ti_t, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + USE_PDL: tl.constexpr, +): + tl.static_assert(topk < BLOCK_SIZE_K) + pid_b = tl.program_id(0) # flattened query-token id + pid_h = tl.program_id(1) + pid_chunk = tl.program_id(2) + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + num_blocks = (kv_len + block_size - 1) // block_size + + # Slice this chunk owns within [0, num_blocks). + chunk_start = pid_chunk * chunk_blocks + chunk_end = tl.minimum(chunk_start + chunk_blocks, num_blocks) + chunk_actual = tl.maximum(chunk_end - chunk_start, 0) + + off_k = tl.arange(0, BLOCK_SIZE_K) + off_t = tl.arange(0, BLOCK_SIZE_T) + + s_ptrs = ( + s_ptr + + pid_b * stride_s_b + + pid_h * stride_s_h + + (chunk_start + off_k) * stride_s_k + ) + + topk_score = tl.full((BLOCK_SIZE_K,), -1e30, dtype=tl.float32) + topk_idx = tl.full((BLOCK_SIZE_K,), 0, dtype=tl.int32) + left_half_mask = tl.arange(0, BLOCK_SIZE_K) < BLOCK_SIZE_K // 2 + + # Streaming top-K within this chunk. tl.range(0, 0) is a no-op so empty + # chunks (chunk_actual == 0) skip the body and store sentinel -1e30 / 0. + for i in tl.range(0, chunk_actual, BLOCK_SIZE_K): + mask = off_k < chunk_actual - i + score = tl.load(s_ptrs, mask=mask, other=-1e30).to(tl.float32) + score = tl.where(score != score, -1e30, score) + s_ptrs = s_ptrs + stride_s_k * BLOCK_SIZE_K + topk_score, last_topk_score = score, topk_score + topk_idx, last_topk_idx = ( + tl.where(mask, chunk_start + i + off_k + 1, 0), # 1-indexed global + topk_idx, + ) + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) + for j in tl.static_range(1, n_dims): + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), j, 2, n_dims + ) + if i != 0: + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), n_dims, False, n_dims + ) + topk_score_new = last_topk_score * left_half_mask + topk_score * ( + 1 - left_half_mask + ) + topk_idx_new = last_topk_idx * left_half_mask + topk_idx * ( + 1 - left_half_mask + ) + topk_score, topk_idx = _bitonic_merge( + topk_score_new, topk_idx_new.to(tl.int32), n_dims, True, n_dims + ) + else: + topk_score, topk_idx = _bitonic_merge( + topk_score, topk_idx.to(tl.int32), n_dims, True, n_dims + ) + + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + # Extract first BLOCK_SIZE_T entries (top-K of this chunk after the sort). + topk_mask_extract = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 + final_score = tl.sum( + topk_mask_extract[:, None] + * tl.reshape(topk_score, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), + axis=0, + ) + final_idx = tl.sum( + topk_mask_extract[:, None] + * tl.reshape(topk_idx, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), + axis=0, + ) + + # Always write all BLOCK_SIZE_T slots — invalid slots carry -1e30 / 0 + # sentinels and lose to real scores in the merge stage. + ts_ptrs = ( + ts_partial_ptr + + pid_chunk * stride_ts_c + + pid_b * stride_ts_b + + pid_h * stride_ts_h + + off_t * stride_ts_t + ) + ti_ptrs = ( + ti_partial_ptr + + pid_chunk * stride_ti_c + + pid_b * stride_ti_b + + pid_h * stride_ti_h + + off_t * stride_ti_t + ) + tl.store(ts_ptrs, final_score) + tl.store(ti_ptrs, final_idx) + + +@triton.heuristics( + { + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["topk"]), + "BLOCK_SIZE_K": lambda args: triton.next_power_of_2( + args["num_topk_chunks"] * triton.next_power_of_2(args["topk"]) + ), + } +) +@triton.jit(do_not_specialize=["num_topk_chunks", "decode_query_len"]) +def _topk_index_merge_kernel( + ts_partial_ptr, # partial scores: [NUM_TOPK_CHUNKS, num_idx_heads, total_q, T] + ti_partial_ptr, # partial idx (1-indexed global, 0=invalid): same shape + ti_final_ptr, # final idx (0-indexed, -1=invalid): [num_idx_heads, total_q, topk] + seq_lens, # [num_reqs] + block_size: tl.constexpr, # sparse block size (128) + topk: tl.constexpr, + decode_query_len, + stride_ts_c, + stride_ts_h, + stride_ts_b, + stride_ts_t, + stride_ti_c, + stride_ti_h, + stride_ti_b, + stride_ti_t, + stride_tif_h, + stride_tif_b, + stride_tif_t, + num_topk_chunks, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + USE_PDL: tl.constexpr, +): + pid_b = tl.program_id(0) # flattened query-token id + pid_h = tl.program_id(1) + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + + if USE_PDL: + tl.extra.cuda.gdc_wait() + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + num_blocks = (kv_len + block_size - 1) // block_size + + # Load NUM_TOPK_CHUNKS * BLOCK_SIZE_T candidates, padded to BLOCK_SIZE_K. + # Candidate at flat position p comes from chunk = p // BLOCK_SIZE_T, + # in_chunk = p % BLOCK_SIZE_T. + off = tl.arange(0, BLOCK_SIZE_K) + chunk_idx = off // BLOCK_SIZE_T + in_chunk_idx = off % BLOCK_SIZE_T + valid = chunk_idx < num_topk_chunks + + score_offset = ( + chunk_idx * stride_ts_c + + pid_h * stride_ts_h + + pid_b * stride_ts_b + + in_chunk_idx * stride_ts_t + ) + idx_offset = ( + chunk_idx * stride_ti_c + + pid_h * stride_ti_h + + pid_b * stride_ti_b + + in_chunk_idx * stride_ti_t + ) + + score = tl.load(ts_partial_ptr + score_offset, mask=valid, other=-1e30).to( + tl.float32 + ) + score = tl.where(score != score, -1e30, score) + idx = tl.load(ti_partial_ptr + idx_offset, mask=valid, other=0).to(tl.int32) + + # Full bitonic descending sort of BLOCK_SIZE_K items. + n_dims: tl.constexpr = tl.standard._log2(BLOCK_SIZE_K) + for j in tl.static_range(1, n_dims): + score, idx = _bitonic_merge(score, idx.to(tl.int32), j, 2, n_dims) + score, idx = _bitonic_merge(score, idx.to(tl.int32), n_dims, True, n_dims) + + # Extract first BLOCK_SIZE_T positions — these are the global top-K. + extract_mask = tl.arange(0, BLOCK_SIZE_K // BLOCK_SIZE_T) == 0 + topk_idx_final = tl.sum( + extract_mask[:, None] + * tl.reshape(idx - 1, [BLOCK_SIZE_K // BLOCK_SIZE_T, BLOCK_SIZE_T]), + axis=0, + ) + + off_t = tl.arange(0, BLOCK_SIZE_T) + tif_ptrs = ( + ti_final_ptr + + pid_h * stride_tif_h + + pid_b * stride_tif_b + + off_t * stride_tif_t + ) + store_mask = off_t < topk + topk_idx_final = tl.where(off_t < tl.minimum(topk, num_blocks), topk_idx_final, -1) + tl.store( + tif_ptrs, topk_idx_final.to(ti_final_ptr.dtype.element_ty), mask=store_mask + ) + + +# --------------------------------------------------------------------------- +# Python wrappers +# --------------------------------------------------------------------------- +@torch.no_grad() +def minimax_m3_index_score( + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] + block_table: torch.Tensor, # [batch, max_blocks] + cu_seqlens_q: torch.Tensor, # [batch+1] int32 + seq_lens: torch.Tensor, # [batch] int32 + prefix_lens: torch.Tensor, # [batch] int32 + max_query_len: int, + max_seq_len: int, + num_kv_heads: int, +) -> torch.Tensor: + """Compute per-token index scores for each visible sparse block. + + Returns score [num_kv_heads, total_q, max_block], where each score is the + max over a 128-token index-K block. M3 has num_idx_heads == num_kv_heads. + """ + total_q, num_idx_heads, head_dim = idx_q.shape + assert num_idx_heads == num_kv_heads, ( + "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)" + ) + batch = cu_seqlens_q.shape[0] - 1 + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) + + # Keep score strides 16-divisible to avoid Triton recompiles. + score_block_stride = round_up(max_block, 16) + score = torch.empty( + (num_idx_heads, total_q, score_block_stride), + dtype=torch.float32, + device=idx_q.device, + ) + BLOCK_SIZE_Q = 64 + grid_score = (triton.cdiv(max_query_len, BLOCK_SIZE_Q), batch * num_idx_heads) + _index_block_score_kernel[grid_score]( + idx_q, + index_kv_cache, + score, + block_table, + cu_seqlens_q, + seq_lens, + prefix_lens, + num_idx_heads, + head_dim, + idx_q.stride(0), + idx_q.stride(1), + idx_q.stride(2), + index_kv_cache.stride(0), + index_kv_cache.stride(1), + index_kv_cache.stride(2), + score.stride(0), + score.stride(1), + score.stride(2), + block_table.stride(0), + BLOCK_SIZE_Q=BLOCK_SIZE_Q, + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + ) + return score + + +@torch.no_grad() +def minimax_m3_index_topk( + score: torch.Tensor, # [num_idx_heads, total_q, max_block] + cu_seqlens_q: torch.Tensor, # [batch+1] int32 + prefix_lens: torch.Tensor, # [batch] int32 + max_query_len: int, + topk: int, + init_blocks: int, + local_blocks: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Select index top-k from a precomputed score tensor. + + When ``out`` is provided (a ``[num_idx_heads, >=total_q, topk]`` buffer), the + result is written into ``out[:, :total_q, :]`` instead of a fresh tensor -- + used to keep the top-k output at a stable address for cudagraph capture. + """ + num_idx_heads = score.shape[0] + batch = cu_seqlens_q.shape[0] - 1 + total_q = score.shape[1] + if out is not None: + topk_idx = out[:, :total_q, :] + else: + topk_idx = torch.empty( + (num_idx_heads, total_q, topk), + dtype=torch.int32, + device=score.device, + ) + # block_size_q == 1 -> query blocks coincide with query tokens. + grid_topk = (max_query_len, batch, num_idx_heads) + _topk_index_kernel[grid_topk]( + score, + topk_idx, + 1, # sample_interval (block_size_q) + SPARSE_BLOCK_SIZE, + cu_seqlens_q, + cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1 + prefix_lens, + topk, + init_blocks, + local_blocks, + score.stride(0), + score.stride(1), + score.stride(2), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + MASK_INIT=False, + MASK_LOCAL=False, + ) + return topk_idx + + +@torch.no_grad() +def minimax_m3_index_decode( + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim] + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + max_seq_len: int, + topk: int, + init_blocks: int, + local_blocks: int, + num_kv_heads: int, + decode_query_len: int, + max_decode_query_len: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Decode index block-score + top-k, both split-K (cudagraph-safe). + + Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad). + When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into + ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating. + """ + total_q, num_idx_heads, head_dim = idx_q.shape + assert num_idx_heads == num_kv_heads, ( + "M3 expects num_idx_heads == num_kv_heads (no topk index reduce)" + ) + assert decode_query_len <= max_decode_query_len + assert total_q == seq_lens.shape[0] * decode_query_len + batch = total_q + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE) + use_pdl = current_platform.is_arch_support_pdl() + # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA + # SM9+); this ROCm Triton rejects it even when False ("Keyword argument + # launch_pdl was specified but unrecognised"). Only pass it when PDL is + # actually supported -- on ROCm use_pdl is always False, so it's omitted. + pdl_kwargs: dict[str, bool | int] = {} + if use_pdl: + pdl_kwargs.update({"launch_pdl": True}) + # TP=1 spec decode scores a wide 4-head x 4-position query tile per K block; + # reduce stages to ease memory/register pressure. Keep no-spec and TP=4 + # single-head codegen unchanged. + score_kwargs = pdl_kwargs.copy() + if num_idx_heads > 1 and max_decode_query_len > 1: + score_kwargs.update({"num_warps": 4, "num_stages": 2}) + + # Keep score strides 16-divisible to avoid Triton recompiles. + score_block_stride = round_up(max_block, 16) + score = torch.empty( + (num_idx_heads, total_q, score_block_stride), + dtype=torch.float32, + device=idx_q.device, + ) + # split-K over seq blocks; chunk count depends only on shape constants so + # the grid is fixed within a cuda graph. + TARGET_GRID = 512 + MAX_NUM_KV_CHUNKS = 256 + # Use the configured max decode length to avoid Triton recompiles when + # switching between qlen=1 and spec-decode verification batches. + BLOCK_SIZE_Q = triton.next_power_of_2(max_decode_query_len) + score_ctas_per_chunk = seq_lens.shape[0] + target = max( + 1, + min(MAX_NUM_KV_CHUNKS, TARGET_GRID // max(1, score_ctas_per_chunk)), + ) + num_kv_chunks = 1 << (target.bit_length() - 1) + grid_score = (seq_lens.shape[0], num_kv_chunks) + _decode_index_score_kernel[grid_score]( + idx_q, + index_kv_cache, + score, + block_table, + seq_lens, + num_idx_heads, + head_dim, + init_blocks, + local_blocks, + decode_query_len, + idx_q.stride(0), + idx_q.stride(1), + idx_q.stride(2), + index_kv_cache.stride(0), + index_kv_cache.stride(1), + index_kv_cache.stride(2), + score.stride(0), + score.stride(1), + score.stride(2), + block_table.stride(0), + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + BLOCK_SIZE_Q=BLOCK_SIZE_Q, + num_kv_chunks=num_kv_chunks, + USE_PDL=use_pdl, + **score_kwargs, + ) + + if out is not None: + topk_idx = out[:, :total_q, :] + else: + topk_idx = torch.empty( + (num_idx_heads, total_q, topk), + dtype=torch.int32, + device=idx_q.device, + ) + # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts + # pow2(num_topk_chunks * pow2(topk)) candidates. + TOPK_TARGET_GRID = 64 + MAX_NUM_TOPK_CHUNKS = 16 + topk_target = max( + 1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads)) + ) + num_topk_chunks = 1 << (topk_target.bit_length() - 1) + block_size_t = triton.next_power_of_2(topk) + chunk_blocks = (max_block + num_topk_chunks - 1) // num_topk_chunks + topk_score_partial = torch.empty( + num_topk_chunks, + num_idx_heads, + batch, + block_size_t, + dtype=torch.float32, + device=idx_q.device, + ) + topk_idx_partial = torch.empty( + num_topk_chunks, + num_idx_heads, + batch, + block_size_t, + dtype=torch.int32, + device=idx_q.device, + ) + _topk_index_partial_kernel[(batch, num_idx_heads, num_topk_chunks)]( + score, + topk_score_partial, + topk_idx_partial, + seq_lens, + SPARSE_BLOCK_SIZE, + topk, + chunk_blocks, + decode_query_len, + score.stride(0), + score.stride(1), + score.stride(2), + topk_score_partial.stride(0), + topk_score_partial.stride(1), + topk_score_partial.stride(2), + topk_score_partial.stride(3), + topk_idx_partial.stride(0), + topk_idx_partial.stride(1), + topk_idx_partial.stride(2), + topk_idx_partial.stride(3), + USE_PDL=use_pdl, + **pdl_kwargs, + ) + _topk_index_merge_kernel[(batch, num_idx_heads)]( + topk_score_partial, + topk_idx_partial, + topk_idx, + seq_lens, + SPARSE_BLOCK_SIZE, + topk, + decode_query_len, + topk_score_partial.stride(0), + topk_score_partial.stride(1), + topk_score_partial.stride(2), + topk_score_partial.stride(3), + topk_idx_partial.stride(0), + topk_idx_partial.stride(1), + topk_idx_partial.stride(2), + topk_idx_partial.stride(3), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + num_topk_chunks=num_topk_chunks, + USE_PDL=use_pdl, + **pdl_kwargs, + ) + return topk_idx diff --git a/vllm/models/minimax_m3/amd/ops/sparse_attn.py b/vllm/models/minimax_m3/amd/ops/sparse_attn.py new file mode 100644 index 00000000000..015a4c516bc --- /dev/null +++ b/vllm/models/minimax_m3/amd/ops/sparse_attn.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm gfx942/gfx950 block-sparse GQA prefill kernel for MiniMax-M3. + +Only the prefill path is specialized on CDNA: each 128-token KV block is split +into SUB_K-token sub-tiles to right-size the per-block QK/PV MFMAs. Everything +else -- the decode split-K kernels, the FP8 dtype set, the sparse block size -- +is reused unchanged from ``common.ops.sparse_attn``. +""" + +import torch + +from vllm.models.minimax_m3.common.ops.sparse_attn import ( + _FP8_DTYPES, + SPARSE_BLOCK_SIZE, + minimax_m3_sparse_attn_decode, +) +from vllm.platforms.rocm import on_gfx950, on_mi3xx +from vllm.triton_utils import tl, triton + +__all__ = ["minimax_m3_sparse_attn", "minimax_m3_sparse_attn_decode"] + + +# Sub-tile width for the prefill kernel's per-block QK/PV GEMMs. gfx950 -> 64, +# gfx942 -> 32 (re-tune with tune_sparse_attn.py). Must divide SPARSE_BLOCK_SIZE. +_SPARSE_ATTN_SUB_K = SPARSE_BLOCK_SIZE // 2 if on_gfx950() else SPARSE_BLOCK_SIZE // 4 + +_SPARSE_ATTN_PREFILL_KWARG: dict | None = None + + +def _sparse_attn_prefill_kwargs() -> dict: + """MFMA + pipeline launch params for the sub-tiled prefill kernel. + + gfx942 and gfx950 share the same params: ``num_warps=1`` keeps one wave + resident on the small per-sub-tile GEMM, ``matrix_instr_nonkdim=16`` / + ``kpack=2`` select the MFMA_16x16 path, and ``num_stages=1`` fits LDS and is + fastest in the sweep. Only the sub-tile width (``_SPARSE_ATTN_SUB_K``) + differs by arch. Empty on other AMD archs. Cached: arch is fixed per process. + """ + global _SPARSE_ATTN_PREFILL_KWARG + if _SPARSE_ATTN_PREFILL_KWARG is None: + kwarg: dict = {} + if on_mi3xx(): + kwarg = { + "num_warps": 1, + "matrix_instr_nonkdim": 16, + "kpack": 2, + "num_stages": 1, + } + _SPARSE_ATTN_PREFILL_KWARG = kwarg + return _SPARSE_ATTN_PREFILL_KWARG + + +# --------------------------------------------------------------------------- +# GQA block-sparse attention (paged). Main heads attend only to the selected +# blocks. BLOCK_SIZE_K == 128 so each selected block is one page. +# --------------------------------------------------------------------------- +# since prefill metadata is sliced from mixed batch metadata, seq_lens and prefix_lens +# might lose pointer alignment, which trigger Triton recompiles. we don't actually +# need pointer alignment for those tensors anyway because we do scalar load. +@triton.heuristics( + { + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]), + "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]), + "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"] + * triton.next_power_of_2(args["gqa_group_size"]), + } +) +@triton.jit(do_not_specialize_on_alignment=["seq_lens", "prefix_lens"]) +def _gqa_sparse_fwd_kernel( + q_ptr, # [total_q, num_heads, head_dim] + kv_cache_ptr, # main cache: [num_blocks, 2, 128, num_kv_heads, head_dim] + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # [total_q, num_heads, head_dim] + block_table_ptr, # [num_reqs, max_blocks] + cu_seqlens_q, + cu_seqblocks_q, + seq_lens, + prefix_lens, + num_kv_heads, + gqa_group_size, + head_dim, + max_topk, + num_q_loop, + sm_scale, + stride_qn, + stride_qh, + stride_qd, + stride_kv_blk, + stride_kv_kv, + stride_kv_pos, + stride_kv_h, + stride_kv_d, + stride_th, + stride_tn, + stride_tk, + stride_on, + stride_oh, + stride_od, + stride_bt_b, + BLOCK_SIZE_Q: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + BLOCK_SIZE_D: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_QH: tl.constexpr, + USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load + SUB_K: tl.constexpr, # CDNA only: KV sub-tile width (see _IS_MI3XX) +): + sm_scale_log2e = sm_scale * 1.4426950409 + pid_q = tl.program_id(0) + pid_kh = tl.program_id(1) + pid_b = tl.program_id(2) + pid_h = pid_kh * gqa_group_size + q_start = tl.load(cu_seqlens_q + pid_b) + q_len = tl.load(cu_seqlens_q + pid_b + 1) - q_start + q_block_start = tl.load(cu_seqblocks_q + pid_b) + q_block_len = tl.load(cu_seqblocks_q + pid_b + 1) - q_block_start + seq_len = tl.load(seq_lens + pid_b) + prefix_len = tl.load(prefix_lens + pid_b) + if pid_q * num_q_loop >= q_block_len: + return + real_q_loop = min(num_q_loop, q_block_len - pid_q * num_q_loop) + bt_row = block_table_ptr + pid_b * stride_bt_b + off_d = tl.arange(0, BLOCK_SIZE_D) + d_mask = off_d < head_dim + for j in range(real_q_loop): + pid_q_j = pid_q * num_q_loop + j + t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th + off_t = tl.arange(0, BLOCK_SIZE_T) + topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1) + real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0) + q_ptrs = tl.make_block_ptr( + base=q_ptr + q_start * stride_qn + pid_h * stride_qh, + shape=(q_len, gqa_group_size, head_dim), + strides=(stride_qn, stride_qh, stride_qd), + offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), + block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), + order=(2, 1, 0), + ) + q = tl.load(q_ptrs, boundary_check=(0, 1, 2), padding_option="zero") + m_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32) + lse_i = tl.full((BLOCK_SIZE_QH,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_QH, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.reshape(q, BLOCK_SIZE_QH, BLOCK_SIZE_D) + + # CDNA: process each 128-token KV block in SUB_K-token sub-tiles so + # each QK/PV MFMA is right-sized. Numerically equivalent to the dense + # path below (flash-softmax reassociation). + NUM_SUB: tl.constexpr = BLOCK_SIZE_K // SUB_K + for _ in tl.range(real_topk): + blk = tl.load(t_ptr_j).to(tl.int32) + t_ptr_j = t_ptr_j + stride_tk + c = blk * BLOCK_SIZE_K + page = tl.load(bt_row + blk).to(tl.int64) + kv_base = kv_cache_ptr + page * stride_kv_blk + pid_kh * stride_kv_h + for sub_i in range(NUM_SUB): + off_sub = tl.arange(0, SUB_K) + sub_i * SUB_K + pos_sub = c + off_sub + pos_mask_sub = pos_sub < seq_len + k_sub = tl.load( + kv_base + + 0 * stride_kv_kv + + off_sub[None, :] * stride_kv_pos + + off_d[:, None] * stride_kv_d, + mask=d_mask[:, None] & pos_mask_sub[None, :], + other=0.0, + ) + if USE_FP8: + k_sub = k_sub.to(q.dtype) + off_q_sub = ( + tl.arange(0, BLOCK_SIZE_Q)[:, None] + + pid_q_j * BLOCK_SIZE_Q + + prefix_len + - off_sub[None, :] + ) + qk_sub = tl.zeros((BLOCK_SIZE_Q, BLOCK_SIZE_H, SUB_K), dtype=tl.float32) + # causal: q_abs_pos - k_off >= block_start (c) + qk_sub += tl.where(off_q_sub[:, None, :] >= c, 0, float("-inf")) + qk_sub = tl.reshape(qk_sub, BLOCK_SIZE_QH, SUB_K) + qk_sub += tl.dot(q, k_sub) * sm_scale_log2e + qk_sub += tl.where(pos_mask_sub[None, :], 0, float("-inf")) + m_ij = tl.maximum(m_i, tl.max(qk_sub, axis=1)) + p_sub = tl.exp2(qk_sub - m_ij[:, None]) + l_ij = tl.sum(p_sub, axis=1) + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + v_sub = tl.load( + kv_base + + 1 * stride_kv_kv + + off_sub[:, None] * stride_kv_pos + + off_d[None, :] * stride_kv_d, + mask=pos_mask_sub[:, None] & d_mask[None, :], + other=0.0, + ) + if USE_FP8: + v_sub = v_sub.to(q.dtype) + acc_o += tl.dot(p_sub.to(v_sub.dtype), v_sub) + m_i = m_ij + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + acc_o = acc_o * tl.exp2(m_i - lse_i)[:, None] + acc_o = tl.reshape(acc_o, BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D) + o_ptrs = tl.make_block_ptr( + base=o_ptr + q_start * stride_on + pid_h * stride_oh, + shape=(q_len, gqa_group_size, head_dim), + strides=(stride_on, stride_oh, stride_od), + offsets=(pid_q_j * BLOCK_SIZE_Q, 0, 0), + block_shape=(BLOCK_SIZE_Q, BLOCK_SIZE_H, BLOCK_SIZE_D), + order=(2, 1, 0), + ) + tl.store(o_ptrs, acc_o.to(o_ptr.dtype.element_ty), boundary_check=(0, 1, 2)) + + +@torch.no_grad() +def minimax_m3_sparse_attn( + q: torch.Tensor, # [total_q, num_heads, head_dim] + kv_cache: torch.Tensor, # [num_blocks, 2, 128, num_kv_heads, head_dim] + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [batch, max_blocks] + cu_seqlens_q: torch.Tensor, # [batch+1] int32 + seq_lens: torch.Tensor, # [batch] int32 + prefix_lens: torch.Tensor, # [batch] int32 + max_query_len: int, + num_kv_heads: int, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] +) -> None: + """GQA block-sparse attention over the selected blocks. block_size_q == 1.""" + total_q, num_heads, head_dim = q.shape + batch = cu_seqlens_q.shape[0] - 1 + topk = topk_idx.shape[-1] + gqa_group_size = num_heads // num_kv_heads + use_fp8 = kv_cache.dtype in _FP8_DTYPES + grid = (max_query_len, num_kv_heads, batch) + _gqa_sparse_fwd_kernel[grid]( + q, + kv_cache, + topk_idx, + output, + block_table, + cu_seqlens_q, + cu_seqlens_q, # cu_seqblocks_q == cu_seqlens_q when block_size_q == 1 + seq_lens, + prefix_lens, + num_kv_heads, + gqa_group_size, + head_dim, + topk, + 1, # num_q_loop + sm_scale, + q.stride(0), + q.stride(1), + q.stride(2), + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + kv_cache.stride(3), + kv_cache.stride(4), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + output.stride(0), + output.stride(1), + output.stride(2), + block_table.stride(0), + BLOCK_SIZE_Q=1, + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + USE_FP8=use_fp8, + SUB_K=_SPARSE_ATTN_SUB_K, + **_sparse_attn_prefill_kwargs(), + ) diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index 4a72b6bc2c9..bb1ed619320 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -27,12 +27,21 @@ from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase -from vllm.models.minimax_m3.common.ops.index_topk import ( - minimax_m3_index_decode, - minimax_m3_index_score, - minimax_m3_index_topk, -) from vllm.platforms import current_platform + +if current_platform.is_rocm(): + from vllm.models.minimax_m3.amd.ops.index_topk import ( + minimax_m3_index_decode, + minimax_m3_index_score, + minimax_m3_index_topk, + ) +else: + from vllm.models.minimax_m3.common.ops.index_topk import ( + minimax_m3_index_decode, + minimax_m3_index_score, + minimax_m3_index_topk, + ) + from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, diff --git a/vllm/models/minimax_m3/common/ops/sparse_attn.py b/vllm/models/minimax_m3/common/ops/sparse_attn.py index 7b6fb73cba9..f04652c89a3 100644 --- a/vllm/models/minimax_m3/common/ops/sparse_attn.py +++ b/vllm/models/minimax_m3/common/ops/sparse_attn.py @@ -31,30 +31,6 @@ _FP8_DTYPES = ( torch.float8_e5m2fnuz, ) -_SPARSE_ATTN_NUM_STAGES_KWARG: dict | None = None - - -def _sparse_attn_num_stages_kwarg() -> dict: - """Triton ``num_stages`` override for the sparse-attn GEMM kernels. - - Forced only where required: CDNA3 (gfx942) caps LDS at - 64 KB, and the default 2-stage pipeline double-buffers the 128x128 K/V tiles - to ~66 KB ("out of resource: shared memory"), so pin gfx942 to a single - stage (~32 KB, which fits). Everywhere else (NVIDIA, CDNA4 gfx950) return an - empty kwarg and let Triton keep its own default -- don't second-guess it. - Cached: the arch is fixed per process. - """ - global _SPARSE_ATTN_NUM_STAGES_KWARG - if _SPARSE_ATTN_NUM_STAGES_KWARG is None: - kwarg: dict = {} - if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx942 - - if on_gfx942(): - kwarg = {"num_stages": 1} - _SPARSE_ATTN_NUM_STAGES_KWARG = kwarg - return _SPARSE_ATTN_NUM_STAGES_KWARG - # --------------------------------------------------------------------------- # GQA block-sparse attention (paged). Main heads attend only to the selected @@ -498,7 +474,6 @@ def minimax_m3_sparse_attn( BLOCK_SIZE_Q=1, BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, USE_FP8=use_fp8, - **_sparse_attn_num_stages_kwarg(), ) @@ -574,7 +549,6 @@ def minimax_m3_sparse_attn_decode( NUM_TOPK_CHUNKS=num_topk_chunks, USE_FP8=use_fp8, USE_PDL=use_pdl, - **_sparse_attn_num_stages_kwarg(), **pdl_launch, ) merge_grid = (total_q, num_heads) diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index 55542230885..88109772eb1 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -24,12 +24,21 @@ from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.forward_context import get_forward_context from vllm.logger import init_logger -from vllm.models.minimax_m3.common.ops.sparse_attn import ( - SPARSE_BLOCK_SIZE, - minimax_m3_sparse_attn, - minimax_m3_sparse_attn_decode, -) +from vllm.models.minimax_m3.common.ops.sparse_attn import SPARSE_BLOCK_SIZE from vllm.platforms import current_platform + +# AMD/ROCm uses the gfx942/gfx950-optimized block-sparse kernels in amd.ops; +# every other platform uses the generic common.ops implementation. +if current_platform.is_rocm(): + from vllm.models.minimax_m3.amd.ops.sparse_attn import ( + minimax_m3_sparse_attn, + minimax_m3_sparse_attn_decode, + ) +else: + from vllm.models.minimax_m3.common.ops.sparse_attn import ( + minimax_m3_sparse_attn, + minimax_m3_sparse_attn_decode, + ) from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index 11fca8e43ab..62275ff2280 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -11,6 +11,10 @@ from functools import cached_property from openai.types.responses import ToolChoiceFunction from pydantic import TypeAdapter, ValidationError +from vllm.entrypoints.chat_utils import ( + get_tool_call_id_type, + make_tool_call_id, +) from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, @@ -24,6 +28,7 @@ from vllm.entrypoints.openai.engine.protocol import ( from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.logger import init_logger from vllm.parser.metrics import record_tool_parser_invocation +from vllm.parser.utils import count_history_tool_calls from vllm.reasoning.abs_reasoning_parsers import ReasoningParser from vllm.sampling_params import StructuredOutputsParams from vllm.tokenizers import TokenizerLike @@ -46,6 +51,7 @@ class StreamState: previous_text: str = "" previous_token_ids: list[int] = field(default_factory=list) history_tool_call_cnt: int = 0 + history_tool_call_cnt_initialized: bool = False tool_call_id_type: str = "random" # only used for "required" and "named tool" choices, # tracks whether function name has been fully returned in the stream yet @@ -108,6 +114,7 @@ class Parser: tokenizer: TokenizerLike, tools: list[Tool] | None = None, *args, + model_config=None, **kwargs, ): self.model_tokenizer = tokenizer @@ -124,7 +131,14 @@ class Parser: self._reasoning_parser is None or self._reasoning_parser.engine_based_streaming ) and (self._tool_parser is None or self._tool_parser.engine_based_streaming) - self._stream_state = StreamState(engine_based=self._engine_based) + self._stream_state = StreamState( + tool_call_id_type=( + get_tool_call_id_type(model_config) + if model_config is not None + else "random" + ), + engine_based=self._engine_based, + ) @cached_property def vocab(self) -> dict[str, int]: @@ -149,6 +163,19 @@ class Parser: def tool_parser(self, parser: ToolParser | None) -> None: self._tool_parser = parser + def _initialize_history_tool_call_cnt( + self, + request: ChatCompletionRequest | ResponsesRequest, + ) -> None: + state = self._stream_state + if state.history_tool_call_cnt_initialized: + return + if state.tool_call_id_type != "kimi_k2": + state.history_tool_call_cnt_initialized = True + return + state.history_tool_call_cnt = count_history_tool_calls(request) + state.history_tool_call_cnt_initialized = True + # ========== Reasoning Parser Methods ========== @abstractmethod @@ -375,6 +402,18 @@ class DelegatingParser(Parser): return request.tool_choice.function.name raise ValueError("Invalid tool_choice for function name extraction.") + def _make_tool_call_id(self, function_name: str) -> str | None: + state = self._stream_state + if state.tool_call_id_type != "kimi_k2": + return None + tool_call_id = make_tool_call_id( + id_type=state.tool_call_id_type, + func_name=function_name, + idx=state.history_tool_call_cnt, + ) + state.history_tool_call_cnt += 1 + return tool_call_id + def _extract_tool_calls( self, content: str | None, @@ -404,9 +443,11 @@ class DelegatingParser(Parser): if is_named_tool_choice and supports_required_and_named: if content is None: return [], None + function_name = self._get_function_name(request) tool_calls.append( FunctionCall( - name=self._get_function_name(request), + id=self._make_tool_call_id(function_name), + name=function_name, arguments=content, ) ) @@ -422,6 +463,7 @@ class DelegatingParser(Parser): for tc in parsed_calls: tool_calls.append( FunctionCall( + id=self._make_tool_call_id(tc.name), name=tc.name, arguments=json.dumps(tc.parameters, ensure_ascii=False), ) @@ -733,6 +775,7 @@ class DelegatingParser(Parser): enable_auto_tools: bool = False, model_output_token_ids: Sequence[int] = (), ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + self._initialize_history_tool_call_cnt(request) reasoning, content = self.extract_reasoning(model_output, request) tool_calls, content = self._extract_tool_calls( content=content, @@ -750,6 +793,7 @@ class DelegatingParser(Parser): *, finished: bool, ) -> DeltaMessage | None: + self._initialize_history_tool_call_cnt(request) state = self._stream_state if not state.prompt_reasoning_checked and prompt_token_ids is not None: diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index 6848a90514c..497eb9039be 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING import regex as re -from vllm.entrypoints.chat_utils import make_tool_call_id +from vllm.entrypoints.chat_utils import get_tool_call_id_type, make_tool_call_id from vllm.entrypoints.openai.engine.protocol import ( DeltaFunctionCall, DeltaMessage, @@ -89,11 +89,18 @@ class ParserEngine(Parser): tools: list[Tool] | None = None, *, parser_engine_config: ParserEngineConfig, + model_config=None, **kwargs, ) -> None: self.model_tokenizer = tokenizer self._tools = tools - self._stream_state = StreamState() + self._stream_state = StreamState( + tool_call_id_type=( + get_tool_call_id_type(model_config) + if model_config is not None + else "random" + ), + ) self._reasoning_parser = None self._tool_parser = None self.parser_engine_config = parser_engine_config @@ -419,6 +426,7 @@ class ParserEngine(Parser): *, finished: bool, ) -> DeltaMessage | None: + self._initialize_history_tool_call_cnt(request) if not self._prompt_streaming_prepared and prompt_token_ids is not None: # NOTE: call the hook BEFORE setting the flag, because the hook # may invoke ``_reset`` (e.g. via ``initialize_streaming``) which @@ -658,6 +666,7 @@ class ParserEngine(Parser): enable_auto_tools: bool = False, model_output_token_ids: Sequence[int] = (), ) -> tuple[str | None, str | None, list[FunctionCall] | None]: + self._initialize_history_tool_call_cnt(request) self._check_skip_tool_parsing(request) reasoning, content, tool_call_info = self._single_pass_parse( model_output, diff --git a/vllm/parser/engine/registered_adapters.py b/vllm/parser/engine/registered_adapters.py index 9d670f30564..c5250abf82a 100644 --- a/vllm/parser/engine/registered_adapters.py +++ b/vllm/parser/engine/registered_adapters.py @@ -13,6 +13,7 @@ from vllm.parser.glm47_moe import Glm47MoeParser from vllm.parser.minimax_m2 import MinimaxM2Parser from vllm.parser.nemotron_v3 import NemotronV3Parser from vllm.parser.qwen3 import Qwen3Parser +from vllm.parser.seed_oss import SeedOssParser ( MinimaxM2ParserReasoningAdapter, @@ -34,6 +35,11 @@ from vllm.parser.qwen3 import Qwen3Parser Qwen3ParserToolAdapter, ) = make_adapters(Qwen3Parser) +( + SeedOssParserReasoningAdapter, + SeedOssParserToolAdapter, +) = make_adapters(SeedOssParser) + ( Glm47MoeParserReasoningAdapter, Glm47MoeParserToolAdapter, diff --git a/vllm/parser/minimax_m2.py b/vllm/parser/minimax_m2.py index d348d5779b4..86fa1d1bad1 100644 --- a/vllm/parser/minimax_m2.py +++ b/vllm/parser/minimax_m2.py @@ -38,13 +38,21 @@ INVOKE_END = "" NAME_END_DQ = '">' NAME_END_SQ = "'>" NAME_END_UNQUOTED = ">" - +PARAM_START = "" _PARAM_RE = re.compile( r"<\s*parameter\s+name\s*=\s*" r"(?:\"(?P[^\"]*)\"|'(?P[^']*)'|(?P[^>\s]+))" r"\s*>" r"(?P.*?)" - r"<\s*/\s*parameter\s*>", + r"(?:<\s*/\s*parameter\s*>|(?=<\s*parameter\s+name\s*=))", + re.DOTALL, +) +_PARTIAL_PARAM_RE = re.compile( + r"<\s*parameter\s+name\s*=\s*" + r"(?:\"(?P[^\"]*)\"|'(?P[^']*)'|(?P[^>\s]+))" + r"\s*>" + r"(?P.*)$", re.DOTALL, ) @@ -63,6 +71,19 @@ def _minimax_m2_arg_converter(raw_args: str, partial: bool) -> str: continue params[name] = match.group("value").strip() + if partial: + remaining = _PARAM_RE.sub("", raw_args) + match = _PARTIAL_PARAM_RE.search(remaining) + if match: + name = ( + match.group("dq_name") + or match.group("sq_name") + or match.group("bare_name") + or "" + ).strip() + if name: + params[name] = match.group("value").strip() + return json.dumps(params, ensure_ascii=False) @@ -75,6 +96,8 @@ def minimax_m2_config() -> ParserEngineConfig: "THINK_START": THINK_START, "THINK_END": THINK_END, "TOOL_START": TOOL_CALL_START, + "PARAM_START": PARAM_START, + "PARAM_END": PARAM_END, "TOOL_END": TOOL_CALL_END, "INVOKE_PREFIX_DQ": INVOKE_PREFIX_DQ, "INVOKE_PREFIX_SQ": INVOKE_PREFIX_SQ, @@ -111,6 +134,14 @@ def minimax_m2_config() -> ParserEngineConfig: ParserState.TOOL_PREAMBLE, (), ), + (ParserState.TOOL_ARGS, "PARAM_START"): Transition( + ParserState.TOOL_ARGS, + (EventType.ARG_VALUE_CHUNK,), + ), + (ParserState.TOOL_ARGS, "PARAM_END"): Transition( + ParserState.TOOL_ARGS, + (EventType.ARG_VALUE_CHUNK,), + ), (ParserState.TOOL_PREAMBLE, "TOOL_END"): Transition( ParserState.CONTENT, (), diff --git a/vllm/parser/qwen3.py b/vllm/parser/qwen3.py index f14da8234c5..f80aa6ff7a2 100644 --- a/vllm/parser/qwen3.py +++ b/vllm/parser/qwen3.py @@ -38,6 +38,8 @@ if TYPE_CHECKING: from vllm.tokenizers import TokenizerLike from vllm.tool_parsers.abstract_tool_parser import Tool +THINK_START = "" +THINK_END = "" TOOL_CALL_START = "" TOOL_CALL_END = "" FUNC_PREFIX = " str: @functools.cache -def qwen3_config(thinking: bool = True) -> ParserEngineConfig: +def qwen3_config( + thinking: bool = True, + *, + name: str = "qwen3", + think_start: str = THINK_START, + think_end: str = THINK_END, + tool_start: str = TOOL_CALL_START, + tool_end: str = TOOL_CALL_END, +) -> ParserEngineConfig: return ParserEngineConfig( - name="qwen3", + name=name, initial_state=ParserState.REASONING if thinking else ParserState.CONTENT, terminals={ # Reasoning terminals - "THINK_START": "", - "THINK_END": "", + "THINK_START": think_start, + "THINK_END": think_end, # Tool call terminals - "TOOL_START": TOOL_CALL_START, - "TOOL_END": TOOL_CALL_END, + "TOOL_START": tool_start, + "TOOL_END": tool_end, "FUNC_PREFIX": FUNC_PREFIX, "FUNC_END": FUNC_END, "PARAM_START": PARAM_START, @@ -93,10 +103,10 @@ def qwen3_config(thinking: bool = True) -> ParserEngineConfig: "CLOSE_ANGLE": ">", }, token_id_terminals={ - "THINK_START": "", - "THINK_END": "", - "TOOL_START": TOOL_CALL_START, - "TOOL_END": TOOL_CALL_END, + "THINK_START": think_start, + "THINK_END": think_end, + "TOOL_START": tool_start, + "TOOL_END": tool_end, }, transitions={ # -- Reasoning transitions -- @@ -185,8 +195,18 @@ class Qwen3Parser(ParserEngine): - ```` as implicit reasoning end - Unpaired ```` token ID detection for ``is_reasoning_end`` + + Subclasses that share the grammar but differ only in the four wrapper + token strings (reasoning + tool-call) override the class attributes + below; everything else is inherited unchanged. """ + CONFIG_NAME = "qwen3" + THINK_START = THINK_START + THINK_END = THINK_END + TOOL_START = TOOL_CALL_START + TOOL_END = TOOL_CALL_END + def __init__( self, tokenizer: TokenizerLike, @@ -197,7 +217,14 @@ class Qwen3Parser(ParserEngine): self.thinking_enabled = chat_kwargs.get("enable_thinking", True) kwargs.setdefault( "parser_engine_config", - qwen3_config(thinking=self.thinking_enabled), + qwen3_config( + thinking=self.thinking_enabled, + name=self.CONFIG_NAME, + think_start=self.THINK_START, + think_end=self.THINK_END, + tool_start=self.TOOL_START, + tool_end=self.TOOL_END, + ), ) super().__init__( tokenizer, @@ -205,8 +232,8 @@ class Qwen3Parser(ParserEngine): **kwargs, ) vocab = self.vocab - self._tool_call_token_id: int | None = vocab.get("") - self._tool_call_end_token_id: int | None = vocab.get("") + self._tool_call_token_id: int | None = vocab.get(self.TOOL_START) + self._tool_call_end_token_id: int | None = vocab.get(self.TOOL_END) def extract_reasoning( self, diff --git a/vllm/parser/seed_oss.py b/vllm/parser/seed_oss.py new file mode 100644 index 00000000000..2f709f0ad67 --- /dev/null +++ b/vllm/parser/seed_oss.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""seed_oss parser for tool calls and reasoning. + +seed_oss shares the Qwen3 XML grammar exactly; only the four wrapper +token strings differ:: + + -> + -> + -> + -> + +```` and ```` are byte-identical, so the +entire transition table and ``_qwen3_arg_converter`` are inherited from +:class:`Qwen3Parser` unchanged. +""" + +from __future__ import annotations + +from vllm.parser.qwen3 import Qwen3Parser + + +class SeedOssParser(Qwen3Parser): + CONFIG_NAME = "seed_oss" + THINK_START = "" + THINK_END = "" + TOOL_START = "" + TOOL_END = "" diff --git a/vllm/parser/utils.py b/vllm/parser/utils.py new file mode 100644 index 00000000000..51382cd2909 --- /dev/null +++ b/vllm/parser/utils.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable, Sequence + +from openai.types.responses import ResponseFunctionToolCall + +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.responses.protocol import ( + ResponseInputOutputItem, + ResponsesRequest, +) + + +def count_tool_calls(tool_calls: object) -> int: + if tool_calls is None: + return 0 + if isinstance(tool_calls, (str, bytes, dict)): + return 1 + if isinstance(tool_calls, Iterable): + return sum(1 for _ in tool_calls) + return 1 + + +def count_chat_history_tool_calls( + messages: Sequence[ChatCompletionMessageParam], +) -> int: + return sum( + count_tool_calls(msg.get("tool_calls")) + for msg in messages + if isinstance(msg, dict) and msg.get("role") == "assistant" + ) + + +def count_response_history_tool_calls( + response_items: Sequence[ResponseInputOutputItem], +) -> int: + count = 0 + for item in response_items: + if isinstance(item, ResponseFunctionToolCall): + count += 1 + continue + + if isinstance(item, dict): + item_type = item.get("type") + if item_type == "function_call": + count += 1 + elif item.get("role") == "assistant": + count += count_tool_calls(item.get("tool_calls")) + + return count + + +def count_history_tool_calls( + request: ChatCompletionRequest | ResponsesRequest, +) -> int: + if isinstance(request, ChatCompletionRequest): + return count_chat_history_tool_calls(request.messages) + + request_input = request.input + if isinstance(request_input, str): + return 0 + + return count_response_history_tool_calls(request_input) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index c529af46df9..c5d7ec2fe71 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -154,23 +154,9 @@ class CpuPlatform(Platform): parallel_config.worker_cls = "vllm.v1.worker.cpu_worker.CPUWorker" # Disable DBO if parallel_config.enable_dbo: - logger.warning("Dual-Batch Overlap is not supported on CPU, disabled.") + logger.warning_once("Dual-Batch Overlap is not supported on CPU, disabled.") parallel_config.enable_dbo = False - if torch.cpu._is_amx_tile_supported() and ( - model_config is not None - and model_config.get_num_layers_by_block_type( - parallel_config, "linear_attention" - ) - > 0 - ): - cache_config.enable_prefix_caching = False - scheduler_config.enable_chunked_prefill = False - logger.warning( - "Disabled unsupported prefix caching and chunked prefill " - "for linear attention on AMX CPU platforms." - ) - # Note: workaround for v1 gpu_model_runner from vllm.config import CompilationMode @@ -309,7 +295,7 @@ class CpuPlatform(Platform): ) if model_config is not None and model_config.use_mla: - logger.info( + logger.info_once( "MLA is enabled on a non-GPU platform; forcing chunked " "prefill and prefix caching to be disabled." ) @@ -431,13 +417,13 @@ class CpuPlatform(Platform): try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %r", e) else: try: import vllm._C_AVX512 # noqa: F401 except ImportError as e: if ignored_msg not in e.msg: - logger.warning( + logger.warning_once( "Failed to import from vllm._C_AVX512: %r", e ) else: @@ -445,12 +431,12 @@ class CpuPlatform(Platform): import vllm._C_AVX2 # noqa: F401 except ImportError as e: if ignored_msg not in e.msg: - logger.warning("Failed to import from vllm._C_AVX2: %r", e) + logger.warning_once("Failed to import from vllm._C_AVX2: %r", e) else: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %r", e) @classmethod def pack_kv_cache( diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index fa96cb8c946..ee73eef8797 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -220,15 +220,17 @@ class CudaPlatformBase(Platform): try: import vllm._C_stable_libtorch # noqa: F401 except ImportError as e: - logger.warning("Failed to import from vllm._C_stable_libtorch: %r", e) + logger.warning_once("Failed to import from vllm._C_stable_libtorch: %r", e) try: import vllm._moe_C_stable_libtorch # noqa: F401 except ImportError as e: - logger.warning("Failed to import from vllm._moe_C_stable_libtorch: %r", e) + logger.warning_once( + "Failed to import from vllm._moe_C_stable_libtorch: %r", e + ) try: import vllm._qutlass_C # noqa: F401 except ImportError as e: - logger.warning("Failed to import from vllm._qutlass_C: %r", e) + logger.warning_once("Failed to import from vllm._qutlass_C: %r", e) @property def supported_dtypes(self) -> list[torch.dtype]: @@ -291,7 +293,7 @@ class CudaPlatformBase(Platform): # kernel with limited pinned memory support for CUDA. version = _get_wsl_kernel_version() if version is None or version < (4, 19, 121): - logger.warning( + logger.warning_once( "Using 'pin_memory=False' as WSL is detected and the " "WSL2 kernel version is below 4.19.121. This may slow " "down performance. Please run `wsl --update`." @@ -320,7 +322,7 @@ class CudaPlatformBase(Platform): and scheduler_config.is_multimodal_model and not scheduler_config.disable_chunked_mm_input ): - logger.warning( + logger.warning_once( "Forcing --disable_chunked_mm_input for models " "with multimodal-bidirectional attention." ) @@ -331,7 +333,7 @@ class CudaPlatformBase(Platform): and vllm_config.offload_config.uva.cpu_offload_gb > 0 and bool(vllm_config.compilation_config.cudagraph_mode) ): - logger.warning( + logger.warning_once( "--cpu-offload-gb is enabled with CUDA graphs on WSL2. " "This combination requires pinned (page-locked) memory " "allocations. WARNING: Windows (WDDM) enforces a hard " diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 82c87416093..a7a0dd52df7 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -355,7 +355,7 @@ class Platform: try: import vllm._C # noqa: F401 except ImportError as e: - logger.warning("Failed to import from vllm._C: %r", e) + logger.warning_once("Failed to import from vllm._C: %r", e) with contextlib.suppress(ImportError): import vllm._moe_C_stable_libtorch # noqa: F401 @@ -859,7 +859,7 @@ class Platform: # Pinned memory support under WSL depends on the vendor and driver # version. Conservative default: return False. Platform subclasses # that can verify support (e.g. CudaPlatformBase) override this. - logger.warning( + logger.warning_once( "Using 'pin_memory=False' as WSL is detected. " "This may slow down performance." ) @@ -1003,7 +1003,7 @@ class Platform: if attr is not None: return attr - logger.warning( + logger.warning_once( "Current platform %s does not have '%s' attribute.", self.device_type, key, diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 06953d504b6..6c3a0fe96ec 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -119,6 +119,14 @@ def _sync_hip_cuda_env_vars(): hip_val = os.environ.get("HIP_VISIBLE_DEVICES") or None cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES") or None + if cuda_val is not None: + logger.warning_once( + "Using CUDA_VISIBLE_DEVICES on ROCm is deprecated and support " + "will be removed in vLLM v0.26.0. Please use HIP_VISIBLE_DEVICES " + "instead.", + scope="process", + ) + if hip_val is not None and cuda_val is not None: if hip_val != cuda_val: raise ValueError( @@ -453,7 +461,6 @@ class RocmPlatform(Platform): "auto_awq", "awq_marlin", # will be overwritten with awq "gptq", - "gptq_marlin", "auto_gptq", "fp8", "deepseek_v4_fp8", diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py index 030b4933bb6..94f5e8e5a89 100644 --- a/vllm/platforms/xpu.py +++ b/vllm/platforms/xpu.py @@ -55,7 +55,7 @@ class XPUPlatform(Platform): from vllm.v1.attention.backends.utils import set_kv_cache_layout set_kv_cache_layout("NHD") - logger.info( + logger.info_once( "Setting VLLM_KV_CACHE_LAYOUT to 'NHD' for XPU; " "only NHD layout is supported by XPU attention kernels." ) @@ -91,7 +91,7 @@ class XPUPlatform(Platform): f"with use_mla: {attn_selector_config.use_mla}" ) - logger.info("Using Flash Attention backend.") + logger.info_once("Using Flash Attention backend.") return AttentionBackendEnum.FLASH_ATTN.get_path() @classmethod @@ -193,13 +193,13 @@ class XPUPlatform(Platform): if not supports_xpu_graph(): compilation_config.cudagraph_mode = CUDAGraphMode.NONE - logger.warning( + logger.warning_once( "XPU Graph is not supported in the current PyTorch version, " "disabling cudagraph_mode." ) elif not envs.VLLM_XPU_ENABLE_XPU_GRAPH: compilation_config.cudagraph_mode = CUDAGraphMode.NONE - logger.warning( + logger.warning_once( "XPU Graph is disabled by environment variable, " "please set VLLM_XPU_ENABLE_XPU_GRAPH=1 to enable it." ) @@ -218,7 +218,7 @@ class XPUPlatform(Platform): if compilation_config.mode != CompilationMode.NONE: for flag, feature_name in fusion_passes_to_disable.items(): if getattr(pass_config, flag): - logger.warning( + logger.warning_once( "Feature %r is not yet supported on XPU and will be disabled.", feature_name, ) diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index cbb1fa350f5..fc74cf2f3f7 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -117,8 +117,8 @@ _REASONING_PARSERS_TO_REGISTER = { "Qwen3ParserReasoningAdapter", ), "seed_oss": ( - "seedoss_reasoning_parser", - "SeedOSSReasoningParser", + "seed_oss_engine_reasoning_parser", + "SeedOssParserReasoningAdapter", ), "step3": ( "step3_reasoning_parser", diff --git a/vllm/reasoning/seed_oss_engine_reasoning_parser.py b/vllm/reasoning/seed_oss_engine_reasoning_parser.py new file mode 100644 index 00000000000..e651d411f43 --- /dev/null +++ b/vllm/reasoning/seed_oss_engine_reasoning_parser.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.engine.registered_adapters import SeedOssParserReasoningAdapter + +__all__ = ["SeedOssParserReasoningAdapter"] diff --git a/vllm/reasoning/seedoss_reasoning_parser.py b/vllm/reasoning/seedoss_reasoning_parser.py deleted file mode 100644 index d3d4d8ec074..00000000000 --- a/vllm/reasoning/seedoss_reasoning_parser.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - - -from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser - - -class SeedOSSReasoningParser(BaseThinkingReasoningParser): - """ - Reasoning parser for SeedOSS model. - - The SeedOSS model uses ... tokens to - denote reasoning content text. This parser extracts - the reasoning content from the model output. - Similar to DeepSeek R1, it supports cases - where the model doesn't generate the start token. - """ - - @property - def start_token(self) -> str: - """The token that starts reasoning content.""" - return "" - - @property - def end_token(self) -> str: - """The token that ends reasoning content.""" - return "" diff --git a/vllm/scalar_type.py b/vllm/scalar_type.py index 05760f3f829..db52e93465c 100644 --- a/vllm/scalar_type.py +++ b/vllm/scalar_type.py @@ -348,6 +348,9 @@ class scalar_types: uint2b2 = ScalarType.uint(2, 2) uint3b4 = ScalarType.uint(3, 4) uint4b8 = ScalarType.uint(4, 8) + uint5b16 = ScalarType.uint(5, 16) + uint6b32 = ScalarType.uint(6, 32) + uint7b64 = ScalarType.uint(7, 64) uint8b128 = ScalarType.uint(8, 128) # colloquial names diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 109189a033a..b9a9c9ad07b 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -163,8 +163,8 @@ _TOOL_PARSERS_TO_REGISTER = { "Qwen3EngineToolParser", ), "seed_oss": ( - "seed_oss_tool_parser", - "SeedOssToolParser", + "seed_oss_engine_tool_parser", + "SeedOssEngineToolParser", ), "step3": ( "step3_tool_parser", diff --git a/vllm/tool_parsers/rust_tool_parser.py b/vllm/tool_parsers/rust_tool_parser.py index 493f765a2c2..05f015369f8 100644 --- a/vllm/tool_parsers/rust_tool_parser.py +++ b/vllm/tool_parsers/rust_tool_parser.py @@ -224,7 +224,7 @@ class RustToolParser(ToolParser): "Error parsing %s tool call output.", self.rust_parser_name ) return None - return output.coalesce_calls(), tool_call_ids + return output.coalesce(), tool_call_ids def extract_tool_calls( self, diff --git a/vllm/tool_parsers/seed_oss_engine_tool_parser.py b/vllm/tool_parsers/seed_oss_engine_tool_parser.py new file mode 100644 index 00000000000..e708afd1710 --- /dev/null +++ b/vllm/tool_parsers/seed_oss_engine_tool_parser.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.engine.registered_adapters import SeedOssParserToolAdapter + + +class SeedOssEngineToolParser(SeedOssParserToolAdapter): # type: ignore[valid-type, misc] + structural_tag_model = None diff --git a/vllm/tool_parsers/seed_oss_tool_parser.py b/vllm/tool_parsers/seed_oss_tool_parser.py deleted file mode 100644 index 42e4ba77691..00000000000 --- a/vllm/tool_parsers/seed_oss_tool_parser.py +++ /dev/null @@ -1,633 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# Adapted from qwen3coder xml parser, All rights reserved. -# ruff: noqa: E501 - -import json -import uuid -from collections.abc import Sequence - -import regex as re - -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, -) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ExtractedToolCallInformation, - FunctionCall, - ToolCall, -) -from vllm.logger import init_logger -from vllm.tokenizers import TokenizerLike -from vllm.tool_parsers.abstract_tool_parser import ( - Tool, - ToolParser, -) -from vllm.tool_parsers.utils import ( - coerce_to_schema_type, - extract_types_from_schema, - find_tool_properties, -) - -logger = init_logger(__name__) - - -class SeedOssToolParser(ToolParser): - TOOL_CALL_START = "" - TOOL_CALL_END = "" - - def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): - super().__init__(tokenizer, tools) - - # --- streaming state --- - self._reset_streaming_state() - self.prev_tool_call_arr: list[dict] = [] - - self.tool_call_start_token: str = self.TOOL_CALL_START - self.tool_call_end_token: str = self.TOOL_CALL_END - # Sentinel tokens for streaming mode - self.tool_call_prefix: str = " or its closing tag." - ) - - tool_start_re = re.escape(self.tool_call_start_token) - tool_end_re = re.escape(self.tool_call_end_token) - - self.tool_call_complete_regex = re.compile( - rf"{tool_start_re}(.*?){tool_end_re}", re.DOTALL - ) - self.tool_call_regex = re.compile( - rf"{tool_start_re}(.*?){tool_end_re}|{tool_start_re}(.*?)$", re.DOTALL - ) - - self.tool_call_function_regex = re.compile( - r"|| str: - """Generate a unique tool call ID.""" - return f"call_{uuid.uuid4().hex[:24]}" - - def _reset_streaming_state(self): - """Reset all streaming state.""" - self.current_tool_index = 0 - self.is_tool_call_started = False - self.header_sent = False - self.current_tool_id = -1 - self.current_function_name = None - self.current_param_name = None - self.current_param_value = "" - self.param_count = 0 - self.in_param = False - self.in_function = False - self.accumulated_text = "" - self.json_started = False - self.json_closed = False - - def _parse_xml_function_call( - self, function_call_str: str, tools: list[Tool] | None - ) -> ToolCall | None: - # Extract function name - end_index = function_call_str.index(">") - function_name = function_call_str[:end_index] - tool_properties = find_tool_properties(tools, function_name) - parameters = function_call_str[end_index + 1 :] - param_dict = {} - for match in self.tool_call_parameter_regex.findall(parameters): - match_text = match[0] if match[0] else match[1] - idx = match_text.index(">") - param_name = match_text[:idx] - param_value = str(match_text[idx + 1 :]) - # Remove prefix and trailing \n - if param_value.startswith("\n"): - param_value = param_value[1:] - if param_value.endswith("\n"): - param_value = param_value[:-1] - - param_types = extract_types_from_schema(tool_properties.get(param_name, {})) - param_dict[param_name] = coerce_to_schema_type(param_value, param_types) - return ToolCall( - type="function", - function=FunctionCall( - name=function_name, arguments=json.dumps(param_dict, ensure_ascii=False) - ), - ) - - def _get_function_calls(self, model_output: str) -> list[str]: - # Find all tool calls - matched_ranges = self.tool_call_regex.findall(model_output) - raw_tool_calls = [ - match[0] if match[0] else match[1] for match in matched_ranges - ] - - # Back-off strategy if no tool_call tags found - if len(raw_tool_calls) == 0: - raw_tool_calls = [model_output] - - raw_function_calls = [] - for tool_call in raw_tool_calls: - raw_function_calls.extend(self.tool_call_function_regex.findall(tool_call)) - - function_calls = [ - match[0] if match[0] else match[1] for match in raw_function_calls - ] - return function_calls - - def extract_tool_calls( - self, - model_output: str, - request: ChatCompletionRequest, - ) -> ExtractedToolCallInformation: - # Quick check to avoid unnecessary processing - if self.tool_call_prefix not in model_output: - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - - # Check if both think start and end tokens are present - if ( - self.think_start_token in model_output - and self.think_end_token in model_output - ): - # Find the position of think end token - think_end_index = model_output.find(self.think_end_token) + len( - self.think_end_token - ) - # Extract content after think end token - result_content = model_output[think_end_index:] - thinking_content = model_output[:think_end_index] - else: - thinking_content = "" - result_content = model_output - - try: - function_calls = self._get_function_calls(result_content) - if len(function_calls) == 0: - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - - tool_calls = [ - self._parse_xml_function_call(function_call_str, self.tools) - for function_call_str in function_calls - ] - - # Populate prev_tool_call_arr for serving layer to set finish_reason - self.prev_tool_call_arr.clear() # Clear previous calls - for tool_call in tool_calls: - if tool_call: - self.prev_tool_call_arr.append( - { - "name": tool_call.function.name, - "arguments": tool_call.function.arguments, - } - ) - - # Extract content before tool calls - tool_call_start_index = result_content.find(self.tool_call_start_token) - tool_call_start_index = ( - tool_call_start_index - if tool_call_start_index >= 0 - else result_content.find(self.tool_call_prefix) - ) - content = thinking_content + result_content[:tool_call_start_index] - - return ExtractedToolCallInformation( - tools_called=(len(tool_calls) > 0), - tool_calls=tool_calls, - content=content if content else None, - ) - - except Exception: - logger.exception("Error in extracting tool call from response.") - return ExtractedToolCallInformation( - tools_called=False, tool_calls=[], content=model_output - ) - - def extract_tool_calls_streaming( - self, - previous_text: str, - current_text: str, - delta_text: str, - previous_token_ids: Sequence[int], - current_token_ids: Sequence[int], - delta_token_ids: Sequence[int], - request: ChatCompletionRequest, - ) -> DeltaMessage | None: - # If no delta text, return None unless - # it's an EOS token after tool calls - if not delta_text: - # Check if this is an EOS token after all tool calls are complete - # We check for tool calls in the text even if is_tool_call_started - # is False because it might have been reset after processing all tools - if delta_token_ids and self.tool_call_end_token_id not in delta_token_ids: - # Count complete tool calls - complete_calls = len( - self.tool_call_complete_regex.findall(current_text) - ) - - # If we have completed tool calls and populated prev_tool_call_arr - if complete_calls > 0 and len(self.prev_tool_call_arr) > 0: - # Check if all tool calls are closed - open_calls = current_text.count( - self.tool_call_start_token - ) - current_text.count(self.tool_call_end_token) - if open_calls == 0: - # Return empty delta message to allow finish_reason processing - return DeltaMessage(content="") - elif not self.is_tool_call_started and current_text: - # This is a regular content response that's now complete - return DeltaMessage(content="") - return None - - # Check if this is the first call (reset state if needed) - if not previous_text: - self._reset_streaming_state() - - # Update accumulated text - self.accumulated_text = current_text - - # Check if we need to advance to next tool - if self.json_closed and not self.in_function: - # Check if this tool call has ended - tool_ends = current_text.count(self.tool_call_end_token) - if tool_ends > self.current_tool_index: - # This tool has ended, advance to next - self.current_tool_index += 1 - self.header_sent = False - self.param_count = 0 - self.json_started = False - self.json_closed = False - - # Check if there are more tool calls - if self.current_tool_index >= current_text.count( - self.tool_call_start_token - ): - # No more tool calls - self.is_tool_call_started = False - # Continue processing next tool - return None - - # Check if end thinking - if not self.is_thinking_end and ( - self.think_end_token_id in delta_token_ids - or self.think_end_token in delta_text - ): - self.is_thinking_end = True - - # If thinking hasn't ended yet, don't process any tool calls - if not self.is_thinking_end: - return DeltaMessage(content=delta_text) - - # Handle normal content before tool calls - if not self.is_tool_call_started: - # Check if tool call is starting - if ( - self.tool_call_start_token_id in delta_token_ids - or self.tool_call_start_token in delta_text - ): - self.is_tool_call_started = True - # Return any content before the tool call - if self.tool_call_start_token in delta_text: - content_before = delta_text[ - : delta_text.index(self.tool_call_start_token) - ] - if content_before: - return DeltaMessage(content=content_before) - return None - else: - # Check if we're between tool calls - skip whitespace - if ( - current_text.rstrip().endswith(self.tool_call_end_token) - and delta_text.strip() == "" - ): - # We just ended a tool call, skip whitespace - return None - # Normal content, no tool call - return DeltaMessage(content=delta_text) - - # Check if we're between tool calls (waiting for next one) - # Count tool calls we've seen vs processed - tool_starts_count = current_text.count(self.tool_call_start_token) - if self.current_tool_index >= tool_starts_count: - # We're past all tool calls, shouldn't be here - return None - - # We're in a tool call, find the current tool call portion - # Need to find the correct tool call based on current_tool_index - # Only process tool calls after think_end_token - think_end_index = ( - current_text.find(self.think_end_token) + len(self.think_end_token) - if self.think_end_token in current_text - else 0 - ) - tool_starts: list[int] = [] - idx = think_end_index - while True: - idx = current_text.find(self.tool_call_start_token, idx) - if idx == -1: - break - tool_starts.append(idx) - idx += len(self.tool_call_start_token) - - if self.current_tool_index >= len(tool_starts): - # No more tool calls to process yet - return None - - tool_start_idx = tool_starts[self.current_tool_index] - # Find where this tool call ends (or current position if not ended yet) - tool_end_idx = current_text.find(self.tool_call_end_token, tool_start_idx) - if tool_end_idx == -1: - tool_text = current_text[tool_start_idx:] - else: - tool_text = current_text[ - tool_start_idx : tool_end_idx + len(self.tool_call_end_token) - ] - - # Looking for function header - if not self.header_sent: - if self.tool_call_prefix in tool_text: - func_start = tool_text.find(self.tool_call_prefix) + len( - self.tool_call_prefix - ) - func_end = tool_text.find(">", func_start) - - if func_end != -1: - # Found complete function name - self.current_function_name = tool_text[func_start:func_end] - self.current_tool_id = self._generate_tool_call_id() # type: ignore - self.header_sent = True - self.in_function = True - - # IMPORTANT: Add to prev_tool_call_arr immediately when we detect a tool call - # This ensures finish_reason="tool_calls" even if parsing isn't complete - already_added = any( - tool.get("name") == self.current_function_name - for tool in self.prev_tool_call_arr - ) - if not already_added: - self.prev_tool_call_arr.append( - { - "name": self.current_function_name, - "arguments": "{}", # Placeholder, will be updated later - } - ) - - # Send header with function info - return DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - id=self.current_tool_id, - function=DeltaFunctionCall( - name=self.current_function_name, arguments="" - ), - type="function", - ) - ] - ) - return None - - # We've sent header, now handle function body - if self.in_function: - # Send opening brace if not sent yet - if not self.json_started and self.parameter_prefix not in delta_text: - self.json_started = True - return DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - function=DeltaFunctionCall(arguments="{"), - ) - ] - ) - - # Make sure json_started is set if we're processing parameters - if not self.json_started: - self.json_started = True - - # Check for function end in accumulated text - if not self.json_closed and self.function_end_token in tool_text: - # Close JSON - self.json_closed = True - - # Extract the complete tool call to update prev_tool_call_arr with final arguments - # Find the function content - func_start = tool_text.find(self.tool_call_prefix) + len( - self.tool_call_prefix - ) - func_content_end = tool_text.find(self.function_end_token, func_start) - if func_content_end != -1: - func_content = tool_text[func_start:func_content_end] - # Parse to get the complete arguments - try: - parsed_tool = self._parse_xml_function_call( - func_content, self.tools - ) - if parsed_tool: - # Update existing entry in prev_tool_call_arr with complete arguments - for i, tool in enumerate(self.prev_tool_call_arr): - if tool.get("name") == parsed_tool.function.name: - self.prev_tool_call_arr[i]["arguments"] = ( - parsed_tool.function.arguments - ) - break - except Exception: - logger.warning( - "Failed to parse tool arguments during streaming.", - exc_info=True, - ) - - result = DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - function=DeltaFunctionCall(arguments="}"), - ) - ] - ) - - # Reset state for next tool - self.in_function = False - self.json_closed = True - - return result - - # Look for parameters - # Count how many complete parameters we have processed - complete_params = tool_text.count(self.parameter_end_token) - - # Check if we should start a new parameter - if not self.in_param and self.param_count < complete_params: - # Find the unprocessed parameter - # Count parameter starts - param_starts = [] - idx = 0 - while True: - idx = tool_text.find(self.parameter_prefix, idx) - if idx == -1: - break - param_starts.append(idx) - idx += len(self.parameter_prefix) - - if len(param_starts) > self.param_count: - # Process the next parameter - param_idx = param_starts[self.param_count] - param_start = param_idx + len(self.parameter_prefix) - remaining = tool_text[param_start:] - - if ">" in remaining: - # We have the complete parameter name - name_end = remaining.find(">") - self.current_param_name = remaining[:name_end] - - # Find the parameter value - value_start = param_start + name_end + 1 - value_text = tool_text[value_start:] - if value_text.startswith("\n"): - value_text = value_text[1:] - - # Find where this parameter ends - param_end_idx = value_text.find(self.parameter_end_token) - if param_end_idx != -1: - # Complete parameter found - param_value = value_text[:param_end_idx] - if param_value.endswith("\n"): - param_value = param_value[:-1] - - # Build complete JSON fragment for this parameter - if self.param_count == 0: - json_fragment = ( - '"' - + self.current_param_name - + '": "' - + json.dumps(param_value, ensure_ascii=False)[1:-1] - + '"' - ) - else: - json_fragment = ( - ', "' - + self.current_param_name - + '": "' - + json.dumps(param_value, ensure_ascii=False)[1:-1] - + '"' - ) - - self.param_count += 1 - - return DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - function=DeltaFunctionCall( - arguments=json_fragment - ), - ) - ] - ) - - # Continue parameter value - if self.in_param: - if self.parameter_end_token in delta_text: - # End of parameter - end_idx = delta_text.find(self.parameter_end_token) - value_chunk = delta_text[:end_idx] - - # Skip past > if at start - if not self.current_param_value and ">" in value_chunk: - gt_idx = value_chunk.find(">") - value_chunk = value_chunk[gt_idx + 1 :] - - if not self.current_param_value and value_chunk.startswith("\n"): - value_chunk = value_chunk[1:] - - # Calculate incremental JSON - full_value = self.current_param_value + value_chunk - prev_escaped = ( - json.dumps(self.current_param_value, ensure_ascii=False)[1:-1] - if self.current_param_value - else "" - ) - full_escaped = json.dumps(full_value, ensure_ascii=False)[1:-1] - delta_escaped = full_escaped[len(prev_escaped) :] - - self.in_param = False - self.current_param_value = "" - - return DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - function=DeltaFunctionCall( - arguments=delta_escaped + '"' - ), - ) - ] - ) - else: - # Continue accumulating value - value_chunk = delta_text - - # Handle first chunk after param name - if not self.current_param_value and ">" in value_chunk: - gt_idx = value_chunk.find(">") - value_chunk = value_chunk[gt_idx + 1 :] - - if not self.current_param_value and value_chunk.startswith("\n"): - value_chunk = value_chunk[1:] - - if value_chunk: - # Stream the escaped delta - prev_escaped = ( - json.dumps(self.current_param_value, ensure_ascii=False)[ - 1:-1 - ] - if self.current_param_value - else "" - ) - self.current_param_value += value_chunk - full_escaped = json.dumps( - self.current_param_value, ensure_ascii=False - )[1:-1] - delta_escaped = full_escaped[len(prev_escaped) :] - - if delta_escaped: - return DeltaMessage( - tool_calls=[ - DeltaToolCall( - index=self.current_tool_index, - function=DeltaFunctionCall( - arguments=delta_escaped - ), - ) - ] - ) - - return None diff --git a/vllm/transformers_utils/configs/speculators/algos.py b/vllm/transformers_utils/configs/speculators/algos.py index 0dc3ccce089..f1dfc8878ff 100644 --- a/vllm/transformers_utils/configs/speculators/algos.py +++ b/vllm/transformers_utils/configs/speculators/algos.py @@ -36,6 +36,8 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None: "norm_before_residual", True ) pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False) + pre_trained_config["fc_norm"] = config_dict.get("fc_norm", False) + pre_trained_config["norm_output"] = config_dict.get("norm_output", False) eagle3_arch_map = { "qwen3": "Eagle3Qwen3ForCausalLM", "llama": "Eagle3LlamaForCausalLM", diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index 657afad838b..3b5012d0175 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/vllm/utils/cpu_triton_utils.py @@ -5,6 +5,7 @@ Contains replacement functions to fallback Triton usages in CPU backend """ +import ctypes from collections.abc import Callable import torch @@ -196,6 +197,133 @@ def _copy_and_expand_eagle_inputs_kernel_impl( out_positions_ptr.copy_(out_pos_i64.to(orig_pos_dtype)) +def _copy_and_expand_dflash_inputs_kernel_impl( + next_token_ids_ptr, + target_positions_ptr, + out_input_ids_ptr, + out_context_positions_ptr, + out_query_positions_ptr, + out_context_slot_mapping_ptr, + out_query_slot_mapping_ptr, + out_token_indices_ptr, + block_table_ptr, + block_table_stride, + query_start_loc_ptr, + num_rejected_tokens_ptr, + parallel_drafting_token_id, + block_size, + num_query_per_req, + num_speculative_tokens, + total_input_tokens, + BLOCK_SIZE=None, + HAS_NUM_REJECTED=False, +): + """Adapter between the DFlash Triton launch and the C++ CPU op.""" + assert block_table_stride == block_table_ptr.stride(0), ( + "block_table_stride mismatch: " + f"{block_table_stride} vs {block_table_ptr.stride(0)}" + ) + + orig_ids_dtype = out_input_ids_ptr.dtype + orig_context_positions_dtype = out_context_positions_ptr.dtype + orig_query_positions_dtype = out_query_positions_ptr.dtype + orig_context_slot_mapping_dtype = out_context_slot_mapping_ptr.dtype + orig_query_slot_mapping_dtype = out_query_slot_mapping_ptr.dtype + out_ids_i64 = _ensure_int64(out_input_ids_ptr) + out_context_positions_i64 = _ensure_int64(out_context_positions_ptr) + out_query_positions_i64 = _ensure_int64(out_query_positions_ptr) + out_context_slot_mapping_i64 = _ensure_int64(out_context_slot_mapping_ptr) + out_query_slot_mapping_i64 = _ensure_int64(out_query_slot_mapping_ptr) + rejected_i64 = _ensure_int64(num_rejected_tokens_ptr) if HAS_NUM_REJECTED else None + + if hasattr(torch.ops._C, "copy_and_expand_dflash_inputs_kernel_impl"): + torch.ops._C.copy_and_expand_dflash_inputs_kernel_impl( + _ensure_int64(next_token_ids_ptr), + _ensure_int64(target_positions_ptr), + out_ids_i64, + out_context_positions_i64, + out_query_positions_i64, + out_context_slot_mapping_i64, + out_query_slot_mapping_i64, + out_token_indices_ptr, + block_table_ptr, + query_start_loc_ptr, + rejected_i64, + parallel_drafting_token_id, + block_size, + num_query_per_req, + num_speculative_tokens, + total_input_tokens, + HAS_NUM_REJECTED, + ) + else: + next_ids_i64 = _ensure_int64(next_token_ids_ptr) + target_positions_i64 = _ensure_int64(target_positions_ptr) + block_table_stride = block_table_ptr.stride(0) + num_reqs = query_start_loc_ptr.shape[0] - 1 + + for req_idx in range(num_reqs): + ctx_start = int(query_start_loc_ptr[req_idx].item()) + ctx_end = int(query_start_loc_ptr[req_idx + 1].item()) + num_ctx = ctx_end - ctx_start + valid_ctx_end = ctx_end + if rejected_i64 is not None: + valid_ctx_end -= int(rejected_i64[req_idx].item()) + # Guard against out-of-bounds: ensure valid_ctx_end > ctx_start. + valid_ctx_end = max(valid_ctx_end, ctx_start + 1) + + last_pos = int(target_positions_i64[valid_ctx_end - 1].item()) + + for j in range(num_ctx): + ctx_idx = ctx_start + j + ctx_pos_idx = min(ctx_idx, total_input_tokens - 1) + position = int(target_positions_i64[ctx_pos_idx].item()) + block_num = min(position // block_size, block_table_stride - 1) + block_id = int(block_table_ptr[req_idx, block_num].item()) + slot = block_id * block_size + (position % block_size) + + out_context_positions_i64[ctx_idx] = position + out_context_slot_mapping_i64[ctx_idx] = slot + + for query_off in range(num_query_per_req): + query_out = req_idx * num_query_per_req + query_off + position = last_pos + 1 + query_off + block_num = min(position // block_size, block_table_stride - 1) + block_id = int(block_table_ptr[req_idx, block_num].item()) + slot = block_id * block_size + (position % block_size) + + out_query_positions_i64[query_out] = position + out_query_slot_mapping_i64[query_out] = slot + out_ids_i64[query_out] = ( + int(next_ids_i64[req_idx].item()) + if query_off == 0 + else parallel_drafting_token_id + ) + + if query_off > 0: + sample_out_idx = req_idx * num_speculative_tokens + (query_off - 1) + out_token_indices_ptr[sample_out_idx] = query_out + + if orig_ids_dtype != torch.int64: + out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype)) + if orig_context_positions_dtype != torch.int64: + out_context_positions_ptr.copy_( + out_context_positions_i64.to(orig_context_positions_dtype) + ) + if orig_query_positions_dtype != torch.int64: + out_query_positions_ptr.copy_( + out_query_positions_i64.to(orig_query_positions_dtype) + ) + if orig_context_slot_mapping_dtype != torch.int64: + out_context_slot_mapping_ptr.copy_( + out_context_slot_mapping_i64.to(orig_context_slot_mapping_dtype) + ) + if orig_query_slot_mapping_dtype != torch.int64: + out_query_slot_mapping_ptr.copy_( + out_query_slot_mapping_i64.to(orig_query_slot_mapping_dtype) + ) + + def _rejection_greedy_sample_kernel_impl( output_token_ids, cu_num_draft_tokens, @@ -302,6 +430,10 @@ def _sample_recovered_tokens_kernel_impl( NO_DRAFT_PROBS=False, USE_FP64_GUMBEL=False, ): + # USE_FP64_GUMBEL only controls the gumbel-noise precision, which the caller + # has already applied to `inv_q` (fp64 vs fp32). The CPU kernel consumes + # `inv_q` directly, so the flag is accepted for interface parity and the + # value is read at its existing dtype. # C++ reads integer tensors as int64_t*; ensure correct dtype. orig_dtype = output_token_ids.dtype output_i64 = _ensure_int64(output_token_ids) @@ -329,6 +461,9 @@ eagle_prepare_next_token_padded_kernel = _FuncWrapper( copy_and_expand_eagle_inputs_kernel = _FuncWrapper( _copy_and_expand_eagle_inputs_kernel_impl ) +copy_and_expand_dflash_inputs_kernel = _FuncWrapper( + _copy_and_expand_dflash_inputs_kernel_impl +) eagle_step_slot_mapping_metadata_kernel = _FuncWrapper( _eagle_step_slot_mapping_metadata_kernel_impl ) @@ -336,3 +471,12 @@ rejection_greedy_sample_kernel = _FuncWrapper(_rejection_greedy_sample_kernel_im rejection_random_sample_kernel = _FuncWrapper(_rejection_random_sample_kernel_impl) expand_kernel = _FuncWrapper(_expand_kernel_impl) sample_recovered_tokens_kernel = _FuncWrapper(_sample_recovered_tokens_kernel_impl) + + +def _batch_memcpy_impl(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=None): + # BLOCK_SIZE is unused; kept for signature parity with the Triton kernel. + for src, dst, size in zip(src_ptrs.tolist(), dst_ptrs.tolist(), sizes.tolist()): + ctypes.memmove(dst, src, size) + + +batch_memcpy_kernel = _FuncWrapper(_batch_memcpy_impl) diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index b2e186ac3b7..056107c364d 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -63,6 +63,10 @@ class CPUAttentionBackend(AttentionBackend): def get_name() -> str: return "CPU_ATTN" + @classmethod + def supports_non_causal(cls) -> bool: + return True + @classmethod def supports_attn_type(cls, attn_type: str) -> bool: """CPU attention supports decoder, diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 35b57b9c2b2..d802f568836 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -81,11 +81,13 @@ class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMet ) self.kv_scale_format = _kv_scale_format_for_model(model_type) - assert indexer is not None, ( - "FLASHINFER_MLA_SPARSE_SM120 requires a sparse-MLA indexer " - "(model with index_topk in its config)." + # Skip-topk layers are built with indexer=None and get the shared + # buffer via mla_args instead (cf. FLASHMLA_SPARSE). + self.topk_indices_buffer: torch.Tensor | None = ( + indexer.topk_indices_buffer + if indexer is not None + else mla_args.get("topk_indices_buffer") ) - self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 if not has_flashinfer_sparse_mla_sm120(): diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py index 533e200cac4..bb6efe59c8c 100644 --- a/vllm/v1/attention/backends/mla/flashmla.py +++ b/vllm/v1/attention/backends/mla/flashmla.py @@ -171,7 +171,10 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]): query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] # we use the max but all should be the same due to uniform length requirement max_query_len = query_lens_cpu.max().item() - num_q_tokens_per_head_k = max_query_len * self.num_q_heads // 1 + num_q_heads = self.num_q_heads + if self.dcp_world_size > 1: + num_q_heads *= self.dcp_world_size + num_q_tokens_per_head_k = max_query_len * num_q_heads // 1 scheduler_metadata, _ = get_mla_metadata( seq_lens_device, num_q_tokens_per_head_k, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index a3822e7fc45..b13c23d8040 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -947,7 +947,9 @@ def may_override_num_blocks(vllm_config: VllmConfig, num_blocks: int) -> int: return num_blocks -def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: +def _pool_bytes_per_block( + vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec] +) -> int: """ Bytes consumed by one block in the worker's shared KV cache pool, mirroring the divisor used by `get_kv_cache_config_from_groups` to convert @@ -958,7 +960,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int: kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs ): return kv_cache_groups[0].kv_cache_spec.page_size_bytes - if _use_packed_kv_cache_groups(kv_cache_groups): + if _use_packed_kv_cache_config(vllm_config, kv_cache_groups): # buckets = {page_size: [[layer_names], [layer_names], ...]} buckets = _bucket_layers_by_page_size(kv_cache_groups) return sum(ps * len(slots) for ps, slots in buckets.items()) @@ -1250,16 +1252,26 @@ def _bucket_layers_by_page_size( return buckets -def _use_packed_kv_cache_groups( +def _use_packed_kv_cache_config( + vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], ) -> bool: is_dsv4 = all( isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs) for group in kv_cache_groups ) - return is_dsv4 or ( - bool(envs.VLLM_USE_PACKED_HMA_KV_CACHE) and len(kv_cache_groups) > 1 + kv_transfer_config = vllm_config.kv_transfer_config + extra_config = ( + kv_transfer_config.kv_connector_extra_config + if kv_transfer_config is not None + else {} ) + # NOTE: enable_cross_layers_blocks is an experimental API and subject to change with + # https://github.com/vllm-project/vllm/issues/42082 + enable_cross_layers = ( + str(extra_config.get("enable_cross_layers_blocks", "False")).lower() == "true" + ) + return is_dsv4 or (enable_cross_layers and len(kv_cache_groups) > 1) def _get_kv_cache_config_packed( @@ -1347,10 +1359,9 @@ def get_kv_cache_config_from_groups( ) for layer_name in kv_cache_groups[0].layer_names ] - elif _use_packed_kv_cache_groups(kv_cache_groups): - # DeepSeek V4 keeps the existing packed layout. Other multi-group - # attention-only HMA layouts can opt in with - # VLLM_USE_PACKED_HMA_KV_CACHE=1. + elif _use_packed_kv_cache_config(vllm_config, kv_cache_groups): + # DeepSeek V4 uses the packed layout by default. Other multi-group + # layouts can opt in with --enable-cross-layers. num_blocks, kv_cache_tensors = _get_kv_cache_config_packed( vllm_config, kv_cache_groups, available_memory ) @@ -2069,7 +2080,7 @@ def get_kv_cache_configs( if not groups: adjusted_memory.append(avail_mem) continue - bytes_per_block = _pool_bytes_per_block(groups) + bytes_per_block = _pool_bytes_per_block(vllm_config, groups) logger.info( "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", avail_mem // bytes_per_block, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 55e40b20436..ab9fd5e3433 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -415,6 +415,8 @@ class Scheduler(SchedulerInterface): encoder_compute_budget = self.max_num_encoder_input_tokens # Spec decode-related. scheduled_spec_decode_tokens: dict[str, list[int]] = {} + # Whether the running batch contains any prefill requests. + prefill_scheduled = False # For logging. scheduled_timestamp = time.monotonic() @@ -573,6 +575,7 @@ class Scheduler(SchedulerInterface): # Schedule the request. scheduled_running_reqs.append(request) + prefill_scheduled |= request.is_prefill_chunk request_id = request.request_id req_to_new_blocks[request_id] = new_blocks num_scheduled_tokens[request_id] = num_new_tokens @@ -778,6 +781,7 @@ class Scheduler(SchedulerInterface): encoder_inputs_to_schedule = None external_load_encoder_input = [] new_encoder_compute_budget = encoder_compute_budget + pad_spec_decode = False if load_kv_async: # KVTransfer: loading remote KV, do not allocate for new work. @@ -793,6 +797,23 @@ class Scheduler(SchedulerInterface): # `request.num_prompt_tokens` to consider the resumed # requests, which have output tokens. num_new_tokens = request.num_tokens - num_computed_tokens + + # Pad new decode requests to uniform spec decoding size to + # preserve full cudagraph for this step. + if ( + (self.num_spec_tokens > 0 and self.dynamic_sd_lookup is None) + and num_new_tokens == 1 + and (scheduled_running_reqs and not prefill_scheduled) + ): + num_new_tokens = 1 + self.num_spec_tokens + if ( + num_new_tokens > token_budget + or num_computed_tokens + num_new_tokens > self.max_model_len + ): + # Prefer to not schedule than schedule un-padded here. + break + pad_spec_decode = True + threshold = self.scheduler_config.long_prefill_token_threshold if 0 < threshold < num_new_tokens: num_new_tokens = threshold @@ -957,6 +978,10 @@ class Scheduler(SchedulerInterface): token_budget -= num_new_tokens request.status = RequestStatus.RUNNING request.num_computed_tokens = num_computed_tokens + if pad_spec_decode: + scheduled_spec_decode_tokens[request_id] = [ + -1 + ] * self.num_spec_tokens # Only track requests that will still be prefilling after this chunk. if num_computed_tokens + num_new_tokens < request.num_tokens: self._inflight_prefills.add(request) diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index a04f080ea6a..38ca8dc6da4 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -156,7 +156,7 @@ class EngineCoreEventType(enum.IntEnum): class EngineCoreEvent(msgspec.Struct): """A timestamped engine core event associated with a request. - The timestamp is a monotonic timestamps and is used for by the engine + The timestamp is a monotonic timestamp and is used by the engine frontend to calculate intervals between engine core events. These timestamps should not be compared with timestamps from other processes. """ diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index cfeec4456ea..57a788631ce 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -1950,13 +1950,8 @@ class DPEngineCoreProc(EngineCoreProc): continue # Execute a dummy pass when no ready requests ran, unless the - # engine is sleeping. self.is_sleeping() also covers the KV-offload - # window before model_executor.is_sleeping flips. - elif not self.is_sleeping(): - with self.log_iteration_details(None): - # We are in a running state and so must execute a dummy pass - # if the model didn't execute any ready requests. - if not self.model_executor.is_sleeping: + # engine is sleeping. + elif not self.model_executor.is_sleeping: with self.log_iteration_details(None): self.execute_dummy_batch() diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index db0940ef386..507e457ac50 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -7,7 +7,7 @@ Core abstractions for KV cache offloading in vLLM v1. from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass -from enum import Enum +from enum import Enum, auto from typing import TYPE_CHECKING, Any, NewType import numpy as np @@ -53,6 +53,15 @@ class ReqContext: kv_transfer_params: dict[str, Any] | None = None +class LookupResult(Enum): + """Result of OffloadingManager.lookup().""" + + MISS = auto() + HIT = auto() + HIT_PENDING = auto() + RETRY = auto() + + class OffloadPolicy(Enum): # Offload only newly-computed blocks as they arrive; prefix-hit # blocks (already offloaded by a prior request) are skipped. @@ -158,7 +167,7 @@ class OffloadingKVEventsConfig: class OffloadingManager(ABC): @abstractmethod - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ Checks whether a single block is offloaded and ready to be read. @@ -167,10 +176,9 @@ class OffloadingManager(ABC): req_context: per-request context (e.g. kv_transfer_params). Returns: - True if the block is offloaded and ready, False if not, - or None if the lookup should be retried later. - Returning None will delay the request handling by the vLLM - scheduler. + HIT if the block is offloaded and ready, MISS if not found, + HIT_PENDING if found but not yet readable, or RETRY if the + lookup should be retried later. """ pass diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index c8b9915a1e5..843e1538f90 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/vllm/v1/kv_offload/cpu/gpu_worker.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +import threading import time from collections import deque from dataclasses import dataclass @@ -120,36 +121,6 @@ def compute_sub_block_ptrs( output[:] = flat[skip_count : skip_count + num_sub_blocks] -def pin_mmap_region(region: SharedOffloadRegion) -> None: - """Register the entire mmap as CUDA pinned memory via cudaHostRegister.""" - if not current_platform.is_cuda_alike(): - logger.info( - "Skipping mmap host registration on %s; cudaHostRegister is only " - "available on CUDA/ROCm.", - current_platform.device_name, - ) - return - - rank = region.rank - - base_ptr = region._base.data_ptr() - result = torch.cuda.cudart().cudaHostRegister(base_ptr, region.total_size_bytes, 0) - if result.value != 0: - logger.warning( - "cudaHostRegister failed for rank=%d (code=%d) — " - "transfers will still work but may be slower (unpinned DMA)", - rank, - result, - ) - else: - logger.debug( - "cudaHostRegister rank=%d %.2f GB", - rank, - region.total_size_bytes / 1e9, - ) - region.is_pinned = True - - def _new_descriptor_buffers( num_copy_ops: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -179,6 +150,8 @@ class SingleDirectionOffloadingHandler: kv_cache_groups_data_refs: list[list[CanonicalKVCacheRef]], gpu_to_cpu: bool, mmap_region: SharedOffloadRegion | None = None, + pin_thread: threading.Thread | None = None, + manually_pinned_tensors: list[torch.Tensor] | None = None, ): """ Initialize a SingleDirectionOffloadingHandler. @@ -226,6 +199,8 @@ class SingleDirectionOffloadingHandler: # mmap_region to clean up on shutdown (gpu_to_cpu handler owns it) self._mmap_region = mmap_region + self._pin_thread = pin_thread + self._manually_pinned_tensors = manually_pinned_tensors # job_id -> event self._transfer_events: dict[int, torch.Event] = {} # queue of transfers (job_id, stream, event) @@ -458,8 +433,23 @@ class SingleDirectionOffloadingHandler: self._stream_pool.clear() self._event_pool.clear() self._buffer_pool.clear() + + if self._pin_thread is not None: + self._pin_thread.join() + self._pin_thread = None + + if self._manually_pinned_tensors is not None: + for tensor in self._manually_pinned_tensors: + result = torch.cuda.cudart().cudaHostUnregister(tensor.data_ptr()) + if result.value != 0: + logger.warning( + "cudaHostUnregister failed for CPU tensor (code=%d)", + result.value, + ) + self.src_tensors.clear() self.dst_tensors.clear() + if self._mmap_region is not None: self._mmap_region.cleanup() self._mmap_region = None @@ -481,12 +471,14 @@ class CPUOffloadingWorker(OffloadingWorker): mmap_region: SharedOffloadRegion | None = None, ): pin_memory = PIN_MEMORY + self.pin_thread: threading.Thread | None = None + self._manually_pinned_tensors: list[torch.Tensor] = [] + logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors)) - if mmap_region is not None and pin_memory: - pin_mmap_region(mmap_region) + self._mmap_region = mmap_region gpu_tensors: list[torch.Tensor] = [] - cpu_tensors: list[torch.Tensor] = [] + self.cpu_tensors: list[torch.Tensor] = [] for kv_cache_tensor in kv_caches.tensors: gpu_page_size_bytes = kv_cache_tensor.page_size_bytes gpu_tensor = kv_cache_tensor.tensor.view(torch.int8).view( @@ -502,10 +494,13 @@ class CPUOffloadingWorker(OffloadingWorker): (num_cpu_blocks, cpu_page_size_bytes), dtype=torch.int8, device="cpu", - pin_memory=pin_memory, + # CUDA/ROCm memory is registered asynchronously below. + # Pinning here would block worker initialization; other + # hardware need PyTorch allocation-time pinning. + pin_memory=PIN_MEMORY and not current_platform.is_cuda_alike(), ) logger.debug( - "torch.zeros pinned tensor %d×%d (%.2f GB): %.3f s", + "torch.zeros tensor %d×%d (%.2f GB): %.3f s", num_cpu_blocks, cpu_page_size_bytes, num_cpu_blocks * cpu_page_size_bytes / 1e9, @@ -513,25 +508,81 @@ class CPUOffloadingWorker(OffloadingWorker): ) gpu_tensors.append(gpu_tensor) - cpu_tensors.append(cpu_tensor) + self.cpu_tensors.append(cpu_tensor) + + if pin_memory: + if not current_platform.is_cuda_alike(): + logger.info( + "Skipping host registration on %s; cudaHostRegister is only " + "available on CUDA/ROCm.", + current_platform.device_name, + ) + else: + self.pin_thread = threading.Thread( + target=self._pin_cpu_tensors, + name="CPUTensorPinThread", + ) + self.pin_thread.start() + logger.info("Starting to pin memory in background...") self._store_handler = SingleDirectionOffloadingHandler( gpu_tensors=gpu_tensors, - cpu_tensors=cpu_tensors, + cpu_tensors=self.cpu_tensors, block_size_factor=block_size_factor, kv_cache_groups_data_refs=kv_caches.group_data_refs, gpu_to_cpu=True, mmap_region=mmap_region, + pin_thread=self.pin_thread, + manually_pinned_tensors=self._manually_pinned_tensors, ) self._load_handler = SingleDirectionOffloadingHandler( gpu_tensors=gpu_tensors, - cpu_tensors=cpu_tensors, + cpu_tensors=self.cpu_tensors, block_size_factor=block_size_factor, kv_cache_groups_data_refs=kv_caches.group_data_refs, gpu_to_cpu=False, ) + def _pin_cpu_tensors(self) -> None: + """Register the CPU offload memory as CUDA pinned memory.""" + + t0 = time.monotonic() + tensors_to_pin = ( + [self._mmap_region._base] + if self._mmap_region is not None + else self.cpu_tensors + ) + num_pinned = 0 + for tensor in tensors_to_pin: + total_size_bytes = tensor.numel() * tensor.element_size() + result = torch.cuda.cudart().cudaHostRegister( + tensor.data_ptr(), total_size_bytes, 0 + ) + if result.value != 0: + logger.warning( + "cudaHostRegister failed for host tensor (code=%d) " + "- transfers will still work but may be slower (unpinned DMA)", + result.value, + ) + continue + if self._mmap_region is not None: + self._mmap_region.is_pinned = True + else: + self._manually_pinned_tensors.append(tensor) + num_pinned += 1 + + logger.debug( + "cudaHostRegister pin %.2f GB", + total_size_bytes / 1e9, + ) + + logger.info( + "Completed CPU memory pinning: %d tensors pinned in %.3f s", + num_pinned, + time.monotonic() - t0, + ) + def submit_store( self, job_id: int, src_spec: GPULoadStoreSpec, dst_spec: LoadStoreSpec ) -> bool: diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index b48abecec1b..0424196c9fd 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( ) from vllm.v1.kv_offload.base import ( LoadStoreSpec, + LookupResult, OffloadingEvent, OffloadingManager, OffloadKey, @@ -112,7 +113,7 @@ class CPUOffloadingManager(OffloadingManager): return RequestOffloadingContext() @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: if self.counts is not None: if key in self.counts: self.counts.move_to_end(key) @@ -123,10 +124,10 @@ class CPUOffloadingManager(OffloadingManager): self.counts[key] = 1 block = self._policy.get(key) if block is None: - return False + return LookupResult.MISS if not block.is_ready: - return None # write in-flight; caller should retry - return True + return LookupResult.HIT_PENDING + return LookupResult.HIT @override def prepare_load( diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index c7927572491..662a826a06d 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any import numpy as np from vllm.v1.kv_offload.base import ( + LookupResult, OffloadingMetricMetadata, OffloadKey, ReqContext, @@ -79,7 +80,7 @@ class SecondaryTierManager(ABC): self.tier_type = tier_type @abstractmethod - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ Check whether a block exists in this secondary tier. @@ -88,9 +89,9 @@ class SecondaryTierManager(ABC): req_context: per-request context (e.g. kv_transfer_params). Returns: - True if the block is present and ready, - False if not found, - or None if the block is being transferred (retry later). + HIT if the block is present and ready, + MISS if not found, + or RETRY if the block is being transferred (retry later). """ pass diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index d352ff54c6e..a9e4e4f689c 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -15,7 +15,12 @@ from typing import TYPE_CHECKING from typing_extensions import override -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + RequestOffloadingContext, +) from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, @@ -67,7 +72,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): self.completed_jobs: list[JobResult] = [] @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ Check whether a block exists in this secondary tier. @@ -76,9 +81,9 @@ class ExampleSecondaryTierManager(SecondaryTierManager): req_context: Per-request context. Returns: - True if the block is present, False if not found. + HIT if the block is present, MISS if not found. """ - return key in self.blocks + return LookupResult.HIT if key in self.blocks else LookupResult.MISS @override def submit_store(self, job_metadata: JobMetadata) -> None: diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index e411f670650..329a24daf34 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -24,7 +24,7 @@ from typing import TYPE_CHECKING from typing_extensions import override from vllm.logger import init_logger -from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( @@ -137,8 +137,11 @@ class FileSystemTierManager(SecondaryTierManager): return RequestOffloadingContext() @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - return self._lookup_manager.lookup(key, req_context) + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + result = self._lookup_manager.lookup(key, req_context) + if result is None: + return LookupResult.RETRY + return LookupResult.HIT if result else LookupResult.MISS @override def submit_store(self, job_metadata: JobMetadata) -> None: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 85346eacb7b..ee5b0b52742 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -32,6 +32,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( LoadStoreSpec, + LookupResult, OffloadingEvent, OffloadingManager, OffloadKey, @@ -233,7 +234,7 @@ class TieringOffloadingManager(OffloadingManager): ) @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ Check whether a single block is offloaded and ready. @@ -248,33 +249,34 @@ class TieringOffloadingManager(OffloadingManager): req_context: Per-request context. Returns: - True — block is ready in the primary tier. - None — block found but not yet ready (primary in-flight, - promotion started, or a secondary tier is busy). - False — block not found in any tier, or primary is full - and cannot accept a promotion. + HIT — block is ready in the primary tier. + HIT_PENDING — block found but not yet readable (write + in-flight on the primary tier). + RETRY — promotion started or a secondary tier is busy. + MISS — block not found in any tier, or primary is full + and cannot accept a promotion. """ self._maybe_process_finished_jobs() primary_hit = self.primary_tier.lookup(key, req_context) - if primary_hit is True: - return True - if primary_hit is None: - return None + if primary_hit is LookupResult.HIT: + return LookupResult.HIT + if primary_hit is LookupResult.HIT_PENDING: + return LookupResult.HIT_PENDING - any_none = False + any_retry = False for tier in self.secondary_tiers: result = tier.lookup(key, req_context) - if result is True: + if result is LookupResult.HIT: if not self._initiate_promotion(tier, key, req_context): - return False # primary full, block unavailable - return None # promotion started, retry later - if result is None: - any_none = True + return LookupResult.MISS + return LookupResult.RETRY + if result is LookupResult.RETRY: + any_retry = True - if any_none: - return None - return False + if any_retry: + return LookupResult.RETRY + return LookupResult.MISS def _initiate_promotion( self, @@ -467,7 +469,9 @@ class TieringOffloadingManager(OffloadingManager): """ # Filter out keys that are not ready in primary (e.g. in-flight) ready_keys = tuple( - k for k in keys if self.primary_tier.lookup(k, req_context) is True + k + for k in keys + if self.primary_tier.lookup(k, req_context) is LookupResult.HIT ) if not ready_keys: return diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index ec032dc1a27..857c3c758a2 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, NamedTuple from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent from vllm.distributed.nixl_utils import nixl_agent_config from vllm.logger import init_logger -from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( @@ -221,8 +221,11 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle) - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - return self._lookup_manager.lookup(key, req_context) + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + result = self._lookup_manager.lookup(key, req_context) + if result is None: + return LookupResult.RETRY + return LookupResult.HIT if result else LookupResult.MISS def submit_store(self, job_metadata: JobMetadata) -> None: obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index d20cac37fcd..c284ff61876 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -111,7 +111,7 @@ def _topk_topp_kernel( pid = tl.program_id(0) num_programs = tl.num_programs(0) for row_id in tl.range(pid, BATCH_SIZE, num_programs): - LOGITS_ROW = LOGITS + row_id * LOGITS_STRIDE_0 + LOGITS_ROW = LOGITS + row_id.to(tl.int64) * LOGITS_STRIDE_0 BUFFER_ROW = BUFFER + pid * VOCAB_SIZE final_pivot = -float("inf") diff --git a/vllm/v1/sample/thinking_budget_state.py b/vllm/v1/sample/thinking_budget_state.py index d32d1b30296..6e4ef0d1278 100644 --- a/vllm/v1/sample/thinking_budget_state.py +++ b/vllm/v1/sample/thinking_budget_state.py @@ -173,6 +173,19 @@ class ThinkingBudgetStateHolder: return i return -1 + @staticmethod + def _find_last_sequence_index_from( + target_list: list[int], token_ids: list[int], search_start: int + ) -> int: + """Last occurrence of ``token_ids`` at or after ``search_start``.""" + if not token_ids: + return -1 + lo = max(0, search_start) + for i in range(len(target_list) - len(token_ids), lo - 1, -1): + if target_list[i : i + len(token_ids)] == token_ids: + return i + return -1 + def _init_state_entry( self, prompt_tok_ids: list[int] | None, thinking_token_budget: int ) -> dict[str, Any]: @@ -226,6 +239,8 @@ class ThinkingBudgetStateHolder: "force_index": [], "start_thinking": start_thinking, "end_thinking": -1, + "start_search_pos": 0, + "end_search_pos": 0, "in_spec_mode": False, "bonus_token_forced": False, "continue_thinking": continue_thinking, @@ -240,16 +255,27 @@ class ThinkingBudgetStateHolder: state["force_index"] = [] return + output_tok_ids = state.get("output_tok_ids", []) if state["start_thinking"] == -1: - start_thinking = self._find_last_sequence_index( - state.get("output_tok_ids", []), self.think_start_token_ids + seq_len = len(self.think_start_token_ids) + start_thinking = self._find_last_sequence_index_from( + output_tok_ids, + self.think_start_token_ids, + state["start_search_pos"] - (seq_len - 1), ) state["start_thinking"] = start_thinking + if start_thinking == -1: + state["start_search_pos"] = len(output_tok_ids) if state["end_thinking"] == -1: - end_thinking = self._find_last_sequence_index( - state.get("output_tok_ids", []), self.think_end_token_ids + seq_len = len(self.think_end_token_ids) + end_thinking = self._find_last_sequence_index_from( + output_tok_ids, + self.think_end_token_ids, + state["end_search_pos"] - (seq_len - 1), ) state["end_thinking"] = end_thinking + if end_thinking == -1: + state["end_search_pos"] = len(output_tok_ids) if state["start_thinking"] == -1: return diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 07978a9dd61..dfaa2234eb9 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -103,9 +103,12 @@ class SimpleCPUOffloadScheduler: assert 0 <= self.fa_gidx < len(self.cpu_kv_cache_config.kv_cache_groups) # FA group's own block_size; divides scheduler_block_size (the LCM) # but is NOT assumed to equal it. - self.fa_block_size: int = self.cpu_kv_cache_config.kv_cache_groups[ - self.fa_gidx - ].kv_cache_spec.block_size + self.fa_block_size: int = ( + self.cpu_kv_cache_config.kv_cache_groups[ + self.fa_gidx + ].kv_cache_spec.block_size + * self.cp_world_size + ) assert self.block_size % self.fa_block_size == 0 logger.info( @@ -348,10 +351,12 @@ class SimpleCPUOffloadScheduler: # the rest will be released along with the temp pin below. cpu_hit_blocks: list[list[KVCacheBlock]] = [] for g in range(num_groups): - g_block_size = kv_cache_groups[g].kv_cache_spec.block_size + g_block_size = ( + kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size + ) assert num_external_tokens % g_block_size == 0, ( f"num_external_tokens={num_external_tokens} not aligned to " - f"group {g} block_size={g_block_size}" + f"group {g} effective block_size={g_block_size}" ) n_take_g = num_external_tokens // g_block_size cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) @@ -599,7 +604,9 @@ class SimpleCPUOffloadScheduler: already_stored_g = state.num_stored_blocks[g] group_gpu_ids = block_ids_by_group[g] - g_block_size = kv_cache_groups[g].kv_cache_spec.block_size + g_block_size = ( + kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size + ) ready_blocks_g = aligned_tokens // g_block_size scannable = group_gpu_ids[already_stored_g:ready_blocks_g] diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index f76305d0857..bae6935cef8 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -10,10 +10,12 @@ from typing_extensions import override from vllm.config import VllmConfig from vllm.forward_context import set_forward_context from vllm.logger import init_logger -from vllm.triton_utils import triton from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer -from vllm.v1.spec_decode.utils import copy_and_expand_dflash_inputs_kernel +from vllm.v1.spec_decode.utils import ( + copy_and_expand_dflash_inputs_kernel, + next_power_of_2, +) logger = init_logger(__name__) @@ -126,8 +128,8 @@ class DFlashProposer(SpecDecodeBaseProposer): # and token_indices_to_sample max_ctx_per_req = cad.max_query_len max_tokens_per_req = max_ctx_per_req + num_query_per_req - BLOCK_SIZE = min(256, triton.next_power_of_2(max_tokens_per_req)) - num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE) + BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req)) + num_blocks = (max_tokens_per_req + BLOCK_SIZE - 1) // BLOCK_SIZE grid = (batch_size, num_blocks) has_num_rejected = num_rejected_tokens_gpu is not None diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index 71ade9c8607..a083f309e0e 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -365,19 +365,23 @@ class RustFrontendProcessManager: cmd.extend(["--coordinator-address", stats_update_address]) from vllm.entrypoints.serve.utils.api_utils import jsonify_non_default_args - args_json = json.dumps( - jsonify_non_default_args( - args, - exclude={ - "api_server_count", - # Python passes the bootstrapped engine range explicitly. - "data_parallel_rank", - "data_parallel_external_lb", - "data_parallel_hybrid_lb", - }, - ), - sort_keys=True, + args_dict = jsonify_non_default_args( + args, + exclude={ + "api_server_count", + # Python passes the bootstrapped engine range explicitly. + "data_parallel_rank", + "data_parallel_external_lb", + "data_parallel_hybrid_lb", + }, ) + # The Rust `frontend` subcommand parses --args-json via serde_json, + # which bypasses clap and therefore ignores any `#[arg(env = ...)]` + # declarations on SharedRuntimeArgs fields. Forward the env-driven + # ready timeout explicitly so VLLM_ENGINE_READY_TIMEOUT_S behaves the + # same on both Python and Rust frontends. + args_dict["engine_ready_timeout_secs"] = envs.VLLM_ENGINE_READY_TIMEOUT_S + args_json = json.dumps(args_dict, sort_keys=True) cmd.extend(["--args-json", args_json]) logger.info("Launching Rust frontend: %s", " ".join(cmd)) diff --git a/vllm/v1/worker/cpu_model_runner.py b/vllm/v1/worker/cpu_model_runner.py index 87b7a9ad220..87f8cb154dc 100644 --- a/vllm/v1/worker/cpu_model_runner.py +++ b/vllm/v1/worker/cpu_model_runner.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import sys from contextlib import contextmanager from typing import Any @@ -12,7 +13,7 @@ from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model from vllm.tracing import instrument from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker.gpu_model_runner import GPUModelRunner @@ -78,7 +79,7 @@ class CPUModelRunner(GPUModelRunner): # Speculative decoding fallbacks import vllm.v1.sample.rejection_sampler import vllm.v1.spec_decode.llm_base_proposer - import vllm.v1.spec_decode.utils + import vllm.v1.spec_decode.utils as spec_decode_utils vllm.v1.spec_decode.llm_base_proposer.eagle_prepare_inputs_padded_kernel = ( cpu_tl.eagle_prepare_inputs_padded_kernel @@ -89,7 +90,18 @@ class CPUModelRunner(GPUModelRunner): vllm.v1.spec_decode.llm_base_proposer.copy_and_expand_eagle_inputs_kernel = ( cpu_tl.copy_and_expand_eagle_inputs_kernel ) - vllm.v1.spec_decode.utils.eagle_step_slot_mapping_metadata_kernel = ( + spec_decode_utils.copy_and_expand_dflash_inputs_kernel = ( + cpu_tl.copy_and_expand_dflash_inputs_kernel + ) + dflash_module = sys.modules.get("vllm.v1.spec_decode.dflash") + if dflash_module is not None: + dflash_kernel_name = "copy_and_expand_dflash_inputs_kernel" + setattr( + dflash_module, + dflash_kernel_name, + cpu_tl.copy_and_expand_dflash_inputs_kernel, + ) + spec_decode_utils.eagle_step_slot_mapping_metadata_kernel = ( cpu_tl.eagle_step_slot_mapping_metadata_kernel ) vllm.v1.sample.rejection_sampler.rejection_greedy_sample_kernel = ( @@ -103,6 +115,10 @@ class CPUModelRunner(GPUModelRunner): cpu_tl.sample_recovered_tokens_kernel ) + import vllm.v1.worker.mamba_utils + + vllm.v1.worker.mamba_utils.batch_memcpy_kernel = cpu_tl.batch_memcpy_kernel + @instrument(span_name="Loading (CPU)") def load_model(self, load_dummy_weights: bool = False) -> None: if load_dummy_weights: @@ -153,9 +169,25 @@ class CPUModelRunner(GPUModelRunner): pass def _zero_block_ids(self, block_ids: list[int]) -> None: - # CPU attention assigns -INF to logits at invalid positions, - # so stale KV cache data never affects computation. - pass + # Zero full-attention blocks to prevent stale data corruption on partial writes. + # Encoder-only (runner-only) layers are not FullAttentionSpec, so the + # spec filter below already excludes them; no runner-only skip needed. + seen_ptrs: set[int] = set() + for group in self.kv_cache_config.kv_cache_groups: + if not isinstance(group.kv_cache_spec, FullAttentionSpec): + continue + for layer_name in group.layer_names: + ctx = self.compilation_config.static_forward_context.get(layer_name) + if ctx is None: + continue + kv = ctx.kv_cache + if not isinstance(kv, torch.Tensor): + continue + if kv.data_ptr() in seen_ptrs: + continue + seen_ptrs.add(kv.data_ptr()) + for block_id in block_ids: + kv[block_id].zero_() # ========================================================================= # CPU-safe overrides for speculative decoding methods diff --git a/vllm/v1/worker/gpu/sample/bad_words.py b/vllm/v1/worker/gpu/sample/bad_words.py index 6286cc38359..b5517dee1b1 100644 --- a/vllm/v1/worker/gpu/sample/bad_words.py +++ b/vllm/v1/worker/gpu/sample/bad_words.py @@ -114,7 +114,7 @@ def _bad_words_kernel( input_ids_ptr, expanded_local_pos_ptr, ): - token_idx = tl.program_id(0) + token_idx = tl.program_id(0).to(tl.int64) bw_idx = tl.program_id(1) req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index fab53fef7ee..6dbb04cd933 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -23,7 +23,7 @@ def _temperature_kernel( vocab_size, BLOCK_SIZE: tl.constexpr, ): - token_idx = tl.program_id(0) + token_idx = tl.program_id(0).to(tl.int64) req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32) if temperature == 0.0 or temperature == 1.0: @@ -91,7 +91,7 @@ def gumbel_block_argmax( USE_FP64: tl.constexpr, PER_TOKEN_COL: tl.constexpr = False, ): - req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) + req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx).to(tl.int64) temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) if temp != 0.0 and APPLY_TEMPERATURE: # Apply temperature. @@ -169,7 +169,7 @@ def _gumbel_sample_kernel( USE_FP64: tl.constexpr, PER_TOKEN_COL: tl.constexpr, ): - token_idx = tl.program_id(0) + token_idx = tl.program_id(0).to(tl.int64) block_idx = tl.program_id(1) block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = block < vocab_size diff --git a/vllm/v1/worker/gpu/sample/logit_bias.py b/vllm/v1/worker/gpu/sample/logit_bias.py index f3f7c29b3f3..6c95ed7aacb 100644 --- a/vllm/v1/worker/gpu/sample/logit_bias.py +++ b/vllm/v1/worker/gpu/sample/logit_bias.py @@ -169,7 +169,7 @@ def _bias_kernel( BLOCK_SIZE: tl.constexpr, LOGITS_BLOCK_SIZE: tl.constexpr, ): - token_idx = tl.program_id(0) + token_idx = tl.program_id(0).to(tl.int64) req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) block = tl.arange(0, BLOCK_SIZE) diff --git a/vllm/v1/worker/gpu/sample/logprob.py b/vllm/v1/worker/gpu/sample/logprob.py index cf24c186e93..0028e8c3a9d 100644 --- a/vllm/v1/worker/gpu/sample/logprob.py +++ b/vllm/v1/worker/gpu/sample/logprob.py @@ -21,7 +21,7 @@ def _topk_log_softmax_kernel( BLOCK_SIZE: tl.constexpr, PADDED_TOPK: tl.constexpr, ): - req_idx = tl.program_id(0) + req_idx = tl.program_id(0).to(tl.int64) row_ptr = logits_ptr + req_idx * logits_stride max_val = float("-inf") @@ -61,7 +61,7 @@ def _ranks_kernel( vocab_size, BLOCK_SIZE: tl.constexpr, ): - req_idx = tl.program_id(0) + req_idx = tl.program_id(0).to(tl.int64) row_ptr = logits_ptr + req_idx * logits_stride token_id = tl.load(token_ids_ptr + req_idx) diff --git a/vllm/v1/worker/gpu/sample/min_p.py b/vllm/v1/worker/gpu/sample/min_p.py index 4f08af2f5a5..b71ae6f3add 100644 --- a/vllm/v1/worker/gpu/sample/min_p.py +++ b/vllm/v1/worker/gpu/sample/min_p.py @@ -14,7 +14,7 @@ def _min_p_kernel( vocab_size, BLOCK_SIZE: tl.constexpr, ): - token_idx = tl.program_id(0) + token_idx = tl.program_id(0).to(tl.int64) req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) min_p = tl.load(min_p_ptr + req_state_idx).to(tl.float32) if min_p == 0.0: diff --git a/vllm/v1/worker/gpu/sample/penalties.py b/vllm/v1/worker/gpu/sample/penalties.py index b2ce2fb812a..25cb2f211d9 100644 --- a/vllm/v1/worker/gpu/sample/penalties.py +++ b/vllm/v1/worker/gpu/sample/penalties.py @@ -120,7 +120,7 @@ def _penalties_kernel( vocab_size, BLOCK_SIZE: tl.constexpr, ): - token_idx = tl.program_id(0) + token_idx = tl.program_id(0).to(tl.int64) req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx) freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 92294e6c7e4..bad70aa0451 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -79,14 +79,14 @@ def _compute_block_stats_kernel( BLOCK_SIZE: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, ): - logit_idx = tl.program_id(0) + logit_idx = tl.program_id(0).to(tl.int64) draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) if draft_step_idx >= num_speculative_steps: # Bonus token. Max/argmax and summed exponentials are not needed. return - req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx) + req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64) temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) block_idx = tl.program_id(1) @@ -206,8 +206,8 @@ def _rejection_kernel( SYNTHETIC_MODE: tl.constexpr, ): req_idx = tl.program_id(0) - req_state_idx = tl.load(idx_mapping_ptr + req_idx) - start_idx = tl.load(cu_num_logits_ptr + req_idx) + req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64) + start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64) end_idx = tl.load(cu_num_logits_ptr + req_idx + 1) num_tokens = end_idx - start_idx seed = tl.load(seed_ptr + req_state_idx) @@ -349,10 +349,10 @@ def _resample_kernel( ): req_idx = tl.program_id(0) resample_idx = tl.load(rejected_step_ptr + req_idx) - start_idx = tl.load(cu_num_logits_ptr + req_idx) + start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64) end_idx = tl.load(cu_num_logits_ptr + req_idx + 1) resample_token_idx = start_idx + resample_idx - req_state_idx = tl.load(expanded_idx_mapping_ptr + resample_token_idx) + req_state_idx = tl.load(expanded_idx_mapping_ptr + resample_token_idx).to(tl.int64) temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) is_bonus = resample_token_idx == end_idx - 1