test: touch tilelang.py for coverage comparison

Signed-off-by: Raj Joshi <rajjoshi@redhat.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

.buildkite/hardware_tests/amd.yaml

@@ -17,6 +17,26 @@ steps:
       --target test
       --no-cache
       --progress plain .
+    - |
+      docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
+        if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
+        if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
+        if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
+        if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
+        command -v python3
+        command -v uv
+        command -v pytest
+        if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
+          echo No ROCm CLI found in image >&2
+          exit 1
+        fi
+        python3 - <<PY
+      import torch, vllm
+      print(torch.__version__)
+      print(vllm.__version__)
+      PY
+        echo AMD image smoke OK
+      '
     - docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
     env:
       DOCKER_BUILDKIT: "1"

.buildkite/hardware_tests/cpu.yaml

@@ -64,7 +64,7 @@ steps:
   - vllm/v1/worker/gpu/
   commands:
     - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      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"
 
@@ -74,11 +74,10 @@ steps:
   no_plugin: true
   source_file_dependencies:
   - csrc/cpu/
-  - vllm/model_executor/layers/quantization/cpu_wna16.py
   - vllm/model_executor/layers/quantization/auto_gptq.py
   - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
-  - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
-  - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
+  - vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+  - vllm/model_executor/kernels/linear/scaled_mm/cpu.py
   - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
   - tests/quantization/test_compressed_tensors.py
   - tests/quantization/test_cpu_wna16.py

.buildkite/image_build/image_build.yaml

@@ -98,3 +98,21 @@ steps:
           limit: 2
         - exit_status: -10  # Agent was lost
           limit: 2
+
+  - label: ":docker: Build arm64 image"
+    key: arm64-image-build
+    depends_on: []
+    source_file_dependencies:
+      - ".buildkite/image_build/image_build.yaml"
+      - ".buildkite/image_build/image_build_arm64.sh"
+      - "docker/Dockerfile"
+    commands:
+    - .buildkite/image_build/image_build_arm64.sh $REGISTRY $REPO $BUILDKITE_COMMIT
+    env:
+      DOCKER_BUILDKIT: "1"
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2

.buildkite/image_build/image_build_arm64.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+set -e
+
+if [[ $# -lt 3 ]]; then
+  echo "Usage: $0 <registry> <repo> <commit>"
+  exit 1
+fi
+
+REGISTRY=$1
+REPO=$2
+BUILDKITE_COMMIT=$3
+
+# authenticate with AWS ECR
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
+
+# skip build if image already exists
+if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64) ]]; then
+  echo "Image not found, proceeding with build..."
+else
+  echo "Image found"
+  exit 0
+fi
+
+# build (Grace/GH200 is the arm64 GPU target; sm_90)
+docker build --file docker/Dockerfile \
+  --platform linux/arm64 \
+  --build-arg max_jobs=16 \
+  --build-arg nvcc_threads=4 \
+  --build-arg torch_cuda_arch_list="9.0" \
+  --build-arg USE_SCCACHE=1 \
+  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
+  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \
+  --target test \
+  --progress plain .
+
+# push
+docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64

.buildkite/release-pipeline.yaml

@@ -737,7 +737,7 @@ steps:
       - "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
     env:
       S3_BUCKET: "vllm-wheels"
-      VARIANT: "rocm722"
+      VARIANT: "rocm723"
 
   # ROCm Job 6: Build ROCm Release Docker Image
   - label: ":docker: Build release image - x86_64 - ROCm"

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -35,25 +35,9 @@ export PYTHONPATH=".."
 # Helper Functions
 ###############################################################################
 
-cleanup_docker() {
-  # Get Docker's root directory
-  docker_root=$(docker info -f '{{.DockerRootDir}}')
-  if [ -z "$docker_root" ]; then
-    echo "Failed to determine Docker root directory."
-    exit 1
-  fi
-  echo "Docker root directory: $docker_root"
-
-  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
-  threshold=70
-  if [ "$disk_usage" -gt "$threshold" ]; then
-    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
-    docker image prune -f
-    docker volume prune -f && docker system prune --force --filter "until=72h" --all
-    echo "Docker images and volumes cleanup completed."
-  else
-    echo "Disk usage is below $threshold%. No cleanup needed."
-  fi
+report_docker_usage() {
+  echo "--- Docker usage"
+  docker system df || true
 }
 
 cleanup_network() {
@@ -254,8 +238,8 @@ re_quote_pytest_markers() {
 echo "--- ROCm info"
 rocminfo
 
-# --- Docker housekeeping ---
-cleanup_docker
+# --- Docker status ---
+report_docker_usage
 
 # --- Pull test image ---
 echo "--- Pulling container"
@@ -264,9 +248,17 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea
 docker pull "${image_name}"
 
 remove_docker_container() {
-  docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
+  # docker run uses --rm, so the container is normally already gone when the
+  # EXIT trap runs. Cleanup is best-effort and must not affect the test result.
+  docker rm -f "${container_name}" >/dev/null 2>&1 || true
 }
-trap remove_docker_container EXIT
+
+on_exit() {
+  local exit_code=$?
+  remove_docker_container
+  exit "$exit_code"
+}
+trap on_exit EXIT
 
 # --- Prepare commands ---
 echo "--- Running container"

.buildkite/test_areas/engine.yaml

@@ -38,7 +38,7 @@ steps:
     - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 40
       depends_on:
       - image-build-amd

.buildkite/test_areas/entrypoints.yaml

@@ -28,7 +28,8 @@ steps:
   - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       depends_on:
       - image-build-amd
 
@@ -45,7 +46,8 @@ steps:
   - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       timeout_in_minutes: 80
       depends_on:
       - image-build-amd
@@ -63,7 +65,8 @@ steps:
   - pytest -v -s entrypoints/test_chat_utils.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -82,7 +85,8 @@ steps:
   - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -104,7 +108,8 @@ steps:
   - pytest -v -s tool_use
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
+      soft_fail: true
       depends_on:
       - image-build-amd
 

.buildkite/test_areas/kernels.yaml

@@ -86,7 +86,7 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       source_file_dependencies:
       - csrc/quantization/
       - vllm/model_executor/layers/quantization

.buildkite/test_areas/misc.yaml

@@ -52,7 +52,7 @@ steps:
     - pytest -v -s v1/test_outputs.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 

.buildkite/test_areas/models_basic.yaml

@@ -58,37 +58,3 @@ steps:
   device: cpu-small
   commands:
     - pytest -v -s models/test_utils.py models/test_vision.py
-
-- label: Transformers Nightly Models
-  device: h200_35gb
-  key: transformers-nightly-models
-  working_dir: "/vllm-workspace/"
-  optional: true
-  soft_fail: true
-  commands:
-    - pip install --upgrade git+https://github.com/huggingface/transformers
-    - pytest -v -s tests/models/test_initialization.py
-    - pytest -v -s tests/models/test_transformers.py
-    - pytest -v -s tests/models/multimodal/processing/
-    - pytest -v -s tests/models/multimodal/test_mapping.py
-    - python3 examples/basic/offline_inference/chat.py
-    - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
-    # Whisper needs spawn method to avoid deadlock
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
-
-- label: Transformers Backward Compatibility Models Test
-  device: h200_35gb
-  key: transformers-backward-compatibility-models-test
-  working_dir: "/vllm-workspace/"
-  optional: true
-  soft_fail: true
-  commands:
-    - pip install transformers==4.57.5
-    - pytest -v -s tests/models/test_initialization.py
-    - pytest -v -s tests/models/test_transformers.py
-    - pytest -v -s tests/models/multimodal/processing/
-    - pytest -v -s tests/models/multimodal/test_mapping.py
-    - python3 examples/basic/offline_inference/chat.py
-    - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
-    # Whisper needs spawn method to avoid deadlock
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper

.buildkite/test_areas/models_language.yaml

@@ -50,7 +50,7 @@ steps:
   mirror:
     torch_nightly: {}
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
       commands:
@@ -96,7 +96,7 @@ steps:
     - pytest -v -s models/language/pooling -m 'not core_model'
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 100
       depends_on:
       - image-build-amd

.buildkite/test_areas/models_multimodal.yaml

@@ -15,7 +15,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -33,7 +33,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -50,7 +50,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -118,7 +118,7 @@ steps:
     - pytest -v -s models/multimodal/test_mapping.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 

.buildkite/test_areas/spec_decode.yaml

@@ -32,6 +32,7 @@ steps:
   source_file_dependencies:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
+    - vllm/v1/attention/backends/
     - vllm/transformers_utils/configs/speculators/
     - tests/v1/e2e/spec_decode/
   commands:

.github/CODEOWNERS

@@ -80,13 +80,13 @@
 /tests/distributed/test_multi_node_assignment.py @youkaichao
 /tests/distributed/test_pipeline_parallel.py @youkaichao
 /tests/distributed/test_same_node.py @youkaichao
-/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche
-/tests/evals @mgoin @vadiklyutiy
-/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye
+/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas
+/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas
+/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas
 /tests/kernels/ir @ProExpertProg @tjtanaa
-/tests/models @DarkLight1337 @ywang96
+/tests/models @DarkLight1337 @ywang96 @AndreasKaratzas
 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche
-/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye
+/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas
 /tests/test_inputs.py @DarkLight1337 @ywang96
 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm
 /tests/v1/structured_output @mgoin @russellb @aarnphm
@@ -171,20 +171,20 @@ mkdocs.yaml @hmellor
 
 # ROCm related: specify owner with write access to notify AMD folks for careful code review
 /vllm/**/*rocm* @tjtanaa @dllehr-amd
-/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd
+/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd @AndreasKaratzas
 /vllm/v1/attention/backends/rocm*.py @tjtanaa @dllehr-amd
 /vllm/v1/attention/backends/mla/rocm*.py @tjtanaa @dllehr-amd
 /vllm/v1/attention/ops/rocm*.py @tjtanaa @dllehr-amd
 /vllm/model_executor/layers/fused_moe/rocm*.py @tjtanaa @dllehr-amd
 /csrc/rocm @tjtanaa @dllehr-amd
-/requirements/*rocm* @tjtanaa
-/tests/**/*rocm* @tjtanaa
+/requirements/*rocm* @tjtanaa @AndreasKaratzas
+/tests/**/*rocm* @tjtanaa @AndreasKaratzas
 /docs/**/*rocm* @tjtanaa
 /vllm/**/*quark* @tjtanaa
-/tests/**/*quark* @tjtanaa
+/tests/**/*quark* @tjtanaa @AndreasKaratzas
 /docs/**/*quark* @tjtanaa
-/vllm/**/*aiter* @tjtanaa
-/tests/**/*aiter* @tjtanaa
+/vllm/**/*aiter* @tjtanaa @AndreasKaratzas
+/tests/**/*aiter* @tjtanaa @AndreasKaratzas
 
 # TPU
 /vllm/v1/worker/tpu* @NickLucche

.github/mergify.yml

@@ -103,6 +103,19 @@ pull_request_rules:
       add:
         - frontend
 
+- name: label-rust
+  description: Automatically apply rust label
+  conditions:
+    - label != stale
+    - or:
+      - files~=(?i)rust
+      - title~=(?i)rust
+      - title~=(?i)vllm-rs
+  actions:
+    label:
+      add:
+        - rust
+
 - name: label-llama
   description: Automatically apply llama label
   conditions:

CMakeLists.txt

@@ -305,14 +305,10 @@ endif()
 #
 
 set(VLLM_EXT_SRC
-  "csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
   "csrc/cache_kernels.cu"
   "csrc/cache_kernels_fused.cu"
   "csrc/attention/paged_attention_v1.cu"
   "csrc/attention/paged_attention_v2.cu"
-  "csrc/attention/merge_attn_states.cu"
-  "csrc/sampler.cu"
-  "csrc/topk.cu"
   "csrc/cuda_view.cu"
   "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
   "csrc/quantization/activation_kernels.cu"
@@ -369,16 +365,30 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # are not supported by Machete yet.
 
   # marlin arches for fp16 output
-  cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
+  # Family-conditional 12.0f (one cubin for SM12x family) requires CUDA >= 13.0;
+  # fall back to architecture-specific 12.0a;12.1a on CUDA < 13.0 (e.g. 12.8).
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MARLIN_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
   # marlin has limited support for turing
   cuda_archs_loose_intersection(MARLIN_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
   # marlin arches for bf16 output (we need 9.0 for bf16 atomicAdd PTX)
-  cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MARLIN_BF16_ARCHS "8.0+PTX;9.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
   # marlin arches for fp8 input
   # - sm80 doesn't support fp8 computation
   # - sm90 and sm100 don't support QMMA.16832.F32.E4M3.E4M3 SAAS instruction
   # so we only enable fp8 computation for SM89 (e.g. RTX 40x0)  and 12.0 (e.g. RTX 50x0)
-  cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0;12.1" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MARLIN_FP8_ARCHS "8.9;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
   # marlin arches for other files
   cuda_archs_loose_intersection(MARLIN_OTHER_ARCHS "7.5;8.0+PTX" "${CUDA_ARCHS}")
 
@@ -633,7 +643,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
     "csrc/libtorch_stable/layernorm_kernels.cu"
     "csrc/libtorch_stable/layernorm_quant_kernels.cu"
-    "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu")
+    "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
+    "csrc/libtorch_stable/attention/merge_attn_states.cu"
+    "csrc/libtorch_stable/sampler.cu"
+    "csrc/libtorch_stable/topk.cu"
+    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
@@ -1121,7 +1135,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # moe marlin arches
   # note that we always set `use_atomic_add=False` for moe marlin now,
   # so we don't need 9.0 for bf16 atomicAdd PTX
-  cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX" "${CUDA_ARCHS}")
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(MARLIN_MOE_ARCHS "8.0+PTX;12.0a;12.1a" "${CUDA_ARCHS}")
+  endif()
   # moe marlin has limited support for turing
   cuda_archs_loose_intersection(MARLIN_MOE_SM75_ARCHS "7.5" "${CUDA_ARCHS}")
   # moe marlin arches for fp8 input

benchmarks/kernels/benchmark_moe_permute_unpermute.py

@@ -10,6 +10,7 @@ from transformers import AutoConfig
 
 from vllm.model_executor.layers.fused_moe import fused_topk
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
+    MoEPermuteScratch,
     moe_permute,
     moe_unpermute,
 )
@@ -54,6 +55,15 @@ def benchmark_permute(
     topk_weights, topk_ids, token_expert_indices = fused_topk(
         qhidden_states, input_gating, topk, False
     )
+    scratch = MoEPermuteScratch(
+        max_num_tokens=num_tokens,
+        topk=topk,
+        num_experts=num_experts,
+        num_local_experts=num_experts,
+        device=qhidden_states.device,
+        hidden_size=hidden_size,
+        hidden_dtype=qhidden_states.dtype,
+    )
 
     def prepare(i: int):
         input_gating.copy_(gating_output[i])
@@ -65,6 +75,7 @@ def benchmark_permute(
             topk_ids=topk_ids,
             n_expert=num_experts,
             expert_map=None,
+            scratch=scratch,
         )
 
     # JIT compilation & warmup
@@ -123,6 +134,15 @@ def benchmark_unpermute(
     topk_weights, topk_ids, token_expert_indices = fused_topk(
         qhidden_states, input_gating, topk, False
     )
+    scratch = MoEPermuteScratch(
+        max_num_tokens=num_tokens,
+        topk=topk,
+        num_experts=num_experts,
+        num_local_experts=num_experts,
+        device=qhidden_states.device,
+        hidden_size=hidden_size,
+        hidden_dtype=qhidden_states.dtype,
+    )
 
     def prepare():
         (
@@ -137,6 +157,7 @@ def benchmark_unpermute(
             topk_ids=topk_ids,
             n_expert=num_experts,
             expert_map=None,
+            scratch=scratch,
         )
         # convert to fp16/bf16 as gemm output
         return (

csrc/libtorch_stable/attention/merge_attn_states.cu

@@ -1,14 +1,14 @@
 #include <optional>
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
 #include <algorithm>
 #include <limits>
 
-#include "attention_dtypes.h"
-#include "attention_utils.cuh"
-#include "../quantization/w8a8/fp8/common.cuh"
+#include "../torch_utils.h"
 #include "../dispatch_utils.h"
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "../../attention/attention_dtypes.h"
+#include "../../attention/attention_utils.cuh"
+#include "../../quantization/w8a8/fp8/common.cuh"
 
 namespace vllm {
 
@@ -196,17 +196,17 @@ __global__ void merge_attn_states_kernel(
 // The following macro is used to dispatch the conversion function based on
 // the output data type. The FN is a macro that calls a function with
 // template<typename scalar_t>.
-#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn)                      \
-  {                                                                     \
-    if (scalar_dtype == at::ScalarType::Float) {                        \
-      fn(float);                                                        \
-    } else if (scalar_dtype == at::ScalarType::Half) {                  \
-      fn(uint16_t);                                                     \
-    } else if (scalar_dtype == at::ScalarType::BFloat16) {              \
-      fn(__nv_bfloat16);                                                \
-    } else {                                                            \
-      TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
-    }                                                                   \
+#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn)                          \
+  {                                                                         \
+    if (scalar_dtype == torch::headeronly::ScalarType::Float) {             \
+      fn(float);                                                            \
+    } else if (scalar_dtype == torch::headeronly::ScalarType::Half) {       \
+      fn(uint16_t);                                                         \
+    } else if (scalar_dtype == torch::headeronly::ScalarType::BFloat16) {   \
+      fn(__nv_bfloat16);                                                    \
+    } else {                                                                \
+      STD_TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
+    }                                                                       \
   }
 
 #define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS,           \
@@ -245,11 +245,14 @@ __global__ void merge_attn_states_kernel(
  */
 template <typename scalar_t>
 void merge_attn_states_launcher(
-    torch::Tensor& output, std::optional<torch::Tensor> output_lse,
-    const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
-    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
+    torch::stable::Tensor& output,
+    std::optional<torch::stable::Tensor> output_lse,
+    const torch::stable::Tensor& prefix_output,
+    const torch::stable::Tensor& prefix_lse,
+    const torch::stable::Tensor& suffix_output,
+    const torch::stable::Tensor& suffix_lse,
     const std::optional<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::Tensor>& output_scale) {
+    const std::optional<torch::stable::Tensor>& output_scale) {
   constexpr uint NUM_THREADS = 128;
   const uint num_tokens = output.size(0);
   const uint num_heads = output.size(1);
@@ -258,23 +261,23 @@ void merge_attn_states_launcher(
   const uint output_head_stride = output.stride(1);
   // Thread mapping is based on input BF16 pack_size
   const uint pack_size = 16 / sizeof(scalar_t);
-  TORCH_CHECK(head_size % pack_size == 0,
-              "headsize must be multiple of pack_size:", pack_size);
+  STD_TORCH_CHECK(head_size % pack_size == 0,
+                  "headsize must be multiple of pack_size:", pack_size);
 
   const uint prefix_num_tokens =
       prefill_tokens_with_context.has_value()
           ? static_cast<uint>(prefill_tokens_with_context.value())
           : num_tokens;
-  TORCH_CHECK(prefix_num_tokens <= num_tokens,
-              "prefix_num_tokens must be <= num_tokens");
+  STD_TORCH_CHECK(prefix_num_tokens <= num_tokens,
+                  "prefix_num_tokens must be <= num_tokens");
 
   float* output_lse_ptr = nullptr;
   if (output_lse.has_value()) {
-    output_lse_ptr = output_lse.value().data_ptr<float>();
+    output_lse_ptr = output_lse.value().mutable_data_ptr<float>();
   }
   float* output_scale_ptr = nullptr;
   if (output_scale.has_value()) {
-    output_scale_ptr = output_scale.value().data_ptr<float>();
+    output_scale_ptr = output_scale.value().mutable_data_ptr<float>();
   }
   // Process one pack elements per thread. for float, the
   // pack_size is 4 for half/bf16, the pack_size is 8.
@@ -284,14 +287,15 @@ void merge_attn_states_launcher(
   dim3 block(NUM_THREADS);
   dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);
 
-  const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
-  auto stream = at::cuda::getCurrentCUDAStream();
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      prefix_output.get_device_index());
+  auto stream = get_current_cuda_stream();
 
   if (output_scale.has_value()) {
     // FP8 output path - dispatch on output FP8 type
-    VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] {
-      LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true);
-    });
+    VLLM_STABLE_DISPATCH_FP8_TYPES(
+        output.scalar_type(), "merge_attn_states_fp8",
+        [&] { LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true); });
   } else {
     // Original BF16/FP16/FP32 output path
     LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false);
@@ -305,26 +309,29 @@ void merge_attn_states_launcher(
         suffix_lse, prefill_tokens_with_context, output_scale);       \
   }
 
-void merge_attn_states(torch::Tensor& output,
-                       std::optional<torch::Tensor> output_lse,
-                       const torch::Tensor& prefix_output,
-                       const torch::Tensor& prefix_lse,
-                       const torch::Tensor& suffix_output,
-                       const torch::Tensor& suffix_lse,
-                       std::optional<int64_t> prefill_tokens_with_context,
-                       const std::optional<torch::Tensor>& output_scale) {
+void merge_attn_states(
+    torch::stable::Tensor& output,
+    std::optional<torch::stable::Tensor> output_lse,
+    const torch::stable::Tensor& prefix_output,
+    const torch::stable::Tensor& prefix_lse,
+    const torch::stable::Tensor& suffix_output,
+    const torch::stable::Tensor& suffix_lse,
+    const std::optional<int64_t> prefill_tokens_with_context,
+    const std::optional<torch::stable::Tensor>& output_scale) {
   if (output_scale.has_value()) {
-    TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn ||
-                    output.scalar_type() == at::ScalarType::Float8_e4m3fnuz,
-                "output must be FP8 when output_scale is provided, got: ",
-                output.scalar_type());
+    STD_TORCH_CHECK(
+        output.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+            output.scalar_type() ==
+                torch::headeronly::ScalarType::Float8_e4m3fnuz,
+        "output must be FP8 when output_scale is provided, got: ",
+        output.scalar_type());
   } else {
-    TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(),
-                "output dtype (", output.scalar_type(),
-                ") must match prefix_output dtype (",
-                prefix_output.scalar_type(), ") when output_scale is not set");
+    STD_TORCH_CHECK(
+        output.scalar_type() == prefix_output.scalar_type(), "output dtype (",
+        output.scalar_type(), ") must match prefix_output dtype (",
+        prefix_output.scalar_type(), ") when output_scale is not set");
   }
   // Always dispatch on prefix_output (input) dtype
-  DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(),
+  DISPATCH_BY_SCALAR_DTYPE(prefix_output.scalar_type(),
                            CALL_MERGE_ATTN_STATES_LAUNCHER);
 }

csrc/libtorch_stable/mamba/selective_scan.h

@@ -12,6 +12,9 @@
     #include <hip/hip_bf16.h>
 #endif
 #include <cuda_fp16.h>
+
+#include <torch/headeronly/util/Half.h>
+#include <torch/headeronly/util/BFloat16.h>
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
 struct SSMParamsBase {
@@ -159,8 +162,8 @@ struct Converter{
 };
 
 template<int N>
-struct Converter<at::Half, N>{
-    static inline __device__ void to_float(const at::Half (&src)[N], float (&dst)[N]) {
+struct Converter<torch::headeronly::Half, N>{
+    static inline __device__ void to_float(const torch::headeronly::Half (&src)[N], float (&dst)[N]) {
         static_assert(N % 2 == 0);
         auto &src2 = reinterpret_cast<const half2 (&)[N / 2]>(src);
         auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
@@ -171,8 +174,8 @@ struct Converter<at::Half, N>{
 
 #if __CUDA_ARCH__ >= 800
 template<int N>
-struct Converter<at::BFloat16, N>{
-    static inline __device__ void to_float(const at::BFloat16 (&src)[N], float (&dst)[N]) {
+struct Converter<torch::headeronly::BFloat16, N>{
+    static inline __device__ void to_float(const torch::headeronly::BFloat16 (&src)[N], float (&dst)[N]) {
         static_assert(N % 2 == 0);
         auto &src2 = reinterpret_cast<const nv_bfloat162 (&)[N / 2]>(src);
         auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);

csrc/libtorch_stable/mamba/selective_scan_fwd.cu

@@ -1,18 +1,9 @@
 // clang-format off
 // adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/selective_scan_fwd_kernel.cuh
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "../torch_utils.h"
+#include <torch/csrc/stable/macros.h>
 #include "selective_scan.h"
 
-#include <c10/util/BFloat16.h>
-#include <c10/util/Half.h>
-#ifdef USE_ROCM
-    #include <c10/hip/HIPException.h>  // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK
-#else
-    #include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
-#endif
-
 #ifndef USE_ROCM
     #include <cub/block/block_load.cuh>
     #include <cub/block/block_store.cuh>
@@ -416,15 +407,15 @@ void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
                 auto kernel = &selective_scan_fwd_kernel<Ktraits>;
                 if (kSmemSize >= 48 * 1024) {
 #ifdef USE_ROCM
-                    C10_HIP_CHECK(hipFuncSetAttribute(
+                    STD_CUDA_CHECK(hipFuncSetAttribute(
                         reinterpret_cast<const void*>(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
 #else
-                    C10_CUDA_CHECK(cudaFuncSetAttribute(
+                    STD_CUDA_CHECK(cudaFuncSetAttribute(
                         kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
 #endif
                 }
                 kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
-                C10_CUDA_KERNEL_LAUNCH_CHECK();
+                STD_CUDA_KERNEL_LAUNCH_CHECK();
             });
         });
     });
@@ -462,46 +453,46 @@ void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
     #endif
 }
 
-template void selective_scan_fwd_cuda<at::BFloat16, float, at::BFloat16>(SSMParamsBase &params, cudaStream_t stream);
-template void selective_scan_fwd_cuda<at::BFloat16, float, float>(SSMParamsBase &params, cudaStream_t stream);
-template void selective_scan_fwd_cuda<at::Half, float, at::Half>(SSMParamsBase &params, cudaStream_t stream);
-template void selective_scan_fwd_cuda<at::Half, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, torch::headeronly::BFloat16>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<torch::headeronly::BFloat16, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<torch::headeronly::Half, float, torch::headeronly::Half>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<torch::headeronly::Half, float, float>(SSMParamsBase &params, cudaStream_t stream);
 template void selective_scan_fwd_cuda<float, float, float>(SSMParamsBase &params, cudaStream_t stream);
 
-#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+#define CHECK_SHAPE(x, ...) STD_TORCH_CHECK(x.sizes().equals(torch::headeronly::IntHeaderOnlyArrayRef({__VA_ARGS__})), #x " must have shape (" #__VA_ARGS__ ")")
 
 #define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, STYPE, NAME, ...)       \
-    if (ITYPE == at::ScalarType::Half) {                                            \
-        using input_t = at::Half;                                                   \
+    if (ITYPE == torch::headeronly::ScalarType::Half) {                             \
+        using input_t = torch::headeronly::Half;                                                  \
         using weight_t = float;                                                     \
-        if (STYPE == at::ScalarType::Half) {                                        \
-            using state_t = at::Half;                                               \
+        if (STYPE == torch::headeronly::ScalarType::Half) {                         \
+            using state_t = torch::headeronly::Half;                                              \
             __VA_ARGS__();                                                          \
-        } else if (STYPE == at::ScalarType::Float) {                                \
+        } else if (STYPE == torch::headeronly::ScalarType::Float) {                 \
             using state_t = float;                                                  \
             __VA_ARGS__();                                                          \
         } else {                                                                    \
-            AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
+            STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \
         }                                                                           \
-    } else if (ITYPE == at::ScalarType::BFloat16) {                                 \
-        using input_t = at::BFloat16;                                               \
+    } else if (ITYPE == torch::headeronly::ScalarType::BFloat16) {                  \
+        using input_t = torch::headeronly::BFloat16;                                              \
         using weight_t = float;                                                     \
-        if (STYPE == at::ScalarType::BFloat16) {                                    \
-            using state_t = at::BFloat16;                                           \
+        if (STYPE == torch::headeronly::ScalarType::BFloat16) {                     \
+            using state_t = torch::headeronly::BFloat16;                                          \
             __VA_ARGS__();                                                          \
-        } else if (STYPE == at::ScalarType::Float) {                                \
+        } else if (STYPE == torch::headeronly::ScalarType::Float) {                 \
             using state_t = float;                                                  \
             __VA_ARGS__();                                                          \
         } else {                                                                    \
-            AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \
+            STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \
         }                                                                           \
-    } else if (ITYPE == at::ScalarType::Float)  {                                   \
+    } else if (ITYPE == torch::headeronly::ScalarType::Float)  {                    \
         using input_t = float;                                                      \
         using weight_t = float;                                                     \
         using state_t = float;                                                      \
         __VA_ARGS__();                                                              \
     } else {                                                                        \
-        AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
+        STD_TORCH_CHECK(false, #NAME " not implemented for input type '", ITYPE, "'"); \
     }
 
 
@@ -518,30 +509,30 @@ void set_ssm_params_fwd(SSMParamsBase &params,
                         const bool is_variable_B,
                         const bool is_variable_C,
                         // device pointers
-                        const torch::Tensor u,
-                        const torch::Tensor delta,
-                        const torch::Tensor A,
-                        const torch::Tensor B,
-                        const torch::Tensor C,
-                        const torch::Tensor out,
-                        const torch::Tensor z,
-                        const torch::Tensor out_z,
-                        const std::optional<at::Tensor>& D,
-                        const std::optional<at::Tensor>& delta_bias,
-                        const torch::Tensor ssm_states,
+                        const torch::stable::Tensor u,
+                        const torch::stable::Tensor delta,
+                        const torch::stable::Tensor A,
+                        const torch::stable::Tensor B,
+                        const torch::stable::Tensor C,
+                        const torch::stable::Tensor out,
+                        const torch::stable::Tensor z,
+                        const torch::stable::Tensor out_z,
+                        const std::optional<torch::stable::Tensor>& D,
+                        const std::optional<torch::stable::Tensor>& delta_bias,
+                        const torch::stable::Tensor ssm_states,
                         bool has_z,
                         bool delta_softplus,
-                        const std::optional<at::Tensor>& query_start_loc,
-                        const std::optional<at::Tensor>& cache_indices,
-                        const std::optional<at::Tensor>& has_initial_state,
+                        const std::optional<torch::stable::Tensor>& query_start_loc,
+                        const std::optional<torch::stable::Tensor>& cache_indices,
+                        const std::optional<torch::stable::Tensor>& has_initial_state,
                         bool varlen,
                         int64_t null_block_id,
                         int64_t block_size,
-                        const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
-                        const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
-                        const std::optional<torch::Tensor> &initial_state_idx,
-                        const std::optional<torch::Tensor> &cu_chunk_seqlen,
-                        const std::optional<torch::Tensor> &last_chunk_indices) {
+                        const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token,
+                        const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token,
+                        const std::optional<torch::stable::Tensor> &initial_state_idx,
+                        const std::optional<torch::stable::Tensor> &cu_chunk_seqlen,
+                        const std::optional<torch::stable::Tensor> &last_chunk_indices) {
 
     // Reset the parameters
     memset(&params, 0, sizeof(params));
@@ -654,45 +645,45 @@ void set_ssm_params_fwd(SSMParamsBase &params,
     }
 }
 
-void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
-                  const torch::Tensor &A, const torch::Tensor &B, const torch::Tensor &C,
-                  const std::optional<torch::Tensor> &D_,
-                  const std::optional<torch::Tensor> &z_,
-                  const std::optional<torch::Tensor> &delta_bias_,
+void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Tensor &delta,
+                  const torch::stable::Tensor &A, const torch::stable::Tensor &B, const torch::stable::Tensor &C,
+                  const std::optional<torch::stable::Tensor> &D_,
+                  const std::optional<torch::stable::Tensor> &z_,
+                  const std::optional<torch::stable::Tensor> &delta_bias_,
                   bool delta_softplus,
-                  const std::optional<torch::Tensor> &query_start_loc,
-                  const std::optional<torch::Tensor> &cache_indices,
-                  const std::optional<torch::Tensor> &has_initial_state,
-                  const torch::Tensor &ssm_states,
+                  const std::optional<torch::stable::Tensor> &query_start_loc,
+                  const std::optional<torch::stable::Tensor> &cache_indices,
+                  const std::optional<torch::stable::Tensor> &has_initial_state,
+                  const torch::stable::Tensor &ssm_states,
                   // used to identify padding entries if cache_indices provided
                   // in case of padding, the kernel will return early
                   int64_t null_block_id,
                   int64_t block_size,
-                  const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
-                  const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
-                  const std::optional<torch::Tensor> &initial_state_idx,
-                  const std::optional<torch::Tensor> &cu_chunk_seqlen,
-                  const std::optional<torch::Tensor> &last_chunk_indices) {
+                  const std::optional<torch::stable::Tensor> &block_idx_first_scheduled_token,
+                  const std::optional<torch::stable::Tensor> &block_idx_last_scheduled_token,
+                  const std::optional<torch::stable::Tensor> &initial_state_idx,
+                  const std::optional<torch::stable::Tensor> &cu_chunk_seqlen,
+                  const std::optional<torch::stable::Tensor> &last_chunk_indices) {
     auto input_type = u.scalar_type();
     auto weight_type = A.scalar_type();
-    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
-    TORCH_CHECK(weight_type == at::ScalarType::Float);
+    STD_TORCH_CHECK(input_type == torch::headeronly::ScalarType::Float || input_type == torch::headeronly::ScalarType::Half || input_type == torch::headeronly::ScalarType::BFloat16);
+    STD_TORCH_CHECK(weight_type == torch::headeronly::ScalarType::Float);
 
     const bool is_variable_B = B.dim() >= 3;
     const bool is_variable_C = C.dim() >= 3;
 
-    TORCH_CHECK(delta.scalar_type() == input_type);
-    TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
-    TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
+    STD_TORCH_CHECK(delta.scalar_type() == input_type);
+    STD_TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type));
+    STD_TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type));
 
-    TORCH_CHECK(u.is_cuda());
-    TORCH_CHECK(delta.is_cuda());
-    TORCH_CHECK(A.is_cuda());
-    TORCH_CHECK(B.is_cuda());
-    TORCH_CHECK(C.is_cuda());
+    STD_TORCH_CHECK(u.is_cuda());
+    STD_TORCH_CHECK(delta.is_cuda());
+    STD_TORCH_CHECK(A.is_cuda());
+    STD_TORCH_CHECK(B.is_cuda());
+    STD_TORCH_CHECK(C.is_cuda());
 
-    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
-    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+    STD_TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    STD_TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
 
     const auto sizes = u.sizes();
     const bool varlen = query_start_loc.has_value();
@@ -702,7 +693,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
     const int dstate = A.size(1);
     const int n_groups = varlen ? B.size(0) : B.size(1);
 
-    TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
+    STD_TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256");
 
     if (varlen) {
         CHECK_SHAPE(u, dim, seqlen);
@@ -712,94 +703,94 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
         CHECK_SHAPE(delta, batch_size, dim, seqlen);
     }
     CHECK_SHAPE(A, dim, dstate);
-    TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size")
+    STD_TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size");
     if (varlen) {
         CHECK_SHAPE(B, n_groups, dstate, seqlen);
     } else {
-        CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen); 
+        CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
     }
-    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    STD_TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
 
-    TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size")
+    STD_TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size");
     if (varlen) {
         CHECK_SHAPE(C, n_groups, dstate, seqlen);
     } else {
-        CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen); 
+        CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
     }
-    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+    STD_TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
 
     if (D_.has_value()) {
         auto D = D_.value();
-        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
-        TORCH_CHECK(D.is_cuda());
-        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        STD_TORCH_CHECK(D.scalar_type() == torch::headeronly::ScalarType::Float);
+        STD_TORCH_CHECK(D.is_cuda());
+        STD_TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
         CHECK_SHAPE(D, dim);
     }
 
     if (delta_bias_.has_value()) {
         auto delta_bias = delta_bias_.value();
-        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
-        TORCH_CHECK(delta_bias.is_cuda());
-        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        STD_TORCH_CHECK(delta_bias.scalar_type() == torch::headeronly::ScalarType::Float);
+        STD_TORCH_CHECK(delta_bias.is_cuda());
+        STD_TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
         CHECK_SHAPE(delta_bias, dim);
     }
 
 
     if (has_initial_state.has_value()) {
         auto has_initial_state_ = has_initial_state.value();
-        TORCH_CHECK(has_initial_state_.scalar_type() == at::ScalarType::Bool);
-        TORCH_CHECK(has_initial_state_.is_cuda());
+        STD_TORCH_CHECK(has_initial_state_.scalar_type() == torch::headeronly::ScalarType::Bool);
+        STD_TORCH_CHECK(has_initial_state_.is_cuda());
         CHECK_SHAPE(has_initial_state_, batch_size);
     }
 
 
     if (query_start_loc.has_value()) {
         auto query_start_loc_ = query_start_loc.value();
-        TORCH_CHECK(query_start_loc_.scalar_type() == at::ScalarType::Int);
-        TORCH_CHECK(query_start_loc_.is_cuda());
+        STD_TORCH_CHECK(query_start_loc_.scalar_type() == torch::headeronly::ScalarType::Int);
+        STD_TORCH_CHECK(query_start_loc_.is_cuda());
     }
 
 
     if (cache_indices.has_value()) {
         auto cache_indices_ = cache_indices.value();
-        TORCH_CHECK(cache_indices_.scalar_type() == at::ScalarType::Int);
-        TORCH_CHECK(cache_indices_.is_cuda());
+        STD_TORCH_CHECK(cache_indices_.scalar_type() == torch::headeronly::ScalarType::Int);
+        STD_TORCH_CHECK(cache_indices_.is_cuda());
 
         // cache_indices can be either 1D (batch_size,) for non-APC mode
         // or 2D (batch_size, max_positions) for APC mode
         const bool is_apc_mode = block_idx_first_scheduled_token.has_value();
         if (is_apc_mode) {
-            TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
-            TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
+            STD_TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode");
+            STD_TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size");
         } else {
             CHECK_SHAPE(cache_indices_, batch_size);
         }
     }
-   
 
-    at::Tensor z, out_z;
+
+    torch::stable::Tensor z, out_z;
     const bool has_z = z_.has_value();
     if (has_z) {
         z = z_.value();
-        TORCH_CHECK(z.scalar_type() == input_type);
-        TORCH_CHECK(z.is_cuda());
-        TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
+        STD_TORCH_CHECK(z.scalar_type() == input_type);
+        STD_TORCH_CHECK(z.is_cuda());
+        STD_TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1);
         if (varlen){
             CHECK_SHAPE(z, dim, seqlen);
         } else {
             CHECK_SHAPE(z, batch_size, dim, seqlen);
         }
-        
+
         out_z = z;
     }
 
     // Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout
-    at::Tensor out = delta;
+    torch::stable::Tensor out = delta;
     // ssm_states can now be either the same as input_type or float32
     auto state_type = ssm_states.scalar_type();
-    TORCH_CHECK(state_type == input_type || state_type == at::ScalarType::Float);
-    TORCH_CHECK(ssm_states.is_cuda());
-    TORCH_CHECK(ssm_states.stride(-1) == 1);
+    STD_TORCH_CHECK(state_type == input_type || state_type == torch::headeronly::ScalarType::Float);
+    STD_TORCH_CHECK(ssm_states.is_cuda());
+    STD_TORCH_CHECK(ssm_states.stride(-1) == 1);
 
     SSMParamsBase params;
     set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, is_variable_B, is_variable_C,
@@ -823,8 +814,8 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
                        );
 
     
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(u));
-    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    const torch::stable::accelerator::DeviceGuard device_guard(u.get_device_index());
+    auto stream = get_current_cuda_stream();
     DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), ssm_states.scalar_type(), "selective_scan_fwd", [&] {
         selective_scan_fwd_cuda<input_t, weight_t, state_t>(params, stream);
     });

csrc/libtorch_stable/ops.h

@@ -164,6 +164,17 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
 
 #endif
 
+// Attention kernels (shared CUDA/ROCm)
+void merge_attn_states(
+    torch::stable::Tensor& output,
+    std::optional<torch::stable::Tensor> output_lse,
+    const torch::stable::Tensor& prefix_output,
+    const torch::stable::Tensor& prefix_lse,
+    const torch::stable::Tensor& suffix_output,
+    const torch::stable::Tensor& suffix_lse,
+    const std::optional<int64_t> prefill_tokens_with_context,
+    const std::optional<torch::stable::Tensor>& output_scale = std::nullopt);
+
 torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
                                          bool inplace);
 
@@ -220,6 +231,48 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
                         torch::stable::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
+// Sampler kernels (shared CUDA/ROCm)
+void apply_repetition_penalties_(
+    torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
+    const torch::stable::Tensor& output_mask,
+    const torch::stable::Tensor& repetition_penalties);
+
+void top_k_per_row_prefill(const torch::stable::Tensor& logits,
+                           const torch::stable::Tensor& rowStarts,
+                           const torch::stable::Tensor& rowEnds,
+                           torch::stable::Tensor& indices, int64_t numRows,
+                           int64_t stride0, int64_t stride1, int64_t topK);
+
+void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
+                          const torch::stable::Tensor& seqLens,
+                          torch::stable::Tensor& indices, int64_t numRows,
+                          int64_t stride0, int64_t stride1, int64_t topK);
+
+void persistent_topk(const torch::stable::Tensor& logits,
+                     const torch::stable::Tensor& lengths,
+                     torch::stable::Tensor& output,
+                     torch::stable::Tensor& workspace, int64_t k,
+                     int64_t max_seq_len);
+
+void selective_scan_fwd(
+    const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
+    const torch::stable::Tensor& A, const torch::stable::Tensor& B,
+    const torch::stable::Tensor& C,
+    const std::optional<torch::stable::Tensor>& D_,
+    const std::optional<torch::stable::Tensor>& z_,
+    const std::optional<torch::stable::Tensor>& delta_bias_,
+    bool delta_softplus,
+    const std::optional<torch::stable::Tensor>& query_start_loc,
+    const std::optional<torch::stable::Tensor>& cache_indices,
+    const std::optional<torch::stable::Tensor>& has_initial_state,
+    const torch::stable::Tensor& ssm_states, int64_t null_block_id,
+    int64_t block_size,
+    const std::optional<torch::stable::Tensor>& block_idx_first_scheduled_token,
+    const std::optional<torch::stable::Tensor>& block_idx_last_scheduled_token,
+    const std::optional<torch::stable::Tensor>& initial_state_idx,
+    const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
+    const std::optional<torch::stable::Tensor>& last_chunk_indices);
+
 // Activation kernels (shared CUDA/ROCm)
 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,

csrc/libtorch_stable/sampler.cu

@@ -1,8 +1,6 @@
-#include "cuda_compat.h"
+#include "../cuda_compat.h"
 #include "dispatch_utils.h"
-
-#include <torch/cuda.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "torch_utils.h"
 
 #ifndef USE_ROCM
   #include <cub/cub.cuh>
@@ -618,14 +616,14 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
 }  // namespace vllm
 
 void apply_repetition_penalties_(
-    torch::Tensor& logits,             // [num_seqs, vocab_size], in-place
-    const torch::Tensor& prompt_mask,  // [num_seqs, vocab_size]
-    const torch::Tensor& output_mask,  // [num_seqs, vocab_size]
-    const torch::Tensor& repetition_penalties) {  // [num_seqs]
-  TORCH_CHECK(logits.is_contiguous());
-  TORCH_CHECK(prompt_mask.is_contiguous());
-  TORCH_CHECK(output_mask.is_contiguous());
-  TORCH_CHECK(repetition_penalties.is_contiguous());
+    torch::stable::Tensor& logits,  // [num_seqs, vocab_size], in-place
+    const torch::stable::Tensor& prompt_mask,  // [num_seqs, vocab_size]
+    const torch::stable::Tensor& output_mask,  // [num_seqs, vocab_size]
+    const torch::stable::Tensor& repetition_penalties) {  // [num_seqs]
+  STD_TORCH_CHECK(logits.is_contiguous());
+  STD_TORCH_CHECK(prompt_mask.is_contiguous());
+  STD_TORCH_CHECK(output_mask.is_contiguous());
+  STD_TORCH_CHECK(repetition_penalties.is_contiguous());
 
   int vocab_size = logits.size(-1);
   int num_seqs = logits.size(0);
@@ -635,7 +633,7 @@ void apply_repetition_penalties_(
   // Get number of SMs on the current device
   int sms = 0;
   cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount,
-                         logits.get_device());
+                         logits.get_device_index());
 
   // Compute tile_num and tile_size
   int tile_num =
@@ -645,27 +643,29 @@ void apply_repetition_penalties_(
   // Each block handles one sequence and a tile of vocab
   dim3 grid(num_seqs, tile_num);
   dim3 block(std::min(tile_size, 1024));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(logits));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  VLLM_DISPATCH_FLOATING_TYPES(
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      logits.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream();
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(
       logits.scalar_type(), "apply_repetition_penalties_kernel", [&] {
         vllm::apply_repetition_penalties_kernel<scalar_t>
             <<<grid, block, 0, stream>>>(
-                logits.data_ptr<scalar_t>(), prompt_mask.data_ptr<bool>(),
-                output_mask.data_ptr<bool>(),
-                repetition_penalties.data_ptr<scalar_t>(), num_seqs, vocab_size,
-                tile_size);
+                logits.mutable_data_ptr<scalar_t>(),
+                prompt_mask.const_data_ptr<bool>(),
+                output_mask.const_data_ptr<bool>(),
+                repetition_penalties.const_data_ptr<scalar_t>(), num_seqs,
+                vocab_size, tile_size);
       });
 }
 
-void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
-                          const torch::Tensor& seqLens, torch::Tensor& indices,
-                          int64_t numRows, int64_t stride0, int64_t stride1,
-                          int64_t topK) {
+void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
+                          const torch::stable::Tensor& seqLens,
+                          torch::stable::Tensor& indices, int64_t numRows,
+                          int64_t stride0, int64_t stride1, int64_t topK) {
   constexpr int kSortingAlgorithmThreshold = 12288;
   constexpr int kSplitWorkThreshold = 200 * 1000;
   constexpr int kNumThreadsPerBlock = 512;
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
   const auto numColumns = logits.size(1);
 
   // True if seqLens is 2D (B, next_n): each logit row has its own pre-computed
@@ -677,73 +677,76 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
     // Use insertion sort
     vllm::topKPerRowDecode<kNumThreadsPerBlock, false>
         <<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
-            logits.data_ptr<float>(), seqLens.data_ptr<int>(),
-            indices.data_ptr<int>(), static_cast<int>(stride0),
+            logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
             static_cast<int>(stride1), static_cast<int>(topK),
             static_cast<int>(next_n), seqLensIs2D);
   } else if (numColumns < kSplitWorkThreshold) {
     // From this threshold, use radix sort instead
     vllm::topKPerRowDecode<kNumThreadsPerBlock, true>
         <<<numRows, kNumThreadsPerBlock, topK * sizeof(int32_t), stream>>>(
-            logits.data_ptr<float>(), seqLens.data_ptr<int>(),
-            indices.data_ptr<int>(), static_cast<int>(stride0),
+            logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
             static_cast<int>(stride1), static_cast<int>(topK),
             static_cast<int>(next_n), seqLensIs2D);
   } else {
     // Long sequences are run in two steps
     constexpr auto multipleBlocksPerRowConfig = 10;
 
-    const auto outIndicesAux =
-        torch::empty({numRows, multipleBlocksPerRowConfig, topK},
-                     torch::dtype(torch::kInt32).device(logits.device()));
-    const auto outLogitsAux =
-        torch::empty({numRows, multipleBlocksPerRowConfig, topK},
-                     torch::dtype(torch::kFloat).device(logits.device()));
+    const auto outIndicesAux = torch::stable::empty(
+        {numRows, multipleBlocksPerRowConfig, topK},
+        torch::headeronly::ScalarType::Int, std::nullopt, logits.device());
+    const auto outLogitsAux = torch::stable::empty(
+        {numRows, multipleBlocksPerRowConfig, topK},
+        torch::headeronly::ScalarType::Float, std::nullopt, logits.device());
 
     vllm::topKPerRowDecode<kNumThreadsPerBlock, true, true>
         <<<dim3(numRows, multipleBlocksPerRowConfig), kNumThreadsPerBlock,
            2 * topK * sizeof(int32_t), stream>>>(
-            logits.data_ptr<float>(), seqLens.data_ptr<int>(),
-            outIndicesAux.data_ptr<int>(), static_cast<int>(stride0),
+            logits.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            outIndicesAux.mutable_data_ptr<int>(), static_cast<int>(stride0),
             static_cast<int>(stride1), static_cast<int>(topK),
             static_cast<int>(next_n), seqLensIs2D,
-            outLogitsAux.data_ptr<float>());
+            outLogitsAux.mutable_data_ptr<float>());
 
     constexpr int kNumThreadsPerBlockMerge = 1024;
     vllm::topKPerRowDecode<kNumThreadsPerBlockMerge, true, false, true>
         <<<numRows, kNumThreadsPerBlockMerge, topK * sizeof(int32_t), stream>>>(
-            outLogitsAux.data_ptr<float>(), seqLens.data_ptr<int>(),
-            indices.data_ptr<int>(), multipleBlocksPerRowConfig * topK, 1,
-            static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,
-            nullptr, multipleBlocksPerRowConfig, outIndicesAux.data_ptr<int>());
+            outLogitsAux.const_data_ptr<float>(), seqLens.const_data_ptr<int>(),
+            indices.mutable_data_ptr<int>(), multipleBlocksPerRowConfig * topK,
+            1, static_cast<int>(topK), static_cast<int>(next_n), seqLensIs2D,
+            nullptr, multipleBlocksPerRowConfig,
+            outIndicesAux.const_data_ptr<int>());
   }
 }
 
-void top_k_per_row_prefill(const torch::Tensor& logits,
-                           const torch::Tensor& rowStarts,
-                           const torch::Tensor& rowEnds, torch::Tensor& indices,
-                           int64_t numRows, int64_t stride0, int64_t stride1,
-                           int64_t topK) {
+void top_k_per_row_prefill(const torch::stable::Tensor& logits,
+                           const torch::stable::Tensor& rowStarts,
+                           const torch::stable::Tensor& rowEnds,
+                           torch::stable::Tensor& indices, int64_t numRows,
+                           int64_t stride0, int64_t stride1, int64_t topK) {
   constexpr int kSortingAlgorithmThreshold = 12288;
   constexpr int kNumThreadsPerBlock = 512;
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
 
   int numInsertionBlocks =
       std::min(static_cast<int>(numRows), kSortingAlgorithmThreshold);
   vllm::topKPerRowPrefill<kNumThreadsPerBlock, false>
       <<<numInsertionBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
-         stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
-                   rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
-                   static_cast<int>(stride0), static_cast<int>(stride1),
-                   static_cast<int>(topK), 0);
+         stream>>>(logits.const_data_ptr<float>(),
+                   rowStarts.const_data_ptr<int>(),
+                   rowEnds.const_data_ptr<int>(),
+                   indices.mutable_data_ptr<int>(), static_cast<int>(stride0),
+                   static_cast<int>(stride1), static_cast<int>(topK), 0);
 
   if (numRows > kSortingAlgorithmThreshold) {
     int numRadixBlocks = numRows - kSortingAlgorithmThreshold;
     vllm::topKPerRowPrefill<kNumThreadsPerBlock, true>
         <<<numRadixBlocks, kNumThreadsPerBlock, topK * sizeof(int32_t),
-           stream>>>(logits.data_ptr<float>(), rowStarts.data_ptr<int>(),
-                     rowEnds.data_ptr<int>(), indices.data_ptr<int>(),
-                     static_cast<int>(stride0), static_cast<int>(stride1),
-                     static_cast<int>(topK), kSortingAlgorithmThreshold);
+           stream>>>(
+            logits.const_data_ptr<float>(), rowStarts.const_data_ptr<int>(),
+            rowEnds.const_data_ptr<int>(), indices.mutable_data_ptr<int>(),
+            static_cast<int>(stride0), static_cast<int>(stride1),
+            static_cast<int>(topK), kSortingAlgorithmThreshold);
   }
 }

csrc/libtorch_stable/topk.cu

@@ -1,49 +1,51 @@
 // Persistent TopK kernel for DeepSeek V3 sparse attention indexer.
 // See persistent_topk.cuh for kernel implementation.
 
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
 #include <cuda_runtime.h>
 #include <algorithm>
 
+#include "torch_utils.h"
+
 #ifndef USE_ROCM
-  #include "persistent_topk.cuh"
+  #include "../persistent_topk.cuh"
 #endif
 
 namespace {
 
 #ifndef USE_ROCM
 template <int TopK>
-void launch_persistent_topk(const torch::Tensor& logits,
-                            const torch::Tensor& lengths, torch::Tensor& output,
-                            torch::Tensor& workspace, int64_t max_seq_len) {
+void launch_persistent_topk(const torch::stable::Tensor& logits,
+                            const torch::stable::Tensor& lengths,
+                            torch::stable::Tensor& output,
+                            torch::stable::Tensor& workspace,
+                            int64_t max_seq_len) {
   namespace P = vllm::persistent;
 
   const int64_t num_rows = logits.size(0);
   const int64_t stride = logits.stride(0);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream();
 
   static int num_sms = 0;
   static int max_smem_per_block = 0;
   if (num_sms == 0) {
-    int device;
-    cudaGetDevice(&device);
-    cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device);
-    cudaDeviceGetAttribute(&max_smem_per_block,
-                           cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
+    const cudaDeviceProp* device_prop = get_device_prop();
+    num_sms = device_prop->multiProcessorCount;
+    max_smem_per_block = device_prop->sharedMemPerBlockOptin;
   }
 
   if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
     cudaError_t status =
         vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-            logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-            lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
+            logits.const_data_ptr<float>(), output.mutable_data_ptr<int32_t>(),
+            lengths.const_data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
             static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream);
-    TORCH_CHECK(status == cudaSuccess,
-                "FilteredTopK failed: ", cudaGetErrorString(status));
+    STD_TORCH_CHECK(status == cudaSuccess,
+                    "FilteredTopK failed: ", cudaGetErrorString(status));
   } else {
-    TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
-    TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
+    STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
+    STD_TORCH_CHECK(
+        workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
+        "workspace must be uint8");
 
     int effective_max_smem;
     if (num_rows <= 4) {
@@ -99,9 +101,9 @@ void launch_persistent_topk(const torch::Tensor& logits,
           &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
           smem_size);
     }
-    TORCH_CHECK(occ_err == cudaSuccess,
-                "persistent_topk occupancy query failed: ",
-                cudaGetErrorString(occ_err));
+    STD_TORCH_CHECK(occ_err == cudaSuccess,
+                    "persistent_topk occupancy query failed: ",
+                    cudaGetErrorString(occ_err));
     if (occupancy < 1) occupancy = 1;
 
     // The cooperative spin-wait barrier only runs when at least one row hits
@@ -131,27 +133,29 @@ void launch_persistent_topk(const torch::Tensor& logits,
     // If the cooperative launch wouldn't fit, fall back to FilteredTopK
     // instead of deadlocking. Only relevant when needs_cooperative.
     if (needs_cooperative && total_ctas > hw_resident_cap) {
-      TORCH_CHECK(max_smem_per_block >= 128 * 1024,
-                  "persistent_topk would oversubscribe and the FilteredTopK "
-                  "fallback requires >=128KB smem per block (have ",
-                  max_smem_per_block, "). total_ctas=", total_ctas,
-                  " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
-                  ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
-                  ", smem=", smem_size, ").");
+      STD_TORCH_CHECK(
+          max_smem_per_block >= 128 * 1024,
+          "persistent_topk would oversubscribe and the FilteredTopK "
+          "fallback requires >=128KB smem per block (have ",
+          max_smem_per_block, "). total_ctas=", total_ctas,
+          " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
+          ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
+          ", smem=", smem_size, ").");
       cudaError_t status =
           vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
-              logits.data_ptr<float>(), output.data_ptr<int32_t>(),
-              lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
-              static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
-              stream);
-      TORCH_CHECK(status == cudaSuccess,
-                  "FilteredTopK fallback failed: ", cudaGetErrorString(status));
+              logits.const_data_ptr<float>(),
+              output.mutable_data_ptr<int32_t>(),
+              lengths.const_data_ptr<int32_t>(),
+              static_cast<uint32_t>(num_rows), static_cast<uint32_t>(TopK),
+              static_cast<uint32_t>(stride), stream);
+      STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ",
+                      cudaGetErrorString(status));
       return;
     }
 
     size_t state_bytes = num_groups * sizeof(P::RadixRowState);
-    TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
-                "workspace too small, need ", state_bytes, " bytes");
+    STD_TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
+                    "workspace too small, need ", state_bytes, " bytes");
 
     // Zero the per-group RadixRowState region before launch.
     //
@@ -179,22 +183,22 @@ void launch_persistent_topk(const torch::Tensor& logits,
     //      first red_release. cudaMemsetAsync is stream-ordered: the zero
     //      is globally visible before any CTA runs.
     {
-      cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
-                                           state_bytes, stream);
-      TORCH_CHECK(mz_err == cudaSuccess,
-                  "row_states memset failed: ", cudaGetErrorString(mz_err));
+      cudaError_t mz_err = cudaMemsetAsync(
+          workspace.mutable_data_ptr<uint8_t>(), 0, state_bytes, stream);
+      STD_TORCH_CHECK(mz_err == cudaSuccess,
+                      "row_states memset failed: ", cudaGetErrorString(mz_err));
     }
 
     P::PersistentTopKParams params;
-    params.input = logits.data_ptr<float>();
-    params.output = output.data_ptr<int32_t>();
-    params.lengths = lengths.data_ptr<int32_t>();
+    params.input = logits.const_data_ptr<float>();
+    params.output = output.mutable_data_ptr<int32_t>();
+    params.lengths = lengths.const_data_ptr<int32_t>();
     params.num_rows = static_cast<uint32_t>(num_rows);
     params.stride = static_cast<uint32_t>(stride);
     params.top_k = static_cast<uint32_t>(TopK);
     params.chunk_size = chunk_size;
-    params.row_states =
-        reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
+    params.row_states = reinterpret_cast<P::RadixRowState*>(
+        workspace.mutable_data_ptr<uint8_t>());
     params.ctas_per_group = ctas_per_group;
     params.max_seq_len = static_cast<uint32_t>(max_seq_len);
 
@@ -203,8 +207,8 @@ void launch_persistent_topk(const torch::Tensor& logits,
       auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>;                 \
       cudaError_t err = cudaFuncSetAttribute(                                 \
           kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);    \
-      TORCH_CHECK(err == cudaSuccess,                                         \
-                  "Failed to set smem: ", cudaGetErrorString(err));           \
+      STD_TORCH_CHECK(err == cudaSuccess,                                     \
+                      "Failed to set smem: ", cudaGetErrorString(err));       \
       kernel<<<total_ctas, P::kThreadsPerBlock, smem_size, stream>>>(params); \
     } while (0)
 
@@ -219,37 +223,42 @@ void launch_persistent_topk(const torch::Tensor& logits,
   }
 
   cudaError_t err = cudaGetLastError();
-  TORCH_CHECK(err == cudaSuccess,
-              "persistent_topk failed: ", cudaGetErrorString(err));
+  STD_TORCH_CHECK(err == cudaSuccess,
+                  "persistent_topk failed: ", cudaGetErrorString(err));
 }
 #endif
 
 }  // anonymous namespace
 
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
-                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
+void persistent_topk(const torch::stable::Tensor& logits,
+                     const torch::stable::Tensor& lengths,
+                     torch::stable::Tensor& output,
+                     torch::stable::Tensor& workspace, int64_t k,
                      int64_t max_seq_len) {
 #ifndef USE_ROCM
-  TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
-  TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
-  TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
-  TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
-  TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
-  TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
-  TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
-  TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
-              "lengths must be 1D or 2D");
-  TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
-  TORCH_CHECK(output.dim() == 2, "output must be 2D");
+  STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "Only float32 supported");
+  STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "lengths must be int32");
+  STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "output must be int32");
+  STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+                  "lengths must be 1D or 2D");
+  STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
 
   const int64_t num_rows = logits.size(0);
-  const int64_t stride = logits.stride(0);
 
-  TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
-  TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
-              "output size mismatch");
-  TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
-              "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
+  STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+                  "output size mismatch");
+  STD_TORCH_CHECK(
+      k == 512 || k == 1024 || k == 2048,
+      "persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
 
   if (k == 512) {
     launch_persistent_topk<512>(logits, lengths, output, workspace,
@@ -262,6 +271,6 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
                                  max_seq_len);
   }
 #else
-  TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
+  STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
 #endif
 }

csrc/libtorch_stable/torch_bindings.cpp

@@ -263,6 +263,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
 #endif
 
+  // Merge attn states
+  // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
+  // can be used to combine partial attention results (in the split-KV case)
+  ops.def(
+      "merge_attn_states("
+      "    Tensor! output,"
+      "    Tensor!? output_lse,"
+      "    Tensor prefix_output,"
+      "    Tensor prefix_lse,"
+      "    Tensor suffix_output,"
+      "    Tensor suffix_lse,"
+      "    int!? prefill_tokens_with_context,"
+      "    Tensor? output_scale=None) -> ()");
+
   // Hadamard transforms
   // conditionally compiled so impl registration is in source file
   ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
@@ -319,6 +333,26 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "bool is_neox, Tensor position_ids, "
       "int forced_token_heads_per_warp=-1) -> ()");
 
+  // Apply repetition penalties to logits in-place.
+  ops.def(
+      "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
+      "Tensor output_mask, Tensor repetition_penalties) -> ()");
+
+  // Optimized top-k per row operations.
+  ops.def(
+      "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
+      "Tensor! indices, int numRows, int stride0, "
+      "int stride1, int topK) -> ()");
+
+  ops.def(
+      "top_k_per_row_decode(Tensor logits, int next_n, "
+      "Tensor seq_lens, Tensor! indices, "
+      "int numRows, int stride0, int stride1, int topK) -> ()");
+
+  ops.def(
+      "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
+      "Tensor workspace, int k, int max_seq_len) -> ()");
+
   // Activation ops
   // Activation function used in SwiGLU.
   ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
@@ -422,6 +456,24 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "int type, SymInt row, SymInt tokens) -> Tensor");
 
   ops.def("ggml_moe_get_block_size(int type) -> int");
+
+  // Mamba selective scan kernel
+  ops.def(
+      "selective_scan_fwd(Tensor! u, Tensor! delta,"
+      "Tensor! A, Tensor! B, Tensor! C,"
+      "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
+      "bool delta_softplus,"
+      "Tensor? query_start_loc,"
+      "Tensor? cache_indices,"
+      "Tensor? has_initial_state,"
+      "Tensor! ssm_states,"
+      "int null_block_id,"
+      "int block_size,"
+      "Tensor? block_idx_first_scheduled_token,"
+      "Tensor? block_idx_last_scheduled_token,"
+      "Tensor? initial_state_idx,"
+      "Tensor? cu_chunk_seqlen,"
+      "Tensor? last_chunk_indices) -> ()");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -469,6 +521,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
 #endif
 
+  ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states));
+
   // Layernorm kernels (shared CUDA/ROCm)
   ops.impl("rms_norm", TORCH_BOX(&rms_norm));
   ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm));
@@ -487,6 +541,13 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
   ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
 
+  // Sampler kernels (shared CUDA/ROCm)
+  ops.impl("apply_repetition_penalties_",
+           TORCH_BOX(&apply_repetition_penalties_));
+  ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill));
+  ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode));
+  ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
+
   // Activation kernels (shared CUDA/ROCm)
   ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
   ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
@@ -519,6 +580,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8));
   ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
   ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
+  ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
 }
 
 // These capability-check functions take only primitive args (no tensors), so

csrc/moe/moe_ops.h

@@ -62,6 +62,9 @@ std::tuple<torch::Tensor, torch::Tensor> grouped_topk(
 
 bool moe_permute_unpermute_supported();
 
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t num_experts);
+
 void shuffle_rows(const torch::Tensor& input_tensor,
                   const torch::Tensor& dst2src_map,
                   torch::Tensor& output_tensor);

csrc/moe/moe_permute_unpermute_op.cu

@@ -8,6 +8,108 @@
 // moe_permute kernels require at least CUDA 12.0
 #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)
 
+namespace {
+
+torch::Tensor maybe_allocate_tensor(
+    const std::optional<torch::Tensor>& maybe_tensor,
+    at::IntArrayRef expected_sizes, torch::ScalarType dtype, c10::Device device,
+    char const* name) {
+  auto expected_numel = c10::multiply_integers(expected_sizes);
+  if (maybe_tensor.has_value()) {
+    auto tensor = maybe_tensor.value();
+    TORCH_CHECK(tensor.device() == device, name, " must be on the same device");
+    TORCH_CHECK(tensor.scalar_type() == dtype, name, " has incorrect dtype");
+    TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+    TORCH_CHECK(tensor.numel() >= expected_numel, name,
+                " is too small for the requested shape");
+    auto flat_tensor = tensor.view({tensor.numel()});
+    return flat_tensor.narrow(0, 0, expected_numel).view(expected_sizes);
+  }
+  return torch::empty(expected_sizes, torch::dtype(dtype).device(device));
+}
+
+}  // namespace
+
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t n_expert) {
+  return static_cast<int64_t>(
+      CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert));
+}
+
+void moe_permute_impl(
+    const torch::Tensor& input,                      // [n_token, hidden]
+    const torch::Tensor& topk_ids,                   // [n_token, topk]
+    const torch::Tensor& token_expert_indices,       // [n_token, topk]
+    const std::optional<torch::Tensor>& expert_map,  // [n_expert]
+    int64_t n_expert, int64_t n_local_expert, int64_t topk,
+    torch::Tensor& permuted_input,             // [permuted_size, hidden]
+    torch::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
+    torch::Tensor& inv_permuted_idx,           // [n_token, topk]
+    torch::Tensor& permuted_idx,               // [permute_size]
+    const std::optional<torch::Tensor>& maybe_sort_workspace,
+    const std::optional<torch::Tensor>& maybe_permuted_experts_id,
+    const std::optional<torch::Tensor>& maybe_sorted_row_idx,
+    const std::optional<torch::Tensor>& maybe_topk_ids_for_sort) {
+  TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
+              "expert_first_token_offset must be int64");
+  TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
+              "topk_ids must be int32");
+  TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int,
+              "token_expert_indices must be int32");
+  TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
+              "inv_permuted_idx must be int32");
+  TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
+              "expert_first_token_offset shape != n_local_expert+1");
+  TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
+              "token_expert_indices shape must be same as inv_permuted_idx");
+  auto device = input.device();
+  auto n_token = input.sizes()[0];
+  auto n_hidden = input.sizes()[1];
+  auto expanded_rows = n_token * topk;
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+  auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert);
+  auto sort_workspace =
+      maybe_allocate_tensor(maybe_sort_workspace, {sorter_size}, torch::kInt8,
+                            device, "sort_workspace");
+  auto permuted_experts_id =
+      maybe_allocate_tensor(maybe_permuted_experts_id, topk_ids.sizes(),
+                            at::ScalarType::Int, device, "permuted_experts_id");
+  auto sorted_row_idx =
+      maybe_allocate_tensor(maybe_sorted_row_idx, inv_permuted_idx.sizes(),
+                            at::ScalarType::Int, device, "sorted_row_idx");
+
+  CubKeyValueSorter sorter{};
+  int64_t* valid_num_ptr = nullptr;
+  torch::Tensor topk_ids_for_sort = topk_ids;
+
+  if (expert_map.has_value()) {
+    const int* expert_map_ptr = get_ptr<int>(expert_map.value());
+    valid_num_ptr =
+        get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
+    topk_ids_for_sort =
+        maybe_allocate_tensor(maybe_topk_ids_for_sort, topk_ids.sizes(),
+                              at::ScalarType::Int, device, "topk_ids_for_sort");
+    topk_ids_for_sort.copy_(topk_ids);
+    preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
+                             expert_map_ptr, n_expert, stream);
+  }
+
+  sortAndScanExpert(
+      get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
+      get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
+      get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
+      n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
+
+  MOE_DISPATCH(input.scalar_type(), [&] {
+    expandInputRowsKernelLauncher<scalar_t>(
+        get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
+        get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
+        get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
+        n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
+  });
+}
+
 void moe_permute(
     const torch::Tensor& input,                      // [n_token, hidden]
     const torch::Tensor& topk_ids,                   // [n_token, topk]
@@ -18,65 +120,26 @@ void moe_permute(
     torch::Tensor& expert_first_token_offset,  // [n_local_expert + 1]
     torch::Tensor& inv_permuted_idx,           // [n_token, topk]
     torch::Tensor& permuted_idx) {             // [permute_size]
-  TORCH_CHECK(expert_first_token_offset.scalar_type() == at::ScalarType::Long,
-              "expert_first_token_offset must be int64");
-  TORCH_CHECK(topk_ids.scalar_type() == at::ScalarType::Int,
-              "topk_ids must be int32");
-  TORCH_CHECK(token_expert_indices.scalar_type() == at::ScalarType::Int,
-              "token_expert_indices must be int32");
-  TORCH_CHECK(inv_permuted_idx.scalar_type() == at::ScalarType::Int,
-              "inv_permuted_idx must be int32");
-  TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1,
-              "expert_first_token_offset shape != n_local_expert+1")
-  TORCH_CHECK(inv_permuted_idx.sizes() == token_expert_indices.sizes(),
-              "token_expert_indices shape must be same as inv_permuted_idx");
-  auto n_token = input.sizes()[0];
-  auto n_hidden = input.sizes()[1];
-  auto stream = at::cuda::getCurrentCUDAStream().stream();
-  const long sorter_size =
-      CubKeyValueSorter::getWorkspaceSize(n_token * topk, n_expert);
-  auto sort_workspace = torch::empty(
-      {sorter_size},
-      torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
-  torch::Tensor topk_ids_for_sort = topk_ids;
-  auto permuted_experts_id = torch::empty_like(topk_ids);
-  auto sorted_row_idx = torch::empty_like(inv_permuted_idx);
+  moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk, permuted_input,
+                   expert_first_token_offset, inv_permuted_idx, permuted_idx,
+                   std::nullopt, std::nullopt, std::nullopt, std::nullopt);
+}
 
-  CubKeyValueSorter sorter{};
-  int64_t* valid_num_ptr = nullptr;
-  // pre-process kernel for expert-parallelism:
-  // no local expert id plus "n_expert" offset for priority to local expert
-  // map local expert id [n, .., n+n_local_expert-1] to [0, n_local_expert -1]
-  // For example, 4 expert with ep_size=2. ep_rank=1 owns global expert id
-  // [2,3] with expert_map[-1, -1, 0, 1], preprocess_topk_id  process topk_ids
-  // and map global expert id [2, 3] to local_expert id [0, 1] and map global
-  // expert id [0, 1] ( not in ep rank=1)  to [4, 5] by plus n_expert. This map
-  // operation is to make local expert high priority in following sort topk_ids
-  // and scan local expert_first_token_offset for each ep rank for next group
-  // gemm.
-  if (expert_map.has_value()) {
-    const int* expert_map_ptr = get_ptr<int>(expert_map.value());
-    valid_num_ptr =
-        get_ptr<int64_t>(expert_first_token_offset) + n_local_expert;
-    topk_ids_for_sort = topk_ids.clone();
-    preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk,
-                             expert_map_ptr, n_expert, stream);
-  }
-  // expert sort topk expert id and scan expert id get expert_first_token_offset
-  sortAndScanExpert(
-      get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices),
-      get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
-      get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert,
-      n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream);
-
-  // dispatch expandInputRowsKernelLauncher
-  MOE_DISPATCH(input.scalar_type(), [&] {
-    expandInputRowsKernelLauncher<scalar_t>(
-        get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
-        get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
-        get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
-        n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
-  });
+void moe_permute_with_scratch(
+    const torch::Tensor& input, const torch::Tensor& topk_ids,
+    const torch::Tensor& token_expert_indices,
+    const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
+    int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
+    torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
+    torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
+    torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
+    torch::Tensor& topk_ids_for_sort) {
+  moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert,
+                   n_local_expert, topk, permuted_input,
+                   expert_first_token_offset, inv_permuted_idx, permuted_idx,
+                   sort_workspace, permuted_experts_id, sorted_row_idx,
+                   topk_ids_for_sort);
 }
 
 void moe_unpermute(
@@ -169,6 +232,12 @@ void shuffle_rows(const torch::Tensor& input_tensor,
 
 #else
 
+int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows,
+                                        int64_t n_expert) {
+  TORCH_CHECK(
+      false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0");
+}
+
 void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
                  const torch::Tensor& token_expert_indices,
                  const std::optional<torch::Tensor>& expert_map,
@@ -179,6 +248,19 @@ void moe_permute(const torch::Tensor& input, const torch::Tensor& topk_ids,
   TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0");
 }
 
+void moe_permute_with_scratch(
+    const torch::Tensor& input, const torch::Tensor& topk_ids,
+    const torch::Tensor& token_expert_indices,
+    const std::optional<torch::Tensor>& expert_map, int64_t n_expert,
+    int64_t n_local_expert, int64_t topk, torch::Tensor& permuted_input,
+    torch::Tensor& expert_first_token_offset, torch::Tensor& inv_permuted_idx,
+    torch::Tensor& permuted_idx, torch::Tensor& sort_workspace,
+    torch::Tensor& permuted_experts_id, torch::Tensor& sorted_row_idx,
+    torch::Tensor& topk_ids_for_sort) {
+  TORCH_CHECK(false,
+              "moe_permute_with_scratch is not supported on CUDA < 12.0");
+}
+
 void moe_unpermute(
     const torch::Tensor& permuted_hidden_states,
     const torch::Tensor& topk_weights, const torch::Tensor& inv_permuted_idx,
@@ -199,5 +281,6 @@ bool moe_permute_unpermute_supported() {
 
 TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
   m.impl("moe_permute", &moe_permute);
+  m.impl("moe_permute_with_scratch", &moe_permute_with_scratch);
   m.impl("moe_unpermute", &moe_unpermute);
 }

csrc/moe/torch_bindings.cpp

@@ -100,13 +100,26 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
       "permuted_idx)->()");
 
+  m.def(
+      "moe_permute_with_scratch(Tensor input, Tensor topk_ids,"
+      "Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
+      "int n_local_expert,"
+      "int topk, Tensor! permuted_input, Tensor! "
+      "expert_first_token_offset, Tensor! inv_permuted_idx, Tensor! "
+      "permuted_idx, Tensor! sort_workspace, Tensor! permuted_experts_id, "
+      "Tensor! sorted_row_idx, Tensor! topk_ids_for_sort)->()");
+
   m.def(
       "moe_unpermute(Tensor permuted_hidden_states, Tensor topk_weights,"
       "Tensor inv_permuted_idx, Tensor? expert_first_token_offset, "
       "int topk, Tensor! hidden_states)->()");
 
   m.def("moe_permute_unpermute_supported() -> bool");
+  m.def(
+      "moe_permute_sort_workspace_size(int num_expanded_rows, int n_expert) -> "
+      "int");
   m.impl("moe_permute_unpermute_supported", &moe_permute_unpermute_supported);
+  m.impl("moe_permute_sort_workspace_size", &moe_permute_sort_workspace_size);
 
   // Row shuffle for MoE
   m.def(

csrc/ops.h

@@ -54,13 +54,6 @@ void paged_attention_v2(
     const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
     const int64_t blocksparse_head_sliding_step);
 
-void merge_attn_states(
-    torch::Tensor& output, std::optional<torch::Tensor> output_lse,
-    const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse,
-    const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
-    const std::optional<int64_t> prefill_tokens_with_context,
-    const std::optional<torch::Tensor>& output_scale = std::nullopt);
-
 // rms_norm and fused_add_rms_norm declarations also exist in
 // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
 // because the CPU build still uses these torch::Tensor declarations.
@@ -76,26 +69,6 @@ torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
     torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
     int64_t cache_block_size);
 
-void apply_repetition_penalties_(torch::Tensor& logits,
-                                 const torch::Tensor& prompt_mask,
-                                 const torch::Tensor& output_mask,
-                                 const torch::Tensor& repetition_penalties);
-
-void top_k_per_row_prefill(const torch::Tensor& logits,
-                           const torch::Tensor& rowStarts,
-                           const torch::Tensor& rowEnds, torch::Tensor& indices,
-                           int64_t numRows, int64_t stride0, int64_t stride1,
-                           int64_t topK);
-
-void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n,
-                          const torch::Tensor& seqLens, torch::Tensor& indices,
-                          int64_t numRows, int64_t stride0, int64_t stride1,
-                          int64_t topK);
-
-void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
-                     torch::Tensor& output, torch::Tensor& workspace, int64_t k,
-                     int64_t max_seq_len);
-
 void silu_and_mul_per_block_quant(torch::Tensor& out,
                                   torch::Tensor const& input,
                                   torch::Tensor& scales, int64_t group_size,
@@ -150,22 +123,6 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
                                torch::Tensor& scales,
                                std::optional<torch::Tensor> const& azp);
 
-void selective_scan_fwd(
-    const torch::Tensor& u, const torch::Tensor& delta, const torch::Tensor& A,
-    const torch::Tensor& B, const torch::Tensor& C,
-    const std::optional<torch::Tensor>& D_,
-    const std::optional<torch::Tensor>& z_,
-    const std::optional<torch::Tensor>& delta_bias_, bool delta_softplus,
-    const std::optional<torch::Tensor>& query_start_loc,
-    const std::optional<torch::Tensor>& cache_indices,
-    const std::optional<torch::Tensor>& has_initial_state,
-    const torch::Tensor& ssm_states, int64_t null_block_id, int64_t block_size,
-    const std::optional<torch::Tensor>& block_idx_first_scheduled_token,
-    const std::optional<torch::Tensor>& block_idx_last_scheduled_token,
-    const std::optional<torch::Tensor>& initial_state_idx,
-    const std::optional<torch::Tensor>& cu_chunk_seqlen,
-    const std::optional<torch::Tensor>& last_chunk_indices);
-
 torch::Tensor dynamic_4bit_int_moe_cpu(
     torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
     torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,

csrc/persistent_topk.cuh

@@ -126,10 +126,10 @@ struct RadixRowState {
 // ============================================================================
 
 struct PersistentTopKParams {
-  const float* __restrict__ input;  // [num_rows, stride]
-  int32_t* __restrict__ output;     // [num_rows, top_k]
-  int32_t* __restrict__ lengths;    // [num_rows]
-  RadixRowState* row_states;        // large path: per-group state
+  const float* __restrict__ input;      // [num_rows, stride]
+  int32_t* __restrict__ output;         // [num_rows, top_k]
+  const int32_t* __restrict__ lengths;  // [num_rows]
+  RadixRowState* row_states;            // large path: per-group state
   uint32_t num_rows;
   uint32_t stride;
   uint32_t top_k;           // actual k value for output stride
@@ -1269,9 +1269,11 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
 }
 
 template <typename DType, typename IdType, uint32_t MAX_K = 2048>
-cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
-                                        IdType* lengths, uint32_t num_rows,
-                                        uint32_t top_k_val, uint32_t max_len,
+cudaError_t FilteredTopKRaggedTransform(const DType* input,
+                                        IdType* output_indices,
+                                        const IdType* lengths,
+                                        uint32_t num_rows, uint32_t top_k_val,
+                                        uint32_t max_len,
                                         cudaStream_t stream = 0) {
   constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC;
   constexpr int MAX_VEC = 16 / sizeof(DType);

csrc/quantization/marlin/marlin.cu

@@ -400,10 +400,12 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
                 "Turing only support FP16 or INT8 activation.");
   }
   if (a_type == vllm::kFE4M3fn) {
+    TORCH_CHECK(major_capability * 10 + minor_capability >= 89,
+                "FP8 only support Ada Lovelace or newer GPUs.");
     TORCH_CHECK(
         major_capability * 10 + minor_capability == 89 ||
-            major_capability * 10 + minor_capability == 120,
-        "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
+            major_capability == 12,
+        "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
         "Marlin W4A16 on other devices).");
   }
 

csrc/rocm/skinny_gemms.cu

@@ -1277,6 +1277,12 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
         else
           WVSPLIT_TILE_CFG(64, 16, sYT, 4)
         break;
+      case 5:
+        if (use_wave32)
+          WVSPLIT_TILE_CFG(32, 16, sYT, 5)
+        else
+          WVSPLIT_TILE_CFG(64, 16, sYT, 5)
+        break;
       default:
         throw std::runtime_error(
             "Unsupported N value: " + std::to_string(M_in) + "," +

csrc/torch_bindings.cpp

@@ -62,21 +62,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "    int blocksparse_head_sliding_step) -> ()");
   ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
 
-  // Merge attn states
-  // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
-  // can be used to combine partial attention results (in the split-KV case)
-  ops.def(
-      "merge_attn_states("
-      "    Tensor! output,"
-      "    Tensor!? output_lse,"
-      "    Tensor prefix_output,"
-      "    Tensor prefix_lse,"
-      "    Tensor suffix_output,"
-      "    Tensor suffix_lse,"
-      "    int!? prefill_tokens_with_context,"
-      "    Tensor? output_scale=None) -> ()");
-  ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
-
   // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
@@ -105,31 +90,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
            &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
 
-  // Apply repetition penalties to logits in-place
-  ops.def(
-      "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
-      "Tensor output_mask, Tensor repetition_penalties) -> ()");
-  ops.impl("apply_repetition_penalties_", torch::kCUDA,
-           &apply_repetition_penalties_);
-
-  // Optimized top-k per row operation
-  ops.def(
-      "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, "
-      "Tensor! indices, int numRows, int stride0, "
-      "int stride1, int topK) -> ()");
-  ops.impl("top_k_per_row_prefill", torch::kCUDA, &top_k_per_row_prefill);
-
-  ops.def(
-      "top_k_per_row_decode(Tensor logits, int next_n, "
-      "Tensor seq_lens, Tensor! indices, "
-      "int numRows, int stride0, int stride1, int topK) -> ()");
-  ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode);
-
-  ops.def(
-      "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
-      "Tensor workspace, int k, int max_seq_len) -> ()");
-  ops.impl("persistent_topk", torch::kCUDA, &persistent_topk);
-
   // Quantization ops
 #ifndef USE_ROCM
 
@@ -230,25 +190,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
 #endif
 
-  // Mamba selective scan kernel
-  ops.def(
-      "selective_scan_fwd(Tensor! u, Tensor! delta,"
-      "Tensor! A, Tensor! B, Tensor! C,"
-      "Tensor? D_, Tensor!? z_, Tensor? delta_bias_,"
-      "bool delta_softplus,"
-      "Tensor? query_start_loc,"
-      "Tensor? cache_indices,"
-      "Tensor? has_initial_state,"
-      "Tensor! ssm_states,"
-      "int null_block_id,"
-      "int block_size,"
-      "Tensor? block_idx_first_scheduled_token,"
-      "Tensor? block_idx_last_scheduled_token,"
-      "Tensor? initial_state_idx,"
-      "Tensor? cu_chunk_seqlen,"
-      "Tensor? last_chunk_indices) -> ()");
-  ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
-
 #ifndef USE_ROCM
   ops.def(
       "minimax_allreduce_rms("

docker/Dockerfile

@@ -868,6 +868,7 @@ FROM vllm-base AS test
 ADD . /vllm-workspace/
 
 ARG PYTHON_VERSION
+ARG TARGETPLATFORM
 
 ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
@@ -907,6 +908,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \
             --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
         else \
             echo "Installing dev requirements..." \
+            && if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
+                echo "Recompiling test requirements for arm64..." \
+                && uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \
+                --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
+            fi \
             && uv pip install --system -r requirements/dev.txt \
             --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
         fi \

docker/Dockerfile.rocm

@@ -135,13 +135,18 @@ ENV PATH="/root/.cargo/bin:${PATH}"
 # Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
 # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
 ENV CARGO_BUILD_JOBS=4
+ENV CARGO_NET_RETRY=10
+ENV RUSTUP_MAX_RETRIES=10
 
-# Build the release binary. Cache cargo registry/git, and copy the binary out
-# so it persists into the image layer for later COPY --from=rust-build.
-RUN --mount=type=cache,target=/root/.cargo/registry \
-    --mount=type=cache,target=/root/.cargo/git \
+# Build the release binary. Cargo's registry/git caches can be written by
+# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
+# keeping the cache benefit. Copy the binary out so it persists into the image
+# layer for later COPY --from=rust-build.
+RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
     cd ${COMMON_WORKDIR}/vllm \
-    && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh
+    && VLLM_RS_TARGET_PATH=/tmp/vllm-rs bash build_rust.sh \
+    && test -x /tmp/vllm-rs
 
 # -----------------------
 # vLLM build stages

docker/Dockerfile.rocm_base

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.2-complete
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete
 ARG TRITON_BRANCH="ba5c1517"
 ARG TRITON_REPO="https://github.com/ROCm/triton.git"
 ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17
@@ -104,28 +104,6 @@ ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
 ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
 ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
 
-# torch profiler hotfix for 7.2.2: rebuild CLR with https://github.com/ROCm/rocm-systems/pull/5062
-# will be removed once we move to ROCm 7.2.3
-RUN apt-get update && apt-get install -y rocm-llvm-dev
-RUN pip install CppHeaderParser
-RUN git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems /tmp/rocm-systems \
-    && cd /tmp/rocm-systems \
-    && git sparse-checkout init --cone \
-    && git sparse-checkout set projects/hip projects/clr \
-    && git checkout 35e8c7bf8911862e5389509800e65fdf125412b3 \
-    && export CLR_DIR=/tmp/rocm-systems/projects/clr \
-    && export HIP_DIR=/tmp/rocm-systems/projects/hip \
-    && mkdir -p $CLR_DIR/build && cd $CLR_DIR/build \
-    && cmake \
-        -DHIP_COMMON_DIR=$HIP_DIR \
-        -DCMAKE_PREFIX_PATH="/opt/rocm/" \
-        -DCLR_BUILD_HIP=ON \
-        -DCLR_BUILD_OCL=OFF \
-        -DHIP_PLATFORM=amd \
-        .. \
-    && make -j$(nproc) \
-    && make install \
-    && rm -rf /tmp/rocm-systems
 
 ###
 ### Triton Build
@@ -259,9 +237,8 @@ ARG AITER_REPO
 ARG USE_SCCACHE
 RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
     pip install /install/*.whl
-RUN git clone --recursive ${AITER_REPO}
+RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
 RUN cd aiter \
-    && git checkout ${AITER_BRANCH} \
     && git submodule update --init --recursive \
     && pip install -r requirements.txt
 RUN pip install pyyaml && cd aiter \

docs/benchmarking/cli.md

@@ -918,6 +918,41 @@ vllm bench serve \
 
 </details>
 
+### Replay Timed Traces
+
+<details class="admonition abstract" markdown="1">
+<summary>Show more</summary>
+
+Example of how to run traces which have timing information
+with them.
+
+#### Running MoonshotAI traces
+
+Start the server:
+
+```bash
+vllm serve Qwen/Qwen3.5-2B \
+--host 127.0.0.1 --port 8000
+```
+
+Run the benchmark:
+
+```bash
+# Download an example trace 
+# curl -L -o conversation_trace.jsonl \
+#https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl 
+
+vllm bench serve --model Qwen/Qwen3.5-2B \  
+--dataset-name=timed_trace --num-prompts 100 --host 127.0.0.1 \
+--port 8000 --dataset-path ./conversation_trace.jsonl \
+--ignore-eos  --self-timed --timed-trace-chunk-hash-size 512 \
+--timed-trace-sec-multiplier 0.001 
+```
+
+This will replay the first 100 lines from the trace file `conversation.jsonl`.  
+
+</details>
+
 ### 🧪 Hashing Benchmarks
 
 <details class="admonition abstract" markdown="1">

docs/contributing/profiling.md

@@ -201,43 +201,45 @@ The profiling traces generated by the continuous profiling workflow are publicly
 
 The Python standard library includes
 [cProfile](https://docs.python.org/3/library/profile.html) for profiling Python
-code. vLLM includes a couple of helpers that make it easy to apply it to a section of vLLM.
-Both the `vllm.utils.profiling.cprofile` and `vllm.utils.profiling.cprofile_context` functions can be
-used to profile a section of code.
+code.
 
-!!! note
-    The `vllm.utils.profiling` helpers are deprecated and will be removed in
-    `v0.21`. Please use Python's `cProfile` module directly instead.
+### Example usage - function call
 
-### Example usage - decorator
-
-The first helper is a Python decorator that can be used to profile a function.
-If a filename is specified, the profile will be saved to that file. If no filename is
-specified, profile data will be printed to stdout.
+If a filename is specified, the profile will be saved to that file. If no
+filename is specified, profile data can be printed to stdout.
 
 ```python
-from vllm.utils.profiling import cprofile
+import cProfile
+
 
-@cprofile("expensive_function.prof")
 def expensive_function():
     # some expensive code
     pass
+
+
+profiler = cProfile.Profile()
+profiler.runcall(expensive_function)
+profiler.dump_stats("expensive_function.prof")
 ```
 
-### Example Usage - context manager
-
-The second helper is a context manager that can be used to profile a block of
-code. Similar to the decorator, the filename is optional.
+### Example usage - context manager style
 
 ```python
-from vllm.utils.profiling import cprofile_context
+import cProfile
+
 
 def another_function():
     # more expensive code
     pass
 
-with cprofile_context("another_function.prof"):
+
+profiler = cProfile.Profile()
+profiler.enable()
+try:
     another_function()
+finally:
+    profiler.disable()
+    profiler.dump_stats("another_function.prof")
 ```
 
 ### Analyzing Profile Results

docs/design/attention_backends.md

@@ -205,8 +205,9 @@ hardware and configuration.
 | `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 | `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |
 
-> **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
-> On other GPUs, FlashAttention is used as the default.
+> **‡** Automatic selection tries FlashAttention first. On Blackwell
+> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
+> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.
 
 ### Decode Backends
 

docs/design/fusions.md

@@ -21,7 +21,6 @@ or just on the low or high end.
 | Fusion                                                                         | `PassConfig` flag            | Fused operations                               | Default at                     | E2E Speedup        | Fullgraph | `num_tokens` |
 | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
 | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms)                  | `fuse_allreduce_rms`         | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20%              | No        | Low          |
-| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm)                       | `fuse_minimax_qk_norm`       | Q/K variance all-reduce → Q/K RMSNorm          | Off by default                 | 2-3%               | No        | Low          |
 | [Attention + Quant](#attention--quantization-fuse_attn_quant)                  | `fuse_attn_quant`            | Attention output → FP8/NVFP4 quant             | Off by default                 | 3-7%               | Yes       | Always       |
 | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant)              | `fuse_attn_quant`            | MLA Attention output → FP8/NVFP4 quant         | Off by default                 | TBD                | Yes       | Always       |
 | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache)             | `fuse_rope_kvcache`          | Rotary embedding → KV cache write              | O2 (ROCm/AITER only)           | 2-4%               | No        | Low          |
@@ -42,7 +41,6 @@ The table below lists the quantization schemes supported by each fusion on each
 | Fusion                       | SM100 (Blackwell)                        | SM90 (Hopper)                            | SM89 (Ada)                               | SM80 (Ampere) | ROCm                                     |
 | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
 | `fuse_allreduce_rms`         | FP16/BF16, FP8 static, NVFP4             | FP16/BF16, FP8 static                    | —                                        | —             | —                                        |
-| `fuse_minimax_qk_norm`\*     | FP16/BF16                                | FP16/BF16                                | FP16/BF16                                | FP16/BF16     | —                                        |
 | `fuse_attn_quant`\*          | FP8 static\*, NVFP4\*                    | FP8 static\*                             | FP8 static\*                             | —             | FP8 static\*                             |
 | `fuse_attn_quant` (MLA)\*    | FP8 static\*, FP8 per-group\*, NVFP4\*   | FP8 static\*, FP8 per-group\*            | FP8 static\*, FP8 per-group\*            | —             | FP8 static\* (untested)                  |
 | `fuse_rope_kvcache`          | —                                        | —                                        | —                                        | —             | FP16/BF16                                |
@@ -58,9 +56,6 @@ The table below lists the quantization schemes supported by each fusion on each
 fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
 for per-backend details.
 
-\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
-tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
-
 † `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
 other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
 SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
@@ -191,35 +186,6 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
 
 - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)
 
-### MiniMax QK Norm (`fuse_minimax_qk_norm`)
-
-!!! info
-    This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
-    the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
-    and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
-    optimization level.
-
-**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
-per-token Q/K variances before applying RMS normalization to Q and K.
-
-This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
-`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
-`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
-
-Example:
-
-```bash
-vllm serve MiniMaxAI/MiniMax-M2.5 \
-  --tensor-parallel-size 4 \
-  --compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
-```
-
-**Code locations.**
-
-- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
-- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
-- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
-
 ### Sequence Parallelism (`enable_sp`)
 
 **What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,

docs/features/mooncake_connector_usage.md

@@ -31,7 +31,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_conne
 ### Proxy
 
 ```bash
-python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
+python examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
 ```
 
 Now you can send requests to the proxy server through port 8000.

docs/features/reasoning_outputs.md

@@ -15,6 +15,7 @@ vLLM currently supports the following reasoning models:
 | ------------ | ----------- | ---------------- | ----------- |
 | [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
 | [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
+| [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ |
 | [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
 | [ERNIE-4.5-VL series](https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-PT) | `ernie45` | `json`, `regex` | ❌ |
 | [ERNIE-4.5-21B-A3B-Thinking](https://huggingface.co/baidu/ERNIE-4.5-21B-A3B-Thinking) | `ernie45` | `json`, `regex` | ✅ |
@@ -29,6 +30,7 @@ vLLM currently supports the following reasoning models:
 !!! note
     IBM Granite 3.2 and DeepSeek-V3.1 reasoning is disabled by default; to enable it, you must also pass `thinking=True` in your `chat_template_kwargs`.
     The reasoning feature for the Qwen3 series is enabled by default. To disable it, you must pass `enable_thinking=False` in your `chat_template_kwargs`.
+    Gemma 4 reasoning is disabled by default; to enable it, pass `enable_thinking=True` in your `chat_template_kwargs` or set `reasoning_effort` (which enables it automatically).
     DeepSeek-V3.1 tool calling is supported in non-thinking mode.
     Holo2 reasoning is enabled by default. To disable it, you must also pass `thinking=False` in your `chat_template_kwargs`.
 
@@ -314,9 +316,44 @@ for output in outputs:
     print("text:", output.outputs[0].text)
 ```
 
+## Automatic `enable_thinking` Activation
+
+Some models (such as Gemma 4, DeepSeek-V4-Pro and IBM Granite 3.2) require `enable_thinking: true` in their chat template kwargs to activate thinking mode — without it, reasoning tokens are never generated regardless of other settings.
+
+When you set `reasoning_effort` in a Chat Completions request (or `reasoning.effort` in a Responses API request), vLLM automatically injects `enable_thinking` into the chat template kwargs:
+
+- `reasoning_effort` = `"low"`, `"medium"`, or `"high"` → `enable_thinking = true`
+- `reasoning_effort` = `"none"` → `enable_thinking = false`
+- `reasoning_effort` not set → `enable_thinking` is not injected (preserves existing behavior)
+
+This means you no longer need to manually pass `chat_template_kwargs: {"enable_thinking": true}` when using `reasoning_effort` — it is handled automatically.
+
+!!! note
+    If you explicitly set `enable_thinking` in `chat_template_kwargs`, your value takes priority over the automatic injection. This allows you to override the behavior if needed.
+
+    For models whose templates don't declare `enable_thinking` (e.g., DeepSeek R1), the injected kwarg is harmlessly filtered out by `resolve_chat_template_kwargs`.
+
+### Example
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")
+
+# reasoning_effort automatically enables thinking for models that need it
+response = client.chat.completions.create(
+    model="google/gemma-4-26B-A4B-it",
+    messages=[{"role": "user", "content": "What is 15 * 37?"}],
+    reasoning_effort="high",  # Automatically sets enable_thinking=true
+)
+
+print(response.choices[0].message.reasoning)
+print(response.choices[0].message.content)
+```
+
 ## Limitations
 
-- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`).
+- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`).
 
 ## How to support a new reasoning model
 

docs/getting_started/quickstart.md

@@ -76,6 +76,15 @@ This guide will help you quickly get started with vLLM to perform:
     !!! note
         For more detailed instructions, including Docker, installing from source, and troubleshooting, please refer to the [vLLM on TPU documentation](https://docs.vllm.ai/projects/tpu/en/latest/).
 
+=== "Ascend NPU"
+
+    If you are using Ascend NPUs, you can run vLLM through [vLLM Ascend](https://github.com/vllm-project/vllm-ascend), a community-maintained hardware plugin.
+
+    Follow the installation instructions in the [vLLM Ascend quick start](https://docs.vllm.ai/projects/ascend/en/latest/quick_start.html).
+
+    !!! note
+        Ascend setup depends on your NPU hardware and CANN version. For supported versions, Docker images, and troubleshooting, please refer to the [vLLM Ascend documentation](https://docs.vllm.ai/projects/ascend/en/latest/).
+
 === "Apple Silicon (Mac)"
 
     If you are using Apple Silicon Macs, you can use vLLM-Metal for GPU-accelerated inference via Apple's Metal framework.

docs/mkdocs/hooks/url_schemes.py

@@ -106,8 +106,8 @@ class UrlSchemesPreprocessor(Preprocessor):
             return f"[{gh_icon} {title}]({url})"
 
         markdown = "\n".join(lines)
-        markdown = relative_link.sub(replace_relative_link, markdown)
         markdown = github_link.sub(replace_github_link, markdown)
+        markdown = relative_link.sub(replace_relative_link, markdown)
         return markdown.split("\n")
 
 

docs/models/pooling_models/classify.md

@@ -299,7 +299,3 @@ Example configuration:
 ### Remove softmax from PoolingParams
 
 We have already removed `softmax` and `activation` from PoolingParams. Instead, use `use_activation`, since we allow `classify` and `token_classify` to use any activation function.
-
-### Remove `logit_bias` and `logit_scale`
-
-`logit_bias` and `logit_scale` are deprecated aliases for `logit_mean` and `logit_sigma` respectively. When using `logit_scale`, it is automatically converted to `logit_sigma = 1/logit_scale`. These deprecated parameters will be removed in v0.21.

docs/models/supported_models.md

@@ -428,7 +428,6 @@ th {
 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
 | `IQuestCoderForCausalLM` | IQuestCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Instruct`, etc. | | |
 | `IQuestLoopCoderForCausalLM` | IQuestLoopCoderV1 | `IQuestLab/IQuest-Coder-V1-40B-Loop-Instruct`, etc. | | |
-| `JAISLMHeadModel` | Jais | `inceptionai/jais-13b`, `inceptionai/jais-13b-chat`, `inceptionai/jais-30b-v3`, `inceptionai/jais-30b-chat-v3`, etc. | | ✅︎ |
 | `Jais2ForCausalLM` | Jais2 | `inceptionai/Jais-2-8B-Chat`, `inceptionai/Jais-2-70B-Chat`, etc. | | ✅︎ |
 | `JambaForCausalLM` | Jamba | `ai21labs/AI21-Jamba-1.5-Large`, `ai21labs/AI21-Jamba-1.5-Mini`, `ai21labs/Jamba-v0.1`, etc. | ✅︎ | ✅︎ |
 | `KimiLinearForCausalLM` | Kimi-Linear-48B-A3B-Base, Kimi-Linear-48B-A3B-Instruct | `moonshotai/Kimi-Linear-48B-A3B-Base`, `moonshotai/Kimi-Linear-48B-A3B-Instruct` | | ✅︎ |
@@ -551,6 +550,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ |
 | `CheersForConditionalGeneration` | Cheers | T + I | `ai9stars/Cheers` | | ✅︎ |
 | `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ |
+| `Cosmos3ForConditionalGeneration` | Cosmos3 (understanding tower) | T + I<sup>E+</sup> + V<sup>E+</sup> | `nvidia/Cosmos3-Nano` | | ✅︎ |
 | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
 | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
 | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |

pyproject.toml

@@ -141,6 +141,7 @@ bbc5b7ede = "bbc5b7ede"
 NOOPs = "NOOPs"
 nin_shortcut = "nin_shortcut"
 cudaDevAttrMaxSharedMemoryPerBlockOptin = "cudaDevAttrMaxSharedMemoryPerBlockOptin"
+sharedMemPerBlockOptin = "sharedMemPerBlockOptin"
 
 depthwise_seperable_out_channel = "depthwise_seperable_out_channel"
 pard_token = "pard_token"
@@ -181,6 +182,8 @@ VALU = "VALU"
 # Walsh-Hadamard Transform
 wht = "wht"
 WHT = "WHT"
+# Huawei Compute Architecture for Neural Networks
+CANN = "CANN"
 
 [tool.uv]
 no-build-isolation-package = ["torch"]

requirements/build/rocm.txt

@@ -16,3 +16,4 @@ wheel
 jinja2>=3.1.6
 amdsmi==7.0.2
 timm>=1.0.17
+tilelang==0.1.10

requirements/cuda.txt

@@ -21,7 +21,7 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
 fastsafetensors >= 0.2.2
 
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
-nvidia-cutlass-dsl[cu13]==4.5.0
+nvidia-cutlass-dsl[cu13]==4.5.2
 quack-kernels>=0.3.3
 
 # Tokenspeed_MLA for faster mla with spec decode

requirements/rocm.txt

@@ -22,3 +22,4 @@ timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
+tilelang==0.1.10

requirements/test/cuda.in

@@ -53,12 +53,12 @@ tritonclient>=2.51.0
 grpcio==1.78.0
 grpcio-reflection==1.78.0
 
-arctic-inference == 0.1.1 # Required for suffix decoding test
+arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test
 numba == 0.65.0 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
-fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
-instanttensor>=0.1.5
+fastsafetensors>=0.2.2; platform_machine == "x86_64" # 0.2.2 contains important fixes for multi-GPU mem usage
+instanttensor>=0.1.5; platform_machine == "x86_64"
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0; platform_machine == "x86_64"
 # terratorch is temporarily disabled while PyPI has the `lightning` package

requirements/test/rocm.in

@@ -43,6 +43,7 @@ schemathesis>=3.39.15 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9
+tilelang==0.1.10
 
 genai_perf>=0.0.8
 tritonclient>=2.51.0

requirements/test/rocm.txt

@@ -43,7 +43,9 @@ anyio==4.13.0
     #   starlette
     #   watchfiles
 apache-tvm-ffi==0.1.10
-    # via xgrammar
+    # via
+    #   tilelang
+    #   xgrammar
 arctic-inference==0.1.1
     # via -r requirements/test/rocm.in
 argcomplete==3.6.3
@@ -129,7 +131,9 @@ click==8.3.1
     #   typer
     #   uvicorn
 cloudpickle==3.1.2
-    # via -r requirements/test/../common.txt
+    # via
+    #   -r requirements/test/../common.txt
+    #   tilelang
 colorama==0.4.6
     # via
     #   perceptron
@@ -511,6 +515,8 @@ mistral-common==1.11.2
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   -r requirements/test/rocm.in
+ml-dtypes==0.5.4
+    # via tilelang
 model-hosting-container-standards==0.1.14
     # via
     #   -c requirements/common.txt
@@ -587,6 +593,7 @@ numpy==2.2.6
     #   lm-eval
     #   matplotlib
     #   mistral-common
+    #   ml-dtypes
     #   mteb
     #   numba
     #   opencv-python-headless
@@ -610,6 +617,7 @@ numpy==2.2.6
     #   statsmodels
     #   tensorizer
     #   tifffile
+    #   tilelang
     #   torchvision
     #   transformers
     #   tritonclient
@@ -811,6 +819,7 @@ psutil==7.2.2
     #   accelerate
     #   peft
     #   tensorizer
+    #   tilelang
 py==1.11.0
     # via pytest-forked
 py-cpuinfo==9.0.0
@@ -1192,6 +1201,10 @@ tiktoken==0.12.0
     #   gpt-oss
     #   lm-eval
     #   mistral-common
+tilelang==0.1.10
+    # via
+    #   -c requirements/rocm.txt
+    #   -r requirements/test/rocm.in
 timm==1.0.17
     # via
     #   -c requirements/rocm.txt
@@ -1208,6 +1221,8 @@ tomli==2.4.0
     # via schemathesis
 tomli-w==1.2.0
     # via schemathesis
+torch-c-dlpack-ext==0.1.5
+    # via tilelang
 tqdm==4.67.3
     # via
     #   -r requirements/test/../common.txt
@@ -1225,6 +1240,7 @@ tqdm==4.67.3
     #   pqdm
     #   segmentation-models-pytorch
     #   sentence-transformers
+    #   tilelang
     #   transformers
 transformers==5.5.3
     # via
@@ -1293,6 +1309,7 @@ typing-extensions==4.15.0
     #   sentence-transformers
     #   sqlalchemy
     #   starlette
+    #   tilelang
     #   torch
     #   typeguard
     #   typing-inspection
@@ -1359,6 +1376,8 @@ yarl==1.23.0
     # via
     #   aiohttp
     #   schemathesis
+z3-solver==4.15.4.0
+    # via tilelang
 zipp==3.23.0
     # via importlib-metadata
 

rust/Cargo.lock

@@ -2371,6 +2371,15 @@ version = "0.2.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981"
 
+[[package]]
+name = "libmimalloc-sys"
+version = "0.1.49"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "libredox"
 version = "0.1.14"
@@ -2569,6 +2578,15 @@ version = "0.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b"
 
+[[package]]
+name = "mimalloc"
+version = "0.1.52"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862"
+dependencies = [
+ "libmimalloc-sys",
+]
+
 [[package]]
 name = "mime"
 version = "0.3.17"
@@ -2591,6 +2609,7 @@ version = "2.18.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "328251e58ad8e415be6198888fc207502727dc77945806421ab34f35bf012e7d"
 dependencies = [
+ "indexmap 2.13.0",
  "memo-map",
  "serde",
  "serde_json",
@@ -5608,6 +5627,7 @@ dependencies = [
  "minijinja",
  "minijinja-contrib",
  "openai-harmony",
+ "paste",
  "reqwest",
  "rmp-serde",
  "serde",
@@ -5642,6 +5662,7 @@ dependencies = [
  "educe",
  "expect-test",
  "itertools 0.14.0",
+ "mimalloc",
  "native-tls",
  "serde",
  "serde_json",
@@ -5740,6 +5761,25 @@ dependencies = [
  "prometheus-client",
 ]
 
+[[package]]
+name = "vllm-mock-engine"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "asynk-strim-attr",
+ "clap",
+ "futures",
+ "rand 0.9.2",
+ "rmpv",
+ "serde",
+ "tokio",
+ "tokio-util",
+ "tracing",
+ "tracing-subscriber",
+ "vllm-engine-core-client",
+ "zeromq",
+]
+
 [[package]]
 name = "vllm-reasoning-parser"
 version = "0.1.0"
@@ -5810,6 +5850,7 @@ dependencies = [
  "serde",
  "serde_json",
  "serde_with",
+ "serial_test",
  "tempfile",
  "thiserror 2.0.18",
  "thiserror-ext",
@@ -5847,6 +5888,7 @@ name = "vllm-tool-parser"
 version = "0.1.0"
 dependencies = [
  "criterion",
+ "easy-ext",
  "expect-test",
  "futures",
  "openai-protocol",

rust/Cargo.toml

@@ -6,6 +6,7 @@ members = [
     "src/llm",
     "src/managed-engine",
     "src/metrics",
+    "src/mock-engine",
     "src/reasoning-parser",
     "src/server",
     "src/text",
@@ -45,17 +46,20 @@ http-body = "1.0.1"
 itertools = "0.14.0"
 libc = "0.2.177"
 llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
-minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls"] }
+mimalloc = "0.1.52"
+minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
 minijinja-contrib = { version = "2.0", features = ["pycompat"] }
 native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] }
 ndarray = { version = "0.16.1", features = ["serde"] }
 openai-harmony = "0.0.8"
 openai-protocol = "1.6.0"
 parking_lot = "0.12.5"
+paste = "1.0.15"
 prometheus-client = "0.24.0"
 prometheus-client-derive-encode = "0.5.0"
 prost = "0.14.3"
 prost-types = "0.14.3"
+rand = "0.9.2"
 reasoning-parser = "1.2.2"
 reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
 riptoken = { version = "0.3.0", default-features = false }
@@ -65,11 +69,11 @@ rustc-hash = "1.1.0"
 serde = { version = "1.0.228", features = ["derive"] }
 serde-json-fmt = "0.1.0"
 serde_default = "0.2.0"
-serde_json = "1.0.145"
+serde_json = { version = "1.0.145", features = ["arbitrary_precision", "preserve_order"] }
 serde_repr = "0.1.20"
 serde_tuple = "1.1.3"
 serde_with = "3.18.0"
-serial_test = "3.2.0"
+serial_test = { version = "3.2.0", features = ["file_locks"] }
 socket2 = "0.6.3"
 subenum = "1.1.3"
 task-local = "0.1.1"

rust/src/chat/Cargo.toml

@@ -39,8 +39,9 @@ anyhow.workspace = true
 bytes.workspace = true
 clap.workspace = true
 expect-test.workspace = true
+paste.workspace = true
 rmp-serde.workspace = true
-serial_test = { workspace = true, features = ["file_locks"] }
+serial_test.workspace = true
 tempfile.workspace = true
 tokio.workspace = true
 tracing-subscriber.workspace = true

rust/src/chat/src/event.rs

@@ -53,6 +53,25 @@ impl AssistantContentBlock {
             _ => None,
         }
     }
+
+    /// Return a copy of this block with leading and trailing whitespace trimmed from all text
+    /// fields and tool call arguments, or `None` if the resulting text would be empty.
+    pub fn trim(mut self) -> Option<Self> {
+        match &mut self {
+            Self::Text { text } | Self::Reasoning { text } => {
+                let trimmed_text = text.trim();
+                if trimmed_text.is_empty() {
+                    return None;
+                } else {
+                    *text = trimmed_text.to_string();
+                }
+            }
+            Self::ToolCall(call) => {
+                call.arguments = call.arguments.trim().to_string();
+            }
+        }
+        Some(self)
+    }
 }
 
 #[easy_ext::ext(AssistantMessageExt)]
@@ -119,6 +138,13 @@ impl AssistantMessage {
     pub(crate) fn push_block(&mut self, block: AssistantContentBlock) {
         self.content.push(block);
     }
+
+    /// Return a copy of this message with leading and trailing whitespace trimmed from all text
+    /// fields and tool call arguments, and with any blocks that are empty after trimming removed.
+    pub fn trim(mut self) -> Self {
+        self.content = self.content.into_iter().filter_map(|block| block.trim()).collect();
+        self
+    }
 }
 
 /// Streamed chat event emitted by [`crate::ChatEventStream`].

rust/src/chat/src/lib.rs

@@ -233,7 +233,7 @@ mod tests {
         )
         .unwrap_err();
 
-        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
     }
 
     #[test]

rust/src/chat/src/multimodal.rs

@@ -392,53 +392,13 @@ impl MultimodalModelInfo {
         prompt_token_ids: &mut Vec<u32>,
         replacements: Vec<PromptReplacement>,
     ) -> Result<Vec<PlaceholderRange>> {
-        let mut cursor = 0;
-        let mut ranges = Vec::with_capacity(replacements.len());
-        for replacement in replacements {
-            if replacement.modality != Modality::Image {
-                bail_multimodal!(
-                    "unsupported prompt replacement modality `{}`",
-                    replacement.modality
-                );
-            }
-            let offset = find_next_token(
-                prompt_token_ids,
-                self.spec.placeholder_marker_token_id,
-                cursor,
-            )
-            .ok_or_else(|| {
-                multimodal!(
-                    "placeholder token `{}` was not found in tokenized prompt",
-                    self.spec.placeholder_token
-                )
-            })?;
-
-            if replacement.tokens.is_empty() {
-                bail_multimodal!(
-                    "placeholder token `{}` expanded to no tokens",
-                    self.spec.placeholder_token
-                );
-            }
-            let replacement_len = replacement.tokens.len();
-            let replacement_tokens =
-                replacement.tokens.iter().map(|&token| token as u32).collect::<Vec<_>>();
-            let is_embed = {
-                let mask = replacement_tokens
-                    .iter()
-                    .map(|&token| token == self.spec.placeholder_embed_token_id)
-                    .collect::<Vec<_>>();
-                WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)?
-            };
-
-            prompt_token_ids.splice(offset..offset + 1, replacement_tokens);
-            ranges.push(PlaceholderRange {
-                offset,
-                length: replacement_len,
-                is_embed: Some(is_embed),
-            });
-            cursor = offset + replacement_len;
-        }
-        Ok(ranges)
+        expand_prompt_token_ids(
+            prompt_token_ids,
+            replacements,
+            self.spec.placeholder_marker_token_id,
+            self.spec.placeholder_embed_token_id,
+            &self.spec.placeholder_token,
+        )
     }
 
     /// Convert preprocessed image tensors into engine-core multimodal features.
@@ -516,6 +476,71 @@ impl MultimodalModelInfo {
     }
 }
 
+fn expand_prompt_token_ids(
+    prompt_token_ids: &mut Vec<u32>,
+    replacements: Vec<PromptReplacement>,
+    placeholder_marker_token_id: u32,
+    placeholder_embed_token_id: u32,
+    placeholder_token: &str,
+) -> Result<Vec<PlaceholderRange>> {
+    if replacements.is_empty() {
+        return Ok(Vec::new());
+    }
+
+    let replacement_growth = replacements.iter().fold(0usize, |total, replacement| {
+        total.saturating_add(replacement.tokens.len().saturating_sub(1))
+    });
+    let mut expanded =
+        Vec::with_capacity(prompt_token_ids.len().saturating_add(replacement_growth));
+    let mut ranges = Vec::with_capacity(replacements.len());
+    let mut cursor = 0usize;
+
+    for replacement in replacements {
+        if replacement.modality != Modality::Image {
+            bail_multimodal!(
+                "unsupported prompt replacement modality `{}`",
+                replacement.modality
+            );
+        }
+
+        let offset = find_next_token(prompt_token_ids, placeholder_marker_token_id, cursor)
+            .ok_or_else(|| {
+                multimodal!(
+                    "placeholder token `{placeholder_token}` was not found in tokenized prompt"
+                )
+            })?;
+
+        if replacement.tokens.is_empty() {
+            bail_multimodal!("placeholder token `{placeholder_token}` expanded to no tokens");
+        }
+
+        let replacement_len = replacement.tokens.len();
+        let is_embed = {
+            let mask = replacement
+                .tokens
+                .iter()
+                .map(|&token| token as u32 == placeholder_embed_token_id)
+                .collect::<Vec<_>>();
+            WireTensor::from_bool(vec![replacement_len], mask).map_err(Error::Multimodal)?
+        };
+
+        expanded.extend_from_slice(&prompt_token_ids[cursor..offset]);
+        let expanded_offset = expanded.len();
+        expanded.extend(replacement.tokens.into_iter().map(|token| token as u32));
+        ranges.push(PlaceholderRange {
+            offset: expanded_offset,
+            length: replacement_len,
+            is_embed: Some(is_embed),
+        });
+        cursor = offset + 1;
+    }
+
+    expanded.extend_from_slice(&prompt_token_ids[cursor..]);
+    *prompt_token_ids = expanded;
+
+    Ok(ranges)
+}
+
 /// Find `needle` in `haystack`, starting at `start`.
 ///
 /// This is intentionally order-preserving rather than a global replace: each
@@ -736,6 +761,53 @@ mod tests {
         assert!(matches!(error, Error::Multimodal(message) if message.contains("not found")));
     }
 
+    #[test]
+    fn expand_prompt_tokens_ignores_empty_replacements() {
+        let info = llama4_info();
+        let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
+        let original_prompt_token_ids = prompt_token_ids.clone();
+
+        let ranges = info.expand_prompt_tokens(&mut prompt_token_ids, Vec::new()).unwrap();
+
+        assert!(ranges.is_empty());
+        assert_eq!(prompt_token_ids, original_prompt_token_ids);
+    }
+
+    #[test]
+    fn expand_prompt_tokens_leaves_prompt_unchanged_when_later_placeholder_missing() {
+        let info = llama4_info();
+        let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
+        let original_prompt_token_ids = prompt_token_ids.clone();
+        let replacements = vec![
+            llama4_single_tile_replacement(),
+            llama4_single_tile_replacement(),
+        ];
+
+        let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err();
+
+        assert!(matches!(error, Error::Multimodal(message) if message.contains("not found")));
+        assert_eq!(prompt_token_ids, original_prompt_token_ids);
+    }
+
+    #[test]
+    fn expand_prompt_tokens_errors_when_replacement_is_empty() {
+        let info = llama4_info();
+        let mut prompt_token_ids = vec![1, LLAMA4_IMAGE_ID, 2];
+        let original_prompt_token_ids = prompt_token_ids.clone();
+        let replacements = vec![PromptReplacement::sequence(
+            Modality::Image,
+            "<|image|>",
+            Vec::new(),
+        )];
+
+        let error = info.expand_prompt_tokens(&mut prompt_token_ids, replacements).unwrap_err();
+
+        assert!(
+            matches!(error, Error::Multimodal(message) if message.contains("expanded to no tokens"))
+        );
+        assert_eq!(prompt_token_ids, original_prompt_token_ids);
+    }
+
     #[test]
     fn expand_prompt_tokens_skips_llama4_image_marker_inside_replacement() {
         let info = llama4_info();

rust/src/chat/src/output/default/tool.rs

@@ -15,7 +15,7 @@ use crate::Result;
 use crate::error::Error;
 use crate::event::AssistantBlockKind;
 use crate::output::generate_tool_call_id;
-use crate::parser::tool::{ToolCallDelta, ToolParseResult, ToolParser};
+use crate::parser::tool::{ToolCallDelta, ToolParser, ToolParserOutput};
 
 /// Per-stream tool parsing state.
 struct ToolState {
@@ -57,46 +57,52 @@ impl ToolState {
             return Ok(events);
         }
 
-        let parse_result = self.parser.push(&delta);
+        let mut output = ToolParserOutput::default();
+        let parse_result = self.parser.parse_into(&delta, &mut output);
 
         match parse_result {
-            Ok(result) => self.process_parse_result(kind, result, &mut events)?,
+            Ok(()) => self.process_parser_output(kind, output, &mut events)?,
             Err(error) => {
-                if !self.parser_failed {
-                    warn!(
-                        error = %error.as_report(),
-                        "tool parser failed; falling back to plain text deltas"
-                    );
-                    self.parser_failed = true;
-                }
+                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;
-                events.push(AssistantEvent::TextDelta { kind, delta });
+                push_text_delta(&mut events, kind, self.parser.reset());
             }
         }
 
         Ok(events)
     }
 
-    /// Apply one parsed tool result to the current stream state.
-    fn process_parse_result(
+    /// Apply one parsed tool output to the current stream state.
+    fn process_parser_output(
         &mut self,
         kind: AssistantBlockKind,
-        result: ToolParseResult,
+        output: ToolParserOutput,
         events: &mut Vec<AssistantEvent>,
     ) -> 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, result.normal_text);
-            self.process_tool_items(result.calls, events)?;
+            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(result.calls, events)?;
-            if !result.normal_text.is_empty() {
+            self.process_tool_items(output.calls, events)?;
+            if !output.normal_text.is_empty() {
                 self.open_call_index = None;
-                push_text_delta(events, kind, result.normal_text);
+                push_text_delta(events, kind, output.normal_text);
             }
         }
         Ok(())
@@ -158,8 +164,8 @@ impl ToolState {
         }
 
         match self.parser.finish() {
-            Ok(result) => {
-                self.process_parse_result(AssistantBlockKind::Text, result, &mut events)?
+            Ok(output) => {
+                self.process_parser_output(AssistantBlockKind::Text, output, &mut events)?
             }
             Err(error) => {
                 warn!(
@@ -265,17 +271,24 @@ mod tests {
     use crate::error::Error;
     use crate::event::{AssistantBlockKind, AssistantMessageExt as _};
     use crate::output::structured::structured_chat_event_stream;
-    use crate::parser::tool::{ToolParseResult, ToolParser, ToolParserError};
+    use crate::parser::tool::{
+        DeepSeekV4ToolParser, ToolParser, ToolParserError, ToolParserOutput,
+    };
     use crate::request::ChatTool;
-    use crate::stream::ChatEventStream;
+    use crate::stream::{ChatEventStream, CollectedAssistantMessage};
 
     struct FailingParser {
         fail_next: bool,
+        buffered: String,
     }
 
     struct ScriptedParser {
-        push_results: Vec<ToolParseResult>,
-        finish_result: ToolParseResult,
+        push_outputs: Vec<ToolParserOutput>,
+        finish_output: ToolParserOutput,
+    }
+
+    struct PartialThenFailParser {
+        buffered: String,
     }
 
     impl ToolParser for FailingParser {
@@ -283,10 +296,14 @@ mod tests {
         where
             Self: Sized + 'static,
         {
-            Ok(Box::new(Self { fail_next: false }))
+            Ok(Box::new(Self {
+                fail_next: false,
+                buffered: String::new(),
+            }))
         }
 
-        fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
+        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 {
@@ -294,7 +311,16 @@ mod tests {
                 });
             }
 
-            Ok(ToolParseResult::default())
+            self.buffered.clear();
+            Ok(())
+        }
+
+        fn finish(&mut self) -> Result<ToolParserOutput> {
+            Ok(ToolParserOutput::default())
+        }
+
+        fn reset(&mut self) -> String {
+            std::mem::take(&mut self.buffered)
         }
     }
 
@@ -304,18 +330,215 @@ mod tests {
             Self: Sized + 'static,
         {
             Ok(Box::new(Self {
-                push_results: Vec::new(),
-                finish_result: ToolParseResult::default(),
+                push_outputs: Vec::new(),
+                finish_output: ToolParserOutput::default(),
             }))
         }
 
-        fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
-            Ok(self.push_results.pop().unwrap_or_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<ToolParseResult> {
-            Ok(std::mem::take(&mut self.finish_result))
+        fn finish(&mut self) -> Result<ToolParserOutput> {
+            Ok(std::mem::take(&mut self.finish_output))
         }
+
+        fn reset(&mut self) -> String {
+            String::new()
+        }
+    }
+
+    impl ToolParser for PartialThenFailParser {
+        fn create(_tools: &[ChatTool]) -> vllm_tool_parser::Result<Box<dyn ToolParser>>
+        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<ToolParserOutput> {
+            Ok(ToolParserOutput::default())
+        }
+
+        fn reset(&mut self) -> String {
+            std::mem::take(&mut self.buffered)
+        }
+    }
+
+    fn deepseek_v4_test_tools() -> Vec<ChatTool> {
+        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<String>) -> CollectedAssistantMessage {
+        let events = chunks
+            .into_iter()
+            .map(|delta| {
+                Ok(ContentEvent::TextDelta {
+                    kind: AssistantBlockKind::Text,
+                    delta,
+                })
+            })
+            .chain(std::iter::once(Ok(ContentEvent::Done {
+                prompt_token_count: 1,
+                output_token_count: 1,
+                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);
+
+        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 {
+                prompt_token_count: 1,
+                output_token_count: 1,
+                finish_reason: FinishReason::stop_eos(),
+                kv_transfer_params: None,
+            }),
+        ]);
+
+        let events = tool_event_stream(
+            events,
+            Some(Box::new(PartialThenFailParser {
+                buffered: String::new(),
+            })),
+        )
+        .collect::<Vec<_>>()
+        .await
+        .into_iter()
+        .collect::<crate::Result<Vec<_>>>()
+        .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 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</｜DSML｜parameter>\n",
+            "</｜DSML｜invoke>",
+        );
+        let malformed_second_tool_call = concat!(
+            "\n<｜DSML｜invoke name=\"add\">\n",
+            "not a parameter\n",
+            "</｜DSML｜invoke>\n",
+            "</｜DSML｜tool_calls>",
+        );
+        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]
@@ -341,10 +564,15 @@ mod tests {
             }),
         ]);
 
-        let collected =
-            tool_event_stream(events, Some(Box::new(FailingParser { fail_next: true })))
-                .collect::<Vec<_>>()
-                .await;
+        let collected = tool_event_stream(
+            events,
+            Some(Box::new(FailingParser {
+                fail_next: true,
+                buffered: String::new(),
+            })),
+        )
+        .collect::<Vec<_>>()
+        .await;
 
         let events = collected
             .into_iter()
@@ -415,12 +643,18 @@ mod tests {
                 kv_transfer_params: None,
             }),
         ]);
-        let events = tool_event_stream(events, Some(Box::new(FailingParser { fail_next: false })))
-            .collect::<Vec<_>>()
-            .await
-            .into_iter()
-            .collect::<crate::Result<Vec<_>>>()
-            .unwrap();
+        let events = tool_event_stream(
+            events,
+            Some(Box::new(FailingParser {
+                fail_next: false,
+                buffered: String::new(),
+            })),
+        )
+        .collect::<Vec<_>>()
+        .await
+        .into_iter()
+        .collect::<crate::Result<Vec<_>>>()
+        .unwrap();
 
         assert_eq!(
             events,
@@ -468,7 +702,7 @@ mod tests {
         ]);
 
         let parser = ScriptedParser {
-            push_results: vec![ToolParseResult {
+            push_outputs: vec![ToolParserOutput {
                 normal_text: String::new(),
                 calls: vec![
                     crate::parser::tool::ToolCallDelta {
@@ -483,14 +717,14 @@ mod tests {
                     },
                 ],
             }],
-            finish_result: ToolParseResult::default(),
+            finish_output: ToolParserOutput::default(),
         };
 
         let err = tool_event_stream(events, Some(Box::new(parser)))
             .collect::<Vec<_>>()
             .await
             .into_iter()
-            .find_map(|result| result.err())
+            .find_map(|output| output.err())
             .expect("expected invariant error");
 
         assert!(matches!(err, Error::ToolCallStreamInvariant { .. }));
@@ -514,8 +748,8 @@ mod tests {
         ]);
 
         let parser = ScriptedParser {
-            push_results: vec![
-                ToolParseResult {
+            push_outputs: vec![
+                ToolParserOutput {
                     normal_text: String::new(),
                     calls: vec![crate::parser::tool::ToolCallDelta {
                         tool_index: 0,
@@ -523,11 +757,11 @@ mod tests {
                         arguments: "}".to_string(),
                     }],
                 },
-                ToolParseResult {
+                ToolParserOutput {
                     normal_text: "plain text".to_string(),
                     calls: Vec::new(),
                 },
-                ToolParseResult {
+                ToolParserOutput {
                     normal_text: String::new(),
                     calls: vec![crate::parser::tool::ToolCallDelta {
                         tool_index: 0,
@@ -536,14 +770,14 @@ mod tests {
                     }],
                 },
             ],
-            finish_result: ToolParseResult::default(),
+            finish_output: ToolParserOutput::default(),
         };
 
         let err = tool_event_stream(events, Some(Box::new(parser)))
             .collect::<Vec<_>>()
             .await
             .into_iter()
-            .find_map(|result| result.err())
+            .find_map(|output| output.err())
             .expect("expected invariant error");
 
         assert!(matches!(
@@ -573,7 +807,7 @@ mod tests {
         ]);
 
         let parser = ScriptedParser {
-            push_results: vec![ToolParseResult {
+            push_outputs: vec![ToolParserOutput {
                 normal_text: String::new(),
                 calls: vec![
                     crate::parser::tool::ToolCallDelta {
@@ -588,7 +822,7 @@ mod tests {
                     },
                 ],
             }],
-            finish_result: ToolParseResult::default(),
+            finish_output: ToolParserOutput::default(),
         };
 
         let events = tool_event_stream(events, Some(Box::new(parser)))

rust/src/chat/src/parser/tool/mod.rs

@@ -4,9 +4,10 @@ use std::sync::LazyLock;
 
 pub use vllm_tool_parser::{
     DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
-    Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, KimiK2ToolParser,
-    Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser, Qwen3CoderToolParser,
-    Qwen3XmlToolParser, ToolCallDelta, ToolParseResult, ToolParser, ToolParserError,
+    Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
+    KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser,
+    Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
+    ToolParserOutput,
 };
 
 use crate::parser::ParserFactory;
@@ -22,6 +23,7 @@ pub mod names {
     pub const GLM47: &str = "glm47";
     pub const GEMMA4: &str = "gemma4";
     pub const HERMES: &str = "hermes";
+    pub const HY_V3: &str = "hy_v3";
     pub const KIMI_K2: &str = "kimi_k2";
     pub const LLAMA3_JSON: &str = "llama3_json";
     pub const LLAMA4_JSON: &str = "llama4_json";
@@ -59,6 +61,7 @@ impl ToolParserFactory {
             .register_parser::<Glm47MoeToolParser>(names::GLM47)
             .register_parser::<Gemma4ToolParser>(names::GEMMA4)
             .register_parser::<HermesToolParser>(names::HERMES)
+            .register_parser::<HyV3ToolParser>(names::HY_V3)
             .register_parser::<KimiK2ToolParser>(names::KIMI_K2)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
@@ -75,6 +78,8 @@ impl ToolParserFactory {
             .register_pattern("qwen3.5", names::QWEN3_CODER)
             .register_pattern("qwen", names::QWEN3_XML)
             .register_pattern("hermes", names::HERMES)
+            .register_pattern("hy3", names::HY_V3)
+            .register_pattern("hy_v3", names::HY_V3)
             .register_pattern("llama-4", names::LLAMA4_JSON)
             .register_pattern("llama-3.2", names::LLAMA3_JSON)
             .register_pattern("llama-3.1", names::LLAMA3_JSON)

rust/src/chat/src/parser/tool/tests.rs

@@ -1,6 +1,6 @@
 use vllm_tool_parser::Result;
 
-use super::{ToolParseResult, ToolParser, ToolParserFactory, names};
+use super::{ToolParser, ToolParserFactory, ToolParserOutput, names};
 use crate::Error;
 use crate::request::ChatTool;
 
@@ -18,8 +18,16 @@ impl ToolParser for FakeToolParser {
         true
     }
 
-    fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
-        Ok(ToolParseResult::default())
+    fn parse_into(&mut self, _chunk: &str, _output: &mut ToolParserOutput) -> Result<()> {
+        Ok(())
+    }
+
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        Ok(ToolParserOutput::default())
+    }
+
+    fn reset(&mut self) -> String {
+        String::new()
     }
 }
 
@@ -141,6 +149,10 @@ fn factory_new_resolves_default_patterns() {
         factory.resolve_name_for_model("NousResearch/Hermes-3-Llama-3.1-8B"),
         Some(names::HERMES)
     );
+    assert_eq!(
+        factory.resolve_name_for_model("tencent/Hy3-preview"),
+        Some(names::HY_V3)
+    );
     assert_eq!(
         factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
         Some(names::MINIMAX_M2)

rust/src/chat/src/renderer/hf/tojson.rs

@@ -1,7 +1,7 @@
 use minijinja::value::{Kwargs, ViaDeserialize};
 use minijinja::{Error as MinijinjaError, ErrorKind, Value};
 use serde::Deserialize;
-use serde_json::{self, Value as JsonValue};
+use serde_json::Value as JsonValue;
 use serde_json_fmt::{JsonFormat, JsonSyntaxError};
 use thiserror_ext::AsReport;
 
@@ -13,7 +13,7 @@ use thiserror_ext::AsReport;
 /// - extra kwargs such as `ensure_ascii`, `separators`, and `sort_keys`
 /// - Python-style `indent` handling
 pub(super) fn hf_tojson_filter(
-    value: Value,
+    ViaDeserialize(value): ViaDeserialize<JsonValue>,
     kwargs: Kwargs,
 ) -> std::result::Result<Value, MinijinjaError> {
     let ensure_ascii = kwargs.get::<Option<bool>>("ensure_ascii")?.unwrap_or(false);
@@ -30,18 +30,11 @@ pub(super) fn hf_tojson_filter(
 
     kwargs.assert_all_used()?;
 
-    let json_value: serde_json::Value = serde_json::to_value(&value).map_err(|e| {
-        MinijinjaError::new(
-            ErrorKind::InvalidOperation,
-            format!("Failed to convert to JSON value: {e}"),
-        )
-    })?;
-
     let json_str = {
         let value_to_serialize = if sort_keys {
-            &sort_json_keys(&json_value)
+            &sort_json_keys(&value)
         } else {
-            &json_value
+            &value
         };
 
         build_json_format(indent, separators.0, separators.1, ensure_ascii)?
@@ -214,6 +207,14 @@ mod tests {
         assert_eq!(rendered, "{\"x\":[1,2]}");
     }
 
+    #[test]
+    fn tojson_preserves_arbitrary_precision_number_spelling() {
+        let payload = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
+        let rendered = render("{{ payload|tojson }}", payload);
+
+        assert_eq!(rendered, "{\"x\": 2, \"y\": 1.00}");
+    }
+
     #[test]
     fn tojson_supports_negative_indent_as_newline_only() {
         let rendered = render("{{ payload|tojson(indent=-1) }}", json!([1, 2]));

rust/src/chat/tests/roundtrip.rs

@@ -0,0 +1,541 @@
+//! Text-level roundtrip tests for the real chat-template and output-processor pairing.
+//!
+//! The invariant under test is that a structured assistant message rendered as history can be
+//! parsed from the generated assistant completion and then rendered back to the exact same
+//! assistant-completion text.
+
+use std::pin::Pin;
+use std::sync::Arc;
+
+use anyhow::{Context as _, Result, bail, ensure};
+use futures::{Stream, StreamExt as _, stream};
+use serde_json_fmt::JsonFormat as JsonFmt;
+use serial_test::file_serial;
+use vllm_chat::{
+    AssistantContentBlock, AssistantMessage, AssistantMessageExt as _, AssistantToolCall,
+    ChatEvent, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice, FinishReason,
+    GenerationPromptMode, LoadModelBackendsOptions, NewChatOutputProcessorOptions, ParserSelection,
+    RendererSelection, load_model_backends,
+};
+use vllm_text::{DecodedTextEvent, Finished, Prompt};
+
+/// One model/parser configuration used to run the fixed roundtrip fixtures.
+struct RoundtripCase {
+    /// Hugging Face model id resolved through the production backend loader.
+    model_id: &'static str,
+    /// Final assistant-history suffix rendered by the chat template but not
+    /// generated by the model body consumed by the output processor.
+    // TODO: we should adopt `ContinueFinalAssistant` mode to naturally handle this.
+    assistant_stop_suffix: &'static str,
+    /// Tool parser selection used by the output processor.
+    tool_call_parser: ParserSelection,
+    /// Reasoning parser selection used by the output processor.
+    reasoning_parser: ParserSelection,
+    /// JSON formatting expected after this model's template has materialized
+    /// tool-call arguments.
+    json_fmt: JsonFmt,
+}
+
+impl RoundtripCase {
+    /// Qwen3 XML tool-call format with `qwen3` reasoning tags.
+    fn qwen3() -> Self {
+        Self {
+            model_id: "Qwen/Qwen3-0.6B",
+            assistant_stop_suffix: "<|im_end|>\n",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            json_fmt: spaced_json_fmt(),
+        }
+    }
+
+    /// Qwen3.5 coder-style JSON tool-call format with `qwen3` reasoning tags.
+    fn qwen35() -> Self {
+        Self {
+            model_id: "Qwen/Qwen3.5-4B",
+            assistant_stop_suffix: "<|im_end|>\n",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            json_fmt: compact_json_fmt(),
+        }
+    }
+
+    /// MiniMax M2.5 XML invoke format with `<think>` reasoning tags.
+    fn minimax_m25() -> Self {
+        Self {
+            model_id: "MiniMaxAI/MiniMax-M2.5",
+            assistant_stop_suffix: "[e~[\n",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            json_fmt: compact_json_fmt(),
+        }
+    }
+
+    /// DeepSeek V4 DSML tool-call format.
+    fn deepseek_v4() -> Self {
+        Self {
+            model_id: "deepseek-ai/DeepSeek-V4-Flash",
+            assistant_stop_suffix: "<｜end▁of▁sentence｜>",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            json_fmt: compact_json_fmt(),
+        }
+    }
+
+    /// GLM-4.7 XML-like argument format with `<think>` reasoning tags.
+    fn glm47() -> Self {
+        Self {
+            model_id: "zai-org/GLM-4.7-Flash",
+            assistant_stop_suffix: "",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            json_fmt: compact_json_fmt(),
+        }
+    }
+
+    /// Kimi K2.5 tool-call format with `<think>` reasoning tags.
+    #[allow(dead_code)]
+    fn kimi_k25() -> Self {
+        Self {
+            model_id: "moonshotai/Kimi-K2.5",
+            assistant_stop_suffix: "<|im_end|>",
+            tool_call_parser: ParserSelection::Auto,
+            reasoning_parser: ParserSelection::Auto,
+            json_fmt: spaced_json_fmt(),
+        }
+    }
+}
+
+macro_rules! roundtrip_tests {
+    ($($case:ident => [$($fixture:ident),* $(,)?]),+ $(,)?) => {
+        paste::paste! {
+            $(
+                $(
+                    #[tokio::test]
+                    #[file_serial([<hf_ $case>])]
+                    async fn [<roundtrip_ $case _ $fixture>]() -> Result<()> {
+                        [<run_roundtrip_ $fixture>](RoundtripCase::$case()).await
+                    }
+                )*
+            )+
+        }
+    };
+}
+
+roundtrip_tests! {
+    qwen3 => [reasoning_and_content, tool_call_mix],
+    qwen35 => [reasoning_and_content, tool_call_mix],
+    minimax_m25 => [reasoning_and_content, tool_call_mix],
+    deepseek_v4 => [reasoning_and_content, tool_call_mix],
+    glm47 => [reasoning_and_content, tool_call_mix],
+
+    // Note: Kimi K2.5 strips the reasoning content in history.
+    // TODO: we don't respect model-generated tool call id now so `tool_call_mix` cannot pass.
+    // kimi_k25 => [tool_call_mix],
+}
+
+/// Run the fixed reasoning+content fixture for one model/parser case.
+async fn run_roundtrip_reasoning_and_content(case: RoundtripCase) -> Result<()> {
+    let backends = load_roundtrip_backends(&case).await?;
+    let request = roundtrip_request(
+        "roundtrip-reasoning-content",
+        vec![ChatMessage::text(ChatRole::User, "What is 2 + 2?")],
+        Vec::new(),
+    );
+    let expected_reasoning = "Need compute 2 + 2 directly.";
+    let expected_text = "The answer is 4.";
+
+    let result = run_roundtrip(
+        &case,
+        &backends,
+        &request,
+        AssistantMessage {
+            content: vec![
+                AssistantContentBlock::Reasoning {
+                    text: expected_reasoning.to_string(),
+                },
+                AssistantContentBlock::Text {
+                    text: expected_text.to_string(),
+                },
+            ],
+        },
+    )
+    .await?;
+
+    assert_eq!(
+        result.parsed_message.reasoning().as_deref().map(str::trim),
+        Some(expected_reasoning)
+    );
+    assert_eq!(result.parsed_message.text().trim(), expected_text);
+    assert_eq!(result.parsed_message.tool_calls().count(), 0);
+
+    assert_eq!(
+        result.rerendered_closed_completion,
+        result.closed_completion
+    );
+
+    Ok(())
+}
+
+/// Run the fixed reasoning+multiple-tools fixture for one model/parser case.
+async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
+    let backends = load_roundtrip_backends(&case).await?;
+    let request = roundtrip_request(
+        "roundtrip-reasoning-tools",
+        vec![ChatMessage::text(
+            ChatRole::User,
+            "Check Shanghai weather and add 1.00 plus 2.",
+        )],
+        test_tools(),
+    );
+    let expected_reasoning = "Need call the weather and add tools.";
+    let expected_text = "I will call the tools.";
+
+    let result = run_roundtrip(
+        &case,
+        &backends,
+        &request,
+        AssistantMessage {
+            content: vec![
+                AssistantContentBlock::Reasoning {
+                    text: expected_reasoning.to_string(),
+                },
+                AssistantContentBlock::Text {
+                    text: expected_text.to_string(),
+                },
+                AssistantContentBlock::ToolCall(AssistantToolCall {
+                    id: "functions.get_weather:0".to_string(),
+                    name: "get_weather".to_string(),
+                    arguments: r#"{"location":"Shanghai"}"#.to_string(),
+                }),
+                AssistantContentBlock::ToolCall(AssistantToolCall {
+                    id: "functions.add:1".to_string(),
+                    name: "add".to_string(),
+                    // Intentionally use a non-lexical order of keys and a different number
+                    // formatting style to verify text-level fidelity of the roundtrip.
+                    arguments: r#"{"y":1.00,"x":2}"#.to_string(),
+                }),
+            ],
+        },
+    )
+    .await?;
+
+    assert_eq!(
+        result.parsed_message.reasoning().as_deref().map(str::trim),
+        Some(expected_reasoning)
+    );
+    assert_eq!(result.parsed_message.text().trim(), expected_text);
+
+    let tool_calls = result.parsed_message.tool_calls().collect::<Vec<_>>();
+    assert_eq!(
+        tool_calls.len(),
+        2,
+        "parsed message: {:#?}",
+        result.parsed_message
+    );
+    assert_eq!(tool_calls[0].name, "get_weather");
+    assert_eq!(
+        tool_calls[0].arguments,
+        expected_arguments(&case, r#"{"location": "Shanghai"}"#)?,
+    );
+    assert_eq!(tool_calls[1].name, "add");
+    assert_eq!(
+        tool_calls[1].arguments,
+        expected_arguments(&case, r#"{"y": 1.00, "x": 2}"#)?,
+    );
+
+    assert_eq!(
+        result.rerendered_closed_completion,
+        result.closed_completion
+    );
+
+    Ok(())
+}
+
+/// Compact JSON argument formatting used by JSON-native parsers/renderers.
+fn compact_json_fmt() -> JsonFmt {
+    JsonFmt::new()
+}
+
+/// Python `json.dumps`-style compact formatting with a space after commas and
+/// colons.
+fn spaced_json_fmt() -> JsonFmt {
+    JsonFmt::new()
+        .comma(", ")
+        .expect("literal comma separator is valid JSON")
+        .colon(": ")
+        .expect("literal colon separator is valid JSON")
+}
+
+/// Parse and format expected tool-call arguments from raw JSON text.
+/// Pass in a raw JSON string instead of a structured value to ensure the exact precision and
+/// formatting of numbers are preserved.
+fn expected_arguments(case: &RoundtripCase, raw_json: &str) -> Result<String> {
+    let value: serde_json::Value =
+        serde_json::from_str(raw_json).context("invalid expected tool-call arguments")?;
+
+    case.json_fmt
+        .format_to_string(&value)
+        .context("failed to format expected tool-call arguments")
+}
+
+/// Load the real model chat/text backend for one roundtrip case.
+async fn load_roundtrip_backends(case: &RoundtripCase) -> Result<vllm_chat::LoadedModelBackends> {
+    load_model_backends(
+        case.model_id,
+        LoadModelBackendsOptions {
+            renderer: RendererSelection::Auto,
+            ..Default::default()
+        },
+    )
+    .await
+    .with_context(|| format!("failed to load HF model files for {}", case.model_id))
+}
+
+/// Roundtrip artifacts needed for semantic and exact-text assertions.
+struct RoundtripResult {
+    /// Final assistant message reconstructed by the output processor.
+    parsed_message: AssistantMessage,
+    /// Assistant-completion suffix cut from rendering the expected assistant as
+    /// history.
+    closed_completion: String,
+    /// Assistant-completion suffix cut after rendering the parsed assistant
+    /// back as history.
+    rerendered_closed_completion: String,
+}
+
+/// Render, parse, and rerender one assistant turn through the production
+/// renderer/output-processor boundary.
+async fn run_roundtrip(
+    case: &RoundtripCase,
+    backends: &vllm_chat::LoadedModelBackends,
+    request: &ChatRequest,
+    assistant: AssistantMessage,
+) -> Result<RoundtripResult> {
+    let renderer = backends.chat_backend.chat_renderer();
+    let (prompt, closed_completion_text) =
+        render_closed_completion(renderer.as_ref(), request, &assistant)?;
+    let completion_body = closed_completion_text
+        .strip_suffix(case.assistant_stop_suffix)
+        .with_context(|| {
+            format!(
+                "closed assistant completion did not end with {:?}: {:?}",
+                case.assistant_stop_suffix, closed_completion_text
+            )
+        })?;
+
+    let parsed_message =
+        parse_completion(case, backends, request, &prompt, completion_body).await?;
+    let (_, rerendered_closed_completion) =
+        render_closed_completion(renderer.as_ref(), request, &parsed_message)?;
+
+    Ok(RoundtripResult {
+        parsed_message,
+        closed_completion: closed_completion_text,
+        rerendered_closed_completion,
+    })
+}
+
+/// Render `history` as a production prompt and `history + assistant` as closed
+/// history, then return the production prompt and assistant-completion suffix.
+fn render_closed_completion(
+    renderer: &dyn vllm_chat::ChatRenderer,
+    base_request: &ChatRequest,
+    assistant: &AssistantMessage,
+) -> Result<(String, String)> {
+    let mut prompt_request = base_request.clone();
+    prompt_request.chat_options.generation_prompt_mode = GenerationPromptMode::StartNewAssistant;
+    let prompt = render_text(renderer, &prompt_request).context("failed to render prompt")?;
+
+    let mut full_request = base_request.clone();
+    full_request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+    full_request.messages.push(ChatMessage::from(assistant.clone()));
+    let full = render_text(renderer, &full_request).context("failed to render full prompt")?;
+
+    ensure!(
+        full.starts_with(&prompt),
+        "full prompt must extend production prompt\nprompt: {prompt:?}\nfull: {full:?}"
+    );
+    let completion = full[prompt.len()..].to_string();
+
+    Ok((prompt, completion))
+}
+
+/// Render one chat request and require a text prompt.
+fn render_text(renderer: &dyn vllm_chat::ChatRenderer, request: &ChatRequest) -> Result<String> {
+    match renderer.render(request)?.prompt {
+        Prompt::Text(text) => Ok(text),
+        other => bail!("roundtrip tests expect text prompts, got {other:?}"),
+    }
+}
+
+/// Feed one rendered assistant completion body into the real output processor
+/// and collect its terminal assistant message.
+async fn parse_completion(
+    case: &RoundtripCase,
+    backends: &vllm_chat::LoadedModelBackends,
+    base_request: &ChatRequest,
+    prompt: &str,
+    completion_body: &str,
+) -> Result<AssistantMessage> {
+    let tokenizer = backends.text_backend.tokenizer();
+    let prompt_token_ids = tokenizer
+        .encode(prompt, base_request.add_special_tokens)
+        .context("failed to encode rendered prompt")?;
+
+    let mut request = base_request.clone();
+    let processor = backends.chat_backend.new_chat_output_processor(
+        &mut request,
+        NewChatOutputProcessorOptions {
+            tool_call_parser: &case.tool_call_parser,
+            reasoning_parser: &case.reasoning_parser,
+        },
+    )?;
+
+    let decoded = decoded_completion_stream(prompt_token_ids, completion_body);
+    let mut events = processor.process(decoded)?;
+
+    while let Some(event) = events.next().await {
+        if let ChatEvent::Done { message, .. } = event? {
+            // TODO: currently our parsers are not very strict about preserving or trimming
+            // whitespace, so we trim here to avoid roundtrip failures due to
+            // insignificant whitespace differences. However, this may hurt token-level
+            // fidelity so we should consider improving them.
+            return Ok(message.trim());
+        }
+    }
+
+    bail!("output processor finished without a Done event")
+}
+
+/// Build a decoded-text stream from an already-rendered completion body.
+///
+/// The first event carries real prompt token ids so reasoning parsers can
+/// initialize from the same prompt boundary production uses. Completion text is
+/// split into small chunks to exercise streaming parser state across marker
+/// and JSON boundaries.
+fn decoded_completion_stream(
+    prompt_token_ids: Vec<u32>,
+    completion_body: &str,
+) -> Pin<Box<dyn Stream<Item = vllm_chat::Result<DecodedTextEvent>> + Send>> {
+    let prompt_token_count = prompt_token_ids.len();
+    let mut events = vec![DecodedTextEvent::Start {
+        prompt_token_ids: Arc::from(prompt_token_ids.into_boxed_slice()),
+        prompt_logprobs: None,
+    }];
+
+    let chunks = split_by_chars(completion_body, 7);
+    if chunks.is_empty() {
+        events.push({
+            DecodedTextEvent::TextDelta {
+                delta: String::new(),
+                token_ids: Vec::new(),
+                logprobs: None,
+                finished: Some(Finished {
+                    prompt_token_count: 0,
+                    output_token_count: 0,
+                    finish_reason: FinishReason::stop_eos(),
+                    kv_transfer_params: None,
+                }),
+            }
+        });
+    } else {
+        let last_index = chunks.len() - 1;
+        for (index, chunk) in chunks.into_iter().enumerate() {
+            let finished = (index == last_index).then(|| Finished {
+                prompt_token_count,
+                output_token_count: completion_body.chars().count(),
+                finish_reason: FinishReason::stop_eos(),
+                kv_transfer_params: None,
+            });
+            events.push(DecodedTextEvent::TextDelta {
+                delta: chunk,
+                token_ids: Vec::new(),
+                logprobs: None,
+                finished,
+            });
+        }
+    }
+
+    stream::iter(events).map(Ok).boxed()
+}
+
+/// Split text into chunks containing at most `chunk_chars` Unicode scalar
+/// values.
+fn split_by_chars(text: &str, chunk_chars: usize) -> Vec<String> {
+    let mut chunks = Vec::new();
+    let mut start = 0;
+    let mut count = 0;
+
+    for (index, _) in text.char_indices() {
+        if count == chunk_chars {
+            chunks.push(text[start..index].to_string());
+            start = index;
+            count = 0;
+        }
+        count += 1;
+    }
+
+    if start < text.len() {
+        chunks.push(text[start..].to_string());
+    }
+
+    chunks
+}
+
+/// Build a chat request fixture with parser-enabling tool-choice semantics.
+fn roundtrip_request(
+    request_id: impl Into<String>,
+    messages: Vec<ChatMessage>,
+    tools: Vec<ChatTool>,
+) -> ChatRequest {
+    let mut request = ChatRequest {
+        request_id: request_id.into(),
+        messages,
+        tool_choice: if tools.is_empty() {
+            ChatToolChoice::None
+        } else {
+            ChatToolChoice::Auto
+        },
+        tools,
+        ..ChatRequest::for_test()
+    };
+
+    // Enable thinking for some models so that rendering and parsing the reasoning block is
+    // exercised in the roundtrip.
+    for key in ["thinking", "enable_thinking"] {
+        request.chat_options.template_kwargs.insert(key.to_string(), true.into());
+    }
+
+    request
+}
+
+/// Return the function tools used by the multiple-tool-call fixture.
+fn test_tools() -> Vec<ChatTool> {
+    vec![
+        ChatTool {
+            name: "get_weather".to_string(),
+            description: Some("Get weather for a location".to_string()),
+            parameters: serde_json::json!({
+                "type": "object",
+                "properties": {
+                    "location": { "type": "string" }
+                },
+                "required": ["location"]
+            }),
+            strict: None,
+        },
+        ChatTool {
+            name: "add".to_string(),
+            description: Some("Add two integers".to_string()),
+            parameters: serde_json::json!({
+                "type": "object",
+                "properties": {
+                    "y": { "type": "number" },
+                    "x": { "type": "number" }
+                },
+                "required": ["y", "x"]
+            }),
+            strict: None,
+        },
+    ]
+}

rust/src/cmd/Cargo.toml

@@ -17,6 +17,7 @@ anyhow.workspace = true
 clap.workspace = true
 educe.workspace = true
 itertools.workspace = true
+mimalloc.workspace = true
 native-tls-vendored = { workspace = true, optional = true }
 serde.workspace = true
 serde_json.workspace = true

rust/src/cmd/src/main.rs

@@ -11,6 +11,9 @@ use vllm_managed_engine::ManagedEngineHandle;
 
 use crate::cli::{Cli, Command};
 
+#[global_allocator]
+static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
+
 const TOKIO_WORKER_THREADS_ENV: &str = "TOKIO_WORKER_THREADS";
 const DEFAULT_MAX_TOKIO_WORKER_THREADS: usize = 32;
 

rust/src/engine-core-client/src/client.rs

@@ -595,9 +595,27 @@ impl EngineCoreClient {
     }
 
     /// Return whether the engine is currently sleeping at any level.
+    ///
+    /// Under data parallel, all engines should agree on the sleep state: a
+    /// divergence signals a control-plane bug. Returns
+    /// `Error::InconsistentUtilityResults` if engines disagree.
     pub async fn is_sleeping(&self) -> Result<bool> {
-        // TODO: we only return the result of the first engine here.
-        Ok(self.call_utility("is_sleeping", ()).await?[0])
+        let results: Vec<bool> = self.call_utility("is_sleeping", ()).await?;
+        // `engine_count >= 1` is enforced during startup handshake, so `results`
+        // is normally non-empty; fall back to a fail-loud error rather than
+        // indexing in case that invariant is ever bypassed.
+        let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults {
+            method: "is_sleeping".to_string(),
+            values: "[]".to_string(),
+        })?;
+        if results.iter().all(|&v| v == first) {
+            Ok(first)
+        } else {
+            Err(Error::InconsistentUtilityResults {
+                method: "is_sleeping".to_string(),
+                values: format!("{results:?}"),
+            })
+        }
     }
 
     /// Reset the multi-modal cache.
@@ -613,18 +631,30 @@ impl EngineCoreClient {
     }
 
     /// Reset the prefix cache and optionally the external connector cache.
+    ///
+    /// Under data parallel, returns `true` only when every engine confirms the
+    /// reset (AND aggregation).
     pub async fn reset_prefix_cache(
         &self,
         reset_running_requests: bool,
         reset_connector: bool,
     ) -> Result<bool> {
-        // TODO: we only return the result of the first engine here.
-        Ok(self
+        let results: Vec<bool> = self
             .call_utility(
                 "reset_prefix_cache",
                 (reset_running_requests, reset_connector),
             )
-            .await?[0])
+            .await?;
+        // `engine_count >= 1` is enforced during startup handshake, so `results`
+        // is normally non-empty; fail loud rather than reporting a vacuous
+        // success (`[].all() == true`) in case that invariant is ever bypassed.
+        if results.is_empty() {
+            return Err(Error::InconsistentUtilityResults {
+                method: "reset_prefix_cache".to_string(),
+                values: "[]".to_string(),
+            });
+        }
+        Ok(results.into_iter().all(|ok| ok))
     }
 
     /// Put the engine to sleep.

rust/src/engine-core-client/src/client/imp.rs

@@ -107,13 +107,13 @@ impl ClientInner {
         Ok(registry.abortable_request_ids(request_ids))
     }
 
-    /// Obtain the stream sender for one output. If it indicates the request is
-    /// finished, it will be removed from the registry.
-    pub fn take_sender_for_output(
+    /// Obtain stream senders for a whole engine output batch with one registry
+    /// lock acquisition.
+    pub fn take_senders_for_outputs<'a>(
         &self,
-        output: &EngineCoreOutput,
-    ) -> Option<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>> {
-        self.request_reg.lock().sender_for_output(output)
+        outputs: impl IntoIterator<Item = &'a EngineCoreOutput>,
+    ) -> Vec<Option<mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>>> {
+        self.request_reg.lock().senders_for_outputs(outputs)
     }
 
     /// Remove a batch of requests that have finished or aborted, returning
@@ -301,9 +301,10 @@ pub(crate) async fn run_output_dispatcher_loop(
 
             match outputs.classify() {
                 ClassifiedEngineCoreOutputs::RequestBatch(batch) => {
-                    for output in batch.outputs {
+                    let senders = inner.take_senders_for_outputs(&batch.outputs);
+                    for (output, sender) in batch.outputs.into_iter().zip(senders) {
                         let request_id = output.request_id.clone();
-                        let Some(sender) = inner.take_sender_for_output(&output) else {
+                        let Some(sender) = sender else {
                             debug!(request_id, "dropping output for inactive request");
                             continue;
                         };

rust/src/engine-core-client/src/client/state.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap};
 use std::sync::atomic::{AtomicU64, Ordering};
 
 use tokio::sync::{mpsc, oneshot};
@@ -80,7 +80,7 @@ impl EngineRoutingState {
 #[derive(Debug)]
 pub struct RequestRegistry {
     closed: bool,
-    requests: BTreeMap<String, TrackedRequest>,
+    requests: HashMap<String, TrackedRequest>,
     routing_per_engine: BTreeMap<EngineId, EngineRoutingState>,
 }
 
@@ -88,7 +88,7 @@ impl RequestRegistry {
     pub fn new(engines: &[ConnectedEngine]) -> Self {
         Self {
             closed: false,
-            requests: BTreeMap::default(),
+            requests: HashMap::default(),
             routing_per_engine: engines
                 .iter()
                 .map(|engine| (engine.engine_id.clone(), EngineRoutingState::default()))
@@ -180,6 +180,15 @@ impl RequestRegistry {
         }
     }
 
+    /// Obtain stream senders for a whole engine output batch under one
+    /// registry lock. Finished outputs are removed before returning.
+    pub fn senders_for_outputs<'a>(
+        &mut self,
+        outputs: impl IntoIterator<Item = &'a EngineCoreOutput>,
+    ) -> Vec<Option<OutputSender>> {
+        outputs.into_iter().map(|output| self.sender_for_output(output)).collect()
+    }
+
     /// Remove a batch of requests that have finished or aborted, returning
     /// their stream senders.
     pub fn finish_many<'a>(

rust/src/engine-core-client/src/error.rs

@@ -83,6 +83,8 @@ pub enum Error {
     },
     #[error("utility call `{method}` closed unexpectedly (call_id={call_id})")]
     UtilityCallClosed { method: String, call_id: u64 },
+    #[error("utility call `{method}` returned inconsistent results across engines: {values}")]
+    InconsistentUtilityResults { method: String, values: String },
 
     /// A special variant to allow cloning the same error.
     #[error(transparent)]

rust/src/engine-core-client/src/lib.rs

@@ -2,6 +2,7 @@ mod client;
 mod coordinator;
 mod error;
 mod metrics;
+pub mod mock_engine;
 pub mod protocol;
 #[cfg(any(test, feature = "test-util"))]
 pub mod test_utils;

rust/src/engine-core-client/src/mock_engine.rs

@@ -0,0 +1,265 @@
+use std::path::Path;
+use std::time::Duration;
+
+use tokio::time::timeout;
+use zeromq::prelude::{Socket, SocketRecv, SocketSend};
+use zeromq::util::PeerIdentity;
+use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
+
+use crate::EngineId;
+use crate::error::{Error, Result, bail_unexpected_handshake_message};
+use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
+use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack};
+
+/// Default model length advertised by reusable mock engine helpers.
+pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024;
+/// Default KV block count advertised by reusable mock engine helpers.
+pub const DEFAULT_MOCK_NUM_GPU_BLOCKS: u64 = 0;
+
+/// Startup behavior for one mock engine joining a frontend.
+#[derive(Debug, Clone)]
+pub struct MockEngineConfig {
+    /// Whether the engine should advertise itself as local to the frontend.
+    pub local: bool,
+    /// Whether the engine should advertise itself as headless.
+    pub headless: bool,
+    /// Engine-ready payload reported after INIT, including max model length,
+    /// KV block count, and dtype.
+    pub ready_response: EngineCoreReadyResponse,
+    /// Maximum time to wait for IPC endpoints to appear before connecting.
+    pub connect_timeout: Duration,
+}
+
+impl Default for MockEngineConfig {
+    fn default() -> Self {
+        Self {
+            local: false,
+            headless: true,
+            ready_response: default_ready_response(),
+            connect_timeout: Duration::from_secs(5),
+        }
+    }
+}
+
+/// Construct the ready response used by the standalone mock engine CLI.
+pub fn default_ready_response() -> EngineCoreReadyResponse {
+    EngineCoreReadyResponse {
+        max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN,
+        num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS,
+        dp_stats_address: None,
+        dtype: Some(ModelDtype::Float32),
+    }
+}
+
+/// Coordinator-side sockets used by one mock engine when coordinator mode
+/// is enabled.
+pub struct MockCoordinatorSockets {
+    /// Subscription socket that receives coordinator broadcasts such as
+    /// `START_DP_WAVE`.
+    pub input_sub: SubSocket,
+    /// Push socket used to send coordinator-only `EngineCoreOutputs` back to
+    /// the frontend.
+    pub output_push: PushSocket,
+}
+
+/// One mock engine's connection to one frontend client.
+///
+/// vLLM launches one engine-client pair per API server process. A remote
+/// engine connects to every advertised input/output pair and uses the request's
+/// `client_index` to route outputs back to the originating API server.
+pub struct MockEngineDataSockets {
+    /// Socket used to receive frontend requests.
+    pub dealer: DealerSocket,
+    /// Socket used to publish normal request outputs back to the frontend.
+    pub push: PushSocket,
+}
+
+/// Frontend-facing sockets owned by one mock engine.
+pub struct MockEngineSockets {
+    /// Decoded INIT message sent by the frontend during handshake.
+    pub init: HandshakeInitMessage,
+    /// Data sockets for all frontend clients in client-index order.
+    ///
+    /// For Rust frontend this will always be one socket, while for Python frontend
+    /// this may be multiple sockets if there are multiple API server processes.
+    pub data_sockets: Vec<MockEngineDataSockets>,
+    /// Optional coordinator sockets when the client enabled the in-process
+    /// coordinator.
+    pub coordinator: Option<MockCoordinatorSockets>,
+}
+
+/// Build a HELLO or READY handshake status payload.
+fn ready_message(status: &str, config: &MockEngineConfig) -> ReadyMessage {
+    ReadyMessage {
+        status: Some(status.to_string()),
+        local: Some(config.local),
+        headless: Some(config.headless),
+        parallel_config_hash: None,
+    }
+}
+
+/// Convert an engine id into a ZMQ DEALER identity.
+fn peer_identity(engine_id: impl Into<EngineId>) -> Result<PeerIdentity> {
+    let engine_id = engine_id.into();
+    PeerIdentity::try_from(engine_id.clone()).map_err(|error| Error::UnexpectedHandshakeMessage {
+        message: format!(
+            "invalid mock engine identity {:?}: {error}",
+            engine_id.to_vec()
+        ),
+    })
+}
+
+/// Wait for an IPC endpoint path to appear before attempting to connect.
+async fn wait_for_ipc_endpoint(endpoint: &str, connect_timeout: Duration) -> Result<()> {
+    let Some(socket_path) = endpoint.strip_prefix("ipc://") else {
+        return Ok(());
+    };
+
+    timeout(connect_timeout, async {
+        while !Path::new(socket_path).exists() {
+            tokio::time::sleep(Duration::from_millis(20)).await;
+        }
+    })
+    .await
+    .map_err(|_| Error::HandshakeTimeout {
+        stage: "mock engine IPC endpoint",
+        timeout: connect_timeout,
+    })
+}
+
+/// Encode the engine-ready response sent on input socket registration.
+fn ready_response_payload(config: &MockEngineConfig) -> Result<Vec<u8>> {
+    encode_msgpack(&config.ready_response)
+}
+
+/// Join a frontend-owned handshake endpoint and open mock engine sockets.
+pub async fn connect_to_frontend(
+    engine_handshake: impl AsRef<str>,
+    engine_id: impl Into<EngineId>,
+    config: MockEngineConfig,
+) -> Result<MockEngineSockets> {
+    let engine_handshake = engine_handshake.as_ref();
+    wait_for_ipc_endpoint(engine_handshake, config.connect_timeout).await?;
+
+    let peer_identity = peer_identity(engine_id)?;
+    let mut options = SocketOptions::default();
+    options.peer_identity(peer_identity.clone());
+    let mut handshake = DealerSocket::with_options(options);
+    handshake.connect(engine_handshake).await?;
+    handshake
+        .send(ZmqMessage::from(encode_msgpack(&ready_message(
+            "HELLO", &config,
+        ))?))
+        .await?;
+
+    let init_frames = handshake.recv().await?.into_vec();
+    if init_frames.len() != 1 {
+        bail_unexpected_handshake_message!(
+            "expected one INIT frame from frontend, got {}",
+            init_frames.len()
+        );
+    }
+    let init: HandshakeInitMessage = decode_msgpack(init_frames[0].as_ref())?;
+
+    if init.addresses.inputs.is_empty() {
+        return Err(Error::UnexpectedHandshakeMessage {
+            message: "frontend INIT did not include an input address".to_string(),
+        });
+    }
+    if init.addresses.inputs.len() != init.addresses.outputs.len() {
+        return Err(Error::UnexpectedHandshakeMessage {
+            message: format!(
+                "frontend INIT input/output address count mismatch: {} inputs, {} outputs",
+                init.addresses.inputs.len(),
+                init.addresses.outputs.len()
+            ),
+        });
+    }
+
+    let mut data_sockets = Vec::with_capacity(init.addresses.inputs.len());
+    for (input_address, output_address) in
+        init.addresses.inputs.iter().zip(init.addresses.outputs.iter())
+    {
+        wait_for_ipc_endpoint(input_address, config.connect_timeout).await?;
+        wait_for_ipc_endpoint(output_address, config.connect_timeout).await?;
+
+        let mut input_options = SocketOptions::default();
+        input_options.peer_identity(peer_identity.clone());
+        let mut dealer = DealerSocket::with_options(input_options);
+        dealer.connect(input_address).await?;
+        dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?;
+
+        let mut push = PushSocket::new();
+        push.connect(output_address).await?;
+
+        data_sockets.push(MockEngineDataSockets { dealer, push });
+    }
+
+    let coordinator = match (
+        init.addresses.coordinator_input.as_deref(),
+        init.addresses.coordinator_output.as_deref(),
+    ) {
+        (Some(coordinator_input), Some(coordinator_output)) => {
+            let mut input_sub = SubSocket::new();
+            input_sub.connect(coordinator_input).await?;
+            input_sub.subscribe("").await?;
+
+            let mut output_push = PushSocket::new();
+            output_push.connect(coordinator_output).await?;
+
+            let ready = input_sub.recv().await?.into_vec();
+            if ready.len() != 1 || ready[0].as_ref() != b"READY" {
+                bail_unexpected_handshake_message!(
+                    "expected coordinator READY marker, got {:?}",
+                    ready
+                );
+            }
+
+            Some(MockCoordinatorSockets {
+                input_sub,
+                output_push,
+            })
+        }
+        (None, None) => None,
+        _ => bail_unexpected_handshake_message!(
+            "coordinator handshake addresses must be both present or both absent"
+        ),
+    };
+
+    handshake
+        .send(ZmqMessage::from(encode_msgpack(&ready_message(
+            "READY", &config,
+        ))?))
+        .await?;
+
+    Ok(MockEngineSockets {
+        init,
+        data_sockets,
+        coordinator,
+    })
+}
+
+/// Join already-bootstrapped frontend input/output sockets directly.
+pub async fn connect_to_bootstrapped_frontend(
+    input_address: impl AsRef<str>,
+    output_address: impl AsRef<str>,
+    engine_id: impl Into<EngineId>,
+    config: MockEngineConfig,
+) -> Result<(DealerSocket, PushSocket)> {
+    let input_address = input_address.as_ref();
+    let output_address = output_address.as_ref();
+    wait_for_ipc_endpoint(input_address, config.connect_timeout).await?;
+    wait_for_ipc_endpoint(output_address, config.connect_timeout).await?;
+
+    let peer_identity = peer_identity(engine_id)?;
+    let mut input_options = SocketOptions::default();
+    input_options.peer_identity(peer_identity);
+    let mut dealer = DealerSocket::with_options(input_options);
+    dealer.connect(input_address).await?;
+    dealer.send(ZmqMessage::from(ready_response_payload(&config)?)).await?;
+
+    let mut push = PushSocket::new();
+    push.connect(output_address).await?;
+
+    Ok((dealer, push))
+}

rust/src/engine-core-client/src/protocol/mod.rs

@@ -36,6 +36,14 @@ fn is_false(v: &bool) -> bool {
     !v
 }
 
+fn default_top_p() -> f32 {
+    1.0
+}
+
+fn default_repetition_penalty() -> f32 {
+    1.0
+}
+
 mod classified_outputs;
 pub mod dtype;
 pub mod handshake;
@@ -65,6 +73,24 @@ pub enum EngineCoreRequestType {
 }
 
 impl EngineCoreRequestType {
+    /// Decode the single-byte request type frame used on the engine input
+    /// socket. Returns `None` for unrecognized values.
+    pub fn from_frame(frame: &[u8]) -> Option<Self> {
+        let [value] = frame else {
+            return None;
+        };
+
+        match value {
+            0 => Some(Self::Add),
+            1 => Some(Self::Abort),
+            2 => Some(Self::StartDpWave),
+            3 => Some(Self::Utility),
+            _ => None,
+        }
+    }
+
+    /// Encode the request type as the single-byte frame used on the engine
+    /// input socket.
     pub fn to_frame(self) -> Bytes {
         Bytes::from_static(match self {
             Self::Add => b"\x00",
@@ -200,14 +226,17 @@ pub struct EngineCoreSamplingParams {
     /// greedy sampling.
     pub temperature: f32,
     /// Cumulative probability threshold for nucleus sampling.
+    #[serde(default = "default_top_p")]
     pub top_p: f32,
     /// Maximum number of top tokens to consider. `0` means all tokens.
+    #[serde(default)]
     pub top_k: u32,
     /// Random seed used by the sampler when present.
     pub seed: Option<i64>,
     /// Maximum number of tokens to generate per output sequence.
     pub max_tokens: u32,
     /// Minimum number of tokens to generate before EOS or stop-token handling.
+    #[serde(default)]
     pub min_tokens: u32,
     /// Number of log probabilities to return per generated token.
     ///
@@ -218,12 +247,14 @@ pub struct EngineCoreSamplingParams {
     /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
     pub prompt_logprobs: Option<i32>,
     /// Minimum probability threshold for token sampling.
+    #[serde(default)]
     pub min_p: f32,
     /// Frequency penalty applied by the sampler.
     pub frequency_penalty: f32,
     /// Presence penalty applied by the sampler.
     pub presence_penalty: f32,
     /// Repetition penalty applied by the sampler.
+    #[serde(default = "default_repetition_penalty")]
     pub repetition_penalty: f32,
     /// Token IDs that stop generation.
     pub stop_token_ids: Vec<u32>,

rust/src/engine-core-client/src/protocol/utility.rs

@@ -102,7 +102,7 @@ impl<'de> Deserialize<'de> for UtilityCallId {
 ///
 /// Original Python payload shape:
 /// `(client_index, call_id, method_name, args)`
-#[derive(Debug, Clone, PartialEq, Serialize_tuple)]
+#[derive(Debug, Clone, PartialEq, Serialize_tuple, Deserialize_tuple)]
 pub struct EngineCoreUtilityRequest {
     pub client_index: u32,
     pub call_id: UtilityCallId,

rust/src/engine-core-client/src/test_utils.rs

@@ -1,17 +1,18 @@
 use std::future::Future;
 use std::path::Path;
 use std::pin::Pin;
-use std::time::Duration;
 
 use tempfile::TempDir;
 use tokio::sync::oneshot;
-use zeromq::prelude::{Socket, SocketRecv, SocketSend};
-use zeromq::util::PeerIdentity;
-use zeromq::{DealerSocket, PushSocket, SocketOptions, SubSocket, ZmqMessage};
+use zeromq::{DealerSocket, PushSocket};
 
 use crate::EngineId;
+pub use crate::mock_engine::{MockCoordinatorSockets, MockEngineSockets};
+use crate::mock_engine::{
+    MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend,
+};
 use crate::protocol::ModelDtype;
-use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage, ReadyMessage};
+use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage};
 
 /// Per-test IPC endpoint namespace backed by a unique temporary directory.
 ///
@@ -52,156 +53,29 @@ impl IpcNamespace {
     }
 }
 
-/// Construct a standard local READY message used by mock engines in tests.
-fn ready_message(status: &str) -> ReadyMessage {
-    ReadyMessage {
-        status: Some(status.to_string()),
-        local: Some(true),
-        headless: Some(true),
-        parallel_config_hash: None,
+fn test_mock_engine_config() -> MockEngineConfig {
+    MockEngineConfig {
+        local: true,
+        headless: true,
+        ready_response: EngineCoreReadyResponse {
+            max_model_len: 4096,
+            num_gpu_blocks: 0,
+            dp_stats_address: None,
+            dtype: Some(ModelDtype::Float32),
+        },
+        ..Default::default()
     }
 }
 
-/// Construct a default ready response payload for mock engine input
-/// registration.
-fn ready_response_payload() -> Vec<u8> {
-    rmp_serde::to_vec_named(&EngineCoreReadyResponse {
-        max_model_len: 4096,
-        num_gpu_blocks: 0,
-        dp_stats_address: None,
-        dtype: Some(ModelDtype::Float32),
-    })
-    .expect("encode ready response payload")
-}
-
-/// Coordinator-side sockets connected by one mock engine when coordinator mode
-/// is enabled.
-pub struct MockCoordinatorConnections {
-    /// Subscription socket that receives coordinator broadcasts such as
-    /// `START_DP_WAVE`.
-    pub input_sub: SubSocket,
-    /// Push socket used to send coordinator-only `EngineCoreOutputs` back to
-    /// the frontend.
-    pub output_push: PushSocket,
-}
-
-/// Fully connected mock engine transport state used by tests.
-pub struct MockEngineConnections {
-    /// Decoded INIT message sent by the frontend during handshake.
-    pub init: HandshakeInitMessage,
-    /// Socket used to receive frontend requests.
-    pub dealer: DealerSocket,
-    /// Socket used to publish normal request outputs back to the frontend.
-    pub push: PushSocket,
-    /// Optional coordinator sockets when the client enabled the in-process
-    /// coordinator.
-    pub coordinator: Option<MockCoordinatorConnections>,
-}
-
 /// Complete the engine-core handshake and connect mock input/output sockets
 /// plus optional coordinator sockets.
-pub async fn setup_mock_engine_connections(
+pub async fn setup_mock_engine_sockets(
     engine_handshake: String,
     engine_id: impl Into<EngineId>,
-) -> MockEngineConnections {
-    // Wait for the client to bind the handshake socket before connecting.
-    // A fixed sleep is racy under CI load; instead poll for the socket file.
-    let socket_path = engine_handshake
-        .strip_prefix("ipc://")
-        .expect("handshake address must be ipc://");
-    for _ in 0..100 {
-        if Path::new(socket_path).exists() {
-            break;
-        }
-        tokio::time::sleep(Duration::from_millis(20)).await;
-    }
-
-    let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id");
-
-    let mut options = SocketOptions::default();
-    options.peer_identity(peer_identity.clone());
-    let mut handshake = DealerSocket::with_options(options);
-    handshake
-        .connect(&engine_handshake)
+) -> MockEngineSockets {
+    connect_to_frontend(engine_handshake, engine_id, test_mock_engine_config())
         .await
-        .expect("connect mock engine handshake socket");
-    handshake
-        .send(ZmqMessage::from(
-            rmp_serde::to_vec_named(&ready_message("HELLO")).expect("encode HELLO ready message"),
-        ))
-        .await
-        .expect("send HELLO ready message");
-
-    let init_frames = handshake.recv().await.expect("receive handshake init message").into_vec();
-    assert_eq!(init_frames.len(), 1);
-    let init: HandshakeInitMessage =
-        rmp_serde::from_slice(init_frames[0].as_ref()).expect("decode handshake init message");
-
-    let mut input_options = SocketOptions::default();
-    input_options.peer_identity(peer_identity);
-    let mut dealer = DealerSocket::with_options(input_options);
-    dealer
-        .connect(&init.addresses.inputs[0])
-        .await
-        .expect("connect mock engine input socket");
-    dealer
-        .send(ZmqMessage::from(ready_response_payload()))
-        .await
-        .expect("send mock engine input ready frame");
-
-    let mut push = PushSocket::new();
-    push.connect(&init.addresses.outputs[0])
-        .await
-        .expect("connect mock engine output socket");
-
-    let coordinator = match (
-        init.addresses.coordinator_input.as_deref(),
-        init.addresses.coordinator_output.as_deref(),
-    ) {
-        (Some(coordinator_input), Some(coordinator_output)) => {
-            let mut input_sub = SubSocket::new();
-            input_sub
-                .connect(coordinator_input)
-                .await
-                .expect("connect mock engine coordinator input socket");
-            input_sub
-                .subscribe("")
-                .await
-                .expect("subscribe mock engine coordinator input socket");
-
-            let mut output_push = PushSocket::new();
-            output_push
-                .connect(coordinator_output)
-                .await
-                .expect("connect mock engine coordinator output socket");
-
-            let ready =
-                input_sub.recv().await.expect("receive coordinator READY marker").into_vec();
-            assert_eq!(ready.len(), 1);
-            assert_eq!(ready[0].as_ref(), b"READY");
-
-            Some(MockCoordinatorConnections {
-                input_sub,
-                output_push,
-            })
-        }
-        (None, None) => None,
-        _ => panic!("coordinator handshake addresses must be both present or both absent"),
-    };
-
-    handshake
-        .send(ZmqMessage::from(
-            rmp_serde::to_vec_named(&ready_message("READY")).expect("encode READY ready message"),
-        ))
-        .await
-        .expect("send READY ready message");
-
-    MockEngineConnections {
-        init,
-        dealer,
-        push,
-        coordinator,
-    }
+        .expect("connect mock engine")
 }
 
 /// Connect one mock engine directly to already-bootstrapped frontend
@@ -211,31 +85,14 @@ pub async fn setup_bootstrapped_mock_engine(
     output_address: String,
     engine_id: impl Into<EngineId>,
 ) -> (DealerSocket, PushSocket) {
-    for endpoint in [&input_address, &output_address] {
-        if let Some(socket_path) = endpoint.strip_prefix("ipc://") {
-            for _ in 0..100 {
-                if Path::new(socket_path).exists() {
-                    break;
-                }
-                tokio::time::sleep(Duration::from_millis(20)).await;
-            }
-        }
-    }
-
-    let peer_identity = PeerIdentity::try_from(engine_id.into()).expect("peer id");
-    let mut input_options = SocketOptions::default();
-    input_options.peer_identity(peer_identity);
-    let mut dealer = DealerSocket::with_options(input_options);
-    dealer.connect(&input_address).await.expect("connect mock engine input socket");
-    dealer
-        .send(ZmqMessage::from(ready_response_payload()))
-        .await
-        .expect("send mock engine input ready frame");
-
-    let mut push = PushSocket::new();
-    push.connect(&output_address).await.expect("connect mock engine output socket");
-
-    (dealer, push)
+    connect_to_bootstrapped_frontend(
+        input_address,
+        output_address,
+        engine_id,
+        test_mock_engine_config(),
+    )
+    .await
+    .expect("connect bootstrapped mock engine")
 }
 
 /// Complete the engine-core handshake and connect mock input/output sockets.
@@ -247,9 +104,11 @@ pub async fn setup_mock_engine_with_init(
     engine_handshake: String,
     engine_id: impl Into<EngineId>,
 ) -> (HandshakeInitMessage, DealerSocket, PushSocket) {
-    let MockEngineConnections {
-        init, dealer, push, ..
-    } = setup_mock_engine_connections(engine_handshake, engine_id).await;
+    let MockEngineSockets {
+        init, data_sockets, ..
+    } = setup_mock_engine_sockets(engine_handshake, engine_id).await;
+    let MockEngineDataSockets { dealer, push } =
+        data_sockets.into_iter().next().expect("mock engine data socket");
     (init, dealer, push)
 }
 

rust/src/engine-core-client/src/tests/client.rs

@@ -30,7 +30,7 @@ use crate::protocol::{
     EngineCoreRequestType, EngineCoreSamplingParams, decode_engine_core_outputs,
 };
 use crate::test_utils::{
-    IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_connections,
+    IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets,
     setup_mock_engine_with_init, spawn_mock_engine_task,
 };
 use crate::{
@@ -477,8 +477,8 @@ async fn coordinator_handshake_includes_engine_control_addresses() {
     let (init_tx, init_rx) = oneshot::channel();
     let (shutdown_tx, shutdown_rx) = oneshot::channel();
     let engine_task = tokio::spawn(async move {
-        let connections = setup_mock_engine_connections(handshake_address, &engine_id).await;
-        let _ = init_tx.send(connections.init.clone());
+        let sockets = setup_mock_engine_sockets(handshake_address, &engine_id).await;
+        let _ = init_tx.send(sockets.init.clone());
         let _ = shutdown_rx.await;
     });
 
@@ -515,14 +515,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
     let engine0_task = tokio::spawn({
         let handshake_address = handshake_address.clone();
         async move {
-            let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
+            let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
             let mut coordinator =
                 engine.coordinator.take().expect("coordinator sockets should be present");
+            let data_socket = engine.data_sockets.first_mut().expect("data socket");
 
             let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
             assert_eq!((wave, exclude_engine), (0, 0));
 
-            let add = recv_engine_message(&mut engine.dealer).await;
+            let add = recv_engine_message(&mut data_socket.dealer).await;
             assert_eq!(add[0].as_ref(), &[0x00]);
             let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
             assert_eq!(request.request_id, "req-1");
@@ -538,7 +539,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
             );
 
             send_outputs(
-                &mut engine.push,
+                &mut data_socket.push,
                 EngineCoreOutputs {
                     engine_index: 0,
                     outputs: vec![request_output(
@@ -565,14 +566,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
             let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
             assert_eq!((wave, exclude_engine), (1, 0));
 
-            let add = recv_engine_message(&mut engine.dealer).await;
+            let add = recv_engine_message(&mut data_socket.dealer).await;
             assert_eq!(add[0].as_ref(), &[0x00]);
             let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
             assert_eq!(request.request_id, "req-3");
             assert_eq!(request.current_wave, 1);
 
             send_outputs(
-                &mut engine.push,
+                &mut data_socket.push,
                 EngineCoreOutputs {
                     engine_index: 0,
                     outputs: vec![request_output(
@@ -594,14 +595,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
     let engine1_task = tokio::spawn({
         let handshake_address = handshake_address.clone();
         async move {
-            let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await;
+            let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await;
             let mut coordinator =
                 engine.coordinator.take().expect("coordinator sockets should be present");
+            let data_socket = engine.data_sockets.first_mut().expect("data socket");
 
             let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
             assert_eq!((wave, exclude_engine), (0, 0));
 
-            let add = recv_engine_message(&mut engine.dealer).await;
+            let add = recv_engine_message(&mut data_socket.dealer).await;
             assert_eq!(add[0].as_ref(), &[0x00]);
             let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
             assert_eq!(request.request_id, "req-2");
@@ -617,7 +619,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
             );
 
             send_outputs(
-                &mut engine.push,
+                &mut data_socket.push,
                 EngineCoreOutputs {
                     engine_index: 1,
                     outputs: vec![request_output(
@@ -637,7 +639,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() {
             assert!(
                 timeout(
                     Duration::from_millis(200),
-                    recv_engine_message(&mut engine.dealer)
+                    recv_engine_message(&mut data_socket.dealer)
                 )
                 .await
                 .is_err()
@@ -712,7 +714,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
     let engine0_task = tokio::spawn({
         let handshake_address = handshake_address.clone();
         async move {
-            let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
+            let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
             let mut coordinator =
                 engine.coordinator.take().expect("coordinator sockets should be present");
 
@@ -727,7 +729,7 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() {
     let engine1_task = tokio::spawn({
         let handshake_address = handshake_address.clone();
         async move {
-            let mut engine = setup_mock_engine_connections(handshake_address, &[0x01, 0x00]).await;
+            let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await;
             let mut coordinator =
                 engine.coordinator.take().expect("coordinator sockets should be present");
 
@@ -778,9 +780,10 @@ async fn coordinator_accepts_stats_only_outputs() {
 
     let (shutdown_tx, shutdown_rx) = oneshot::channel();
     let engine_task = tokio::spawn(async move {
-        let mut engine = setup_mock_engine_connections(handshake_address, &[0x00, 0x00]).await;
+        let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await;
         let mut coordinator =
             engine.coordinator.take().expect("coordinator sockets should be present");
+        let data_socket = engine.data_sockets.first_mut().expect("data socket");
 
         let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await;
         assert_eq!((wave, exclude_engine), (0, 0));
@@ -799,13 +802,13 @@ async fn coordinator_accepts_stats_only_outputs() {
         )
         .await;
 
-        let add = recv_engine_message(&mut engine.dealer).await;
+        let add = recv_engine_message(&mut data_socket.dealer).await;
         assert_eq!(add[0].as_ref(), &[0x00]);
         let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
         assert_eq!(request.request_id, "req-stats");
 
         send_outputs(
-            &mut engine.push,
+            &mut data_socket.push,
             EngineCoreOutputs {
                 engine_index: 0,
                 outputs: vec![request_output(
@@ -2116,6 +2119,226 @@ async fn collective_rpc_flattens_results_from_all_engines() {
     client.shutdown().await.unwrap();
 }
 
+/// Spawn a mock engine that handles a single utility call, asserts the method
+/// name and serialized args match, and replies with `result`.
+fn spawn_mock_utility_engine(
+    handshake_address: String,
+    engine_id: Vec<u8>,
+    expected_method: &'static str,
+    expected_args: Value,
+    result: bool,
+) -> (
+    tokio::sync::oneshot::Sender<()>,
+    tokio::task::JoinHandle<()>,
+) {
+    spawn_mock_engine_task(handshake_address, engine_id, move |dealer, push| {
+        Box::pin(async move {
+            let utility = recv_engine_message(dealer).await;
+            assert_eq!(utility[0].as_ref(), &[0x03]);
+            let payload = decode_value(&utility[1]);
+            let array = match payload {
+                Value::Array(array) => array,
+                other => panic!("expected utility payload array, got {other:?}"),
+            };
+            // Utility requests serialize as `(client_index, call_id, method, args)`.
+            let call_id = array[1].as_u64().expect("call_id");
+            assert_eq!(array[2], Value::from(expected_method));
+            assert_eq!(array[3], expected_args, "unexpected utility args");
+            send_outputs(
+                push,
+                EngineCoreOutputs {
+                    utility_output: Some(UtilityOutput {
+                        call_id: call_id.into(),
+                        failure_message: None,
+                        result: Some(utility_result_value(result)),
+                    }),
+                    ..Default::default()
+                },
+            )
+            .await;
+        })
+    })
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn is_sleeping_returns_error_when_engines_disagree() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+
+    let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
+        handshake_address.clone(),
+        b"engine-0".to_vec(),
+        "is_sleeping",
+        Value::Array(vec![]),
+        true,
+    );
+    let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
+        handshake_address.clone(),
+        b"engine-1".to_vec(),
+        "is_sleeping",
+        Value::Array(vec![]),
+        false,
+    );
+
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            2,
+            "test-model",
+            Duration::from_secs(2),
+            5,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+
+    let error = client.is_sleeping().await.unwrap_err();
+    assert!(
+        matches!(
+            &error,
+            Error::InconsistentUtilityResults { method, .. } if method == "is_sleeping"
+        ),
+        "expected InconsistentUtilityResults, got {error:?}",
+    );
+
+    let _ = shutdown_tx_0.send(());
+    let _ = shutdown_tx_1.send(());
+    engine_task_0.await.unwrap();
+    engine_task_1.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn is_sleeping_returns_value_when_all_engines_agree() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+
+    let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
+        handshake_address.clone(),
+        b"engine-0".to_vec(),
+        "is_sleeping",
+        Value::Array(vec![]),
+        true,
+    );
+    let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
+        handshake_address.clone(),
+        b"engine-1".to_vec(),
+        "is_sleeping",
+        Value::Array(vec![]),
+        true,
+    );
+
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            2,
+            "test-model",
+            Duration::from_secs(2),
+            5,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+
+    assert!(client.is_sleeping().await.unwrap());
+
+    let _ = shutdown_tx_0.send(());
+    let _ = shutdown_tx_1.send(());
+    engine_task_0.await.unwrap();
+    engine_task_1.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn reset_prefix_cache_returns_true_when_all_engines_succeed() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+
+    let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
+        handshake_address.clone(),
+        b"engine-0".to_vec(),
+        "reset_prefix_cache",
+        Value::Array(vec![Value::from(false), Value::from(false)]),
+        true,
+    );
+    let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
+        handshake_address.clone(),
+        b"engine-1".to_vec(),
+        "reset_prefix_cache",
+        Value::Array(vec![Value::from(false), Value::from(false)]),
+        true,
+    );
+
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            2,
+            "test-model",
+            Duration::from_secs(2),
+            5,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+
+    assert!(client.reset_prefix_cache(false, false).await.unwrap());
+
+    let _ = shutdown_tx_0.send(());
+    let _ = shutdown_tx_1.send(());
+    engine_task_0.await.unwrap();
+    engine_task_1.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn reset_prefix_cache_returns_false_when_any_engine_fails() {
+    init_tracing();
+    let ipc = IpcNamespace::new().unwrap();
+    let handshake_address = ipc.handshake_endpoint();
+
+    let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine(
+        handshake_address.clone(),
+        b"engine-0".to_vec(),
+        "reset_prefix_cache",
+        Value::Array(vec![Value::from(false), Value::from(false)]),
+        true,
+    );
+    let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine(
+        handshake_address.clone(),
+        b"engine-1".to_vec(),
+        "reset_prefix_cache",
+        Value::Array(vec![Value::from(false), Value::from(false)]),
+        false,
+    );
+
+    let client = connect_client_with_ipc(
+        handshake_test_config(
+            handshake_address,
+            2,
+            "test-model",
+            Duration::from_secs(2),
+            5,
+            None,
+        ),
+        &ipc,
+    )
+    .await;
+
+    assert!(!client.reset_prefix_cache(false, false).await.unwrap());
+
+    let _ = shutdown_tx_0.send(());
+    let _ = shutdown_tx_1.send(());
+    engine_task_0.await.unwrap();
+    engine_task_1.await.unwrap();
+    client.shutdown().await.unwrap();
+}
+
 #[test]
 fn python_msgpack_fixtures_match_rust_encoding() {
     init_tracing();

rust/src/mock-engine/Cargo.toml

@@ -0,0 +1,30 @@
+[package]
+name = "vllm-mock-engine"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[[bin]]
+name = "vllm-mock-engine"
+path = "src/main.rs"
+
+[dependencies]
+anyhow.workspace = true
+asynk-strim-attr.workspace = true
+clap.workspace = true
+futures.workspace = true
+rand.workspace = true
+rmpv.workspace = true
+serde.workspace = true
+tokio = { workspace = true, features = ["signal"] }
+tokio-util.workspace = true
+tracing.workspace = true
+tracing-subscriber.workspace = true
+vllm-engine-core-client.workspace = true
+zeromq.workspace = true
+
+[dev-dependencies]
+vllm-engine-core-client = { workspace = true, features = ["test-util"] }
+
+[lints]
+workspace = true

rust/src/mock-engine/README.md

@@ -0,0 +1,104 @@
+# vLLM Mock Engine
+
+`vllm-mock-engine` is a small engine-side process for frontend stress testing. It
+joins a frontend-owned startup handshake, reports a large ready response, treats
+prefill as instant, and emits random decode tokens until each request reaches
+its `max_tokens`.
+
+The frontend must own the handshake socket. Start the frontend first, then start
+the mock engine with the same handshake address.
+
+## Start the mock engine
+
+```bash
+cargo run -p vllm-mock-engine -- \
+  --handshake-address tcp://127.0.0.1:29550 \
+  --engine-count 1 \
+  --output-token-chunk-size 1 \
+  --vocab-size 32000 \
+  --seed 0 \
+  --log-requests
+```
+
+Useful knobs:
+
+- `--engine-count` must match the frontend's expected data-parallel engine
+  count.
+- `--output-token-chunk-size` controls how many token IDs appear in one
+  `EngineCoreOutput`; values greater than 1 are useful for MTP/spec-decode
+  shaped frontend tests.
+- `--vocab-size` should stay within the tokenizer vocabulary of the model used
+  by the frontend.
+
+Stop it with Ctrl-C.
+
+## Rust Frontend
+
+Terminal 1:
+
+```bash
+cargo run --bin vllm-rs -- serve Qwen/Qwen3-0.6B \
+  --data-parallel-size 1 \
+  --data-parallel-size-local 0 \
+  --handshake-port 29550
+```
+
+Terminal 2:
+
+```bash
+cargo run -p vllm-mock-engine -- \
+  --handshake-address tcp://127.0.0.1:29550
+```
+
+For multiple mock engines, set both sides to the same count:
+
+```bash
+cargo run --bin vllm-rs -- serve Qwen/Qwen3-0.6B \
+  --data-parallel-size 4 \
+  --data-parallel-size-local 0 \
+  --handshake-port 29550
+
+cargo run -p vllm-mock-engine -- \
+  --handshake-address tcp://127.0.0.1:29550 \
+  --engine-count 4
+```
+
+## Python Frontend
+
+Use `vllm serve` with `--data-parallel-size-local 0` so the Python process runs
+as a frontend/API server and waits for external engines on
+`--data-parallel-rpc-port`.
+
+Terminal 1:
+
+```bash
+vllm serve Qwen/Qwen3-0.6B \
+  --data-parallel-address 127.0.0.1 \
+  --data-parallel-rpc-port 29550 \
+  --data-parallel-size 1 \
+  --data-parallel-size-local 0
+```
+
+Terminal 2:
+
+```bash
+cargo run -p vllm-mock-engine -- \
+  --handshake-address tcp://127.0.0.1:29550
+```
+
+## Smoke Request
+
+After either frontend is ready:
+
+```bash
+curl http://127.0.0.1:8000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-0.6B",
+    "messages": [{"role": "user", "content": "hello"}],
+    "max_tokens": 16,
+    "stream": true
+  }'
+```
+
+Always pass `max_tokens`; the mock engine stops by length.

rust/src/mock-engine/src/engine.rs

@@ -0,0 +1,416 @@
+use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::hash::{Hash as _, Hasher as _};
+use std::time::{SystemTime, UNIX_EPOCH};
+
+use anyhow::{Result, anyhow};
+use rand::rngs::StdRng;
+use rand::{Rng as _, SeedableRng as _};
+use rmpv::Value;
+use serde::Serialize;
+use tokio::sync::mpsc;
+use tokio::task::yield_now;
+use tokio_util::sync::CancellationToken;
+use tracing::{debug, info, warn};
+use vllm_engine_core_client::protocol::utility::{
+    EngineCoreUtilityRequest, UtilityOutput, UtilityResultEnvelope,
+};
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, EngineCoreRequest,
+};
+
+use super::Opt;
+
+/// Derive a stable per-request seed from the CLI seed, engine, and request id.
+fn request_seed(base_seed: u64, engine_index: u32, request_id: &str) -> u64 {
+    let mut hasher = std::hash::DefaultHasher::new();
+    base_seed.hash(&mut hasher);
+    engine_index.hash(&mut hasher);
+    request_id.hash(&mut hasher);
+    hasher.finish()
+}
+
+/// Current UNIX timestamp in seconds for engine-core output envelopes.
+fn now_secs() -> f64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|duration| duration.as_secs_f64())
+        .unwrap_or_default()
+}
+
+/// Build one request output with only token IDs and terminal status populated.
+fn request_output(
+    request_id: String,
+    new_token_ids: Vec<u32>,
+    finish_reason: Option<EngineCoreFinishReason>,
+) -> EngineCoreOutput {
+    EngineCoreOutput {
+        request_id,
+        new_token_ids,
+        finish_reason,
+        ..Default::default()
+    }
+}
+
+/// Produce an empty output with a terminal finish reason for an invalid request.
+fn empty_finish_outputs(
+    engine_index: u32,
+    request_id: String,
+    finish_reason: EngineCoreFinishReason,
+) -> EngineCoreOutputs {
+    let output = request_output(request_id, Vec::new(), Some(finish_reason));
+    let finished_requests = BTreeSet::from([output.request_id.clone()]);
+
+    EngineCoreOutputs {
+        engine_index,
+        outputs: vec![output],
+        timestamp: now_secs(),
+        finished_requests: Some(finished_requests),
+        ..Default::default()
+    }
+}
+
+/// Encode a utility result into the protocol's msgpack value envelope.
+fn utility_envelope<T>(value: T) -> Result<UtilityResultEnvelope>
+where
+    T: Serialize,
+{
+    Ok(UtilityResultEnvelope::without_type_info(
+        rmpv::ext::to_value(value)?,
+    ))
+}
+
+/// Produce the minimal utility responses needed by the Rust frontend.
+fn utility_response(
+    engine_index: u32,
+    request: EngineCoreUtilityRequest,
+) -> Result<EngineCoreOutputs> {
+    let result = match request.method_name.as_str() {
+        "get_supported_tasks" => utility_envelope(vec!["generate"]),
+        "is_sleeping" => utility_envelope(false),
+        "reset_prefix_cache" => utility_envelope(true),
+        "reset_mm_cache"
+        | "reset_encoder_cache"
+        | "profile"
+        | "sleep"
+        | "wake_up"
+        | "execute_dummy_batch" => utility_envelope(()),
+        _ => utility_envelope(Value::Nil),
+    }?;
+
+    Ok(EngineCoreOutputs {
+        engine_index,
+        utility_output: Some(UtilityOutput {
+            call_id: request.call_id,
+            failure_message: None,
+            result: Some(result),
+        }),
+        timestamp: now_secs(),
+        ..Default::default()
+    })
+}
+
+/// Message sent from the frontend to the mock engine task to drive the engine loop.
+pub(crate) enum EngineInput {
+    Request(Box<EngineCoreRequest>),
+    Abort(Vec<String>),
+    Utility(EngineCoreUtilityRequest),
+    StartDpWave,
+}
+
+/// Message sent from the mock engine task to the frontend for one engine output batch.
+pub(crate) struct EngineOutput {
+    pub client_index: u32,
+    pub outputs: EngineCoreOutputs,
+}
+
+/// Per-request decode state owned by one mock engine.
+#[derive(Debug)]
+struct ActiveRequest {
+    request_id: String,
+    client_index: u32,
+    prompt_len: usize,
+    max_tokens: usize,
+    generated: usize,
+    rng: StdRng,
+}
+
+impl ActiveRequest {
+    /// Create a new active request from an incoming EngineCoreRequest, or return an immediate
+    /// finish reason if the request is invalid.
+    fn new(
+        engine_index: u32,
+        request: Box<EngineCoreRequest>,
+        opt: &Opt,
+    ) -> Result<Self, EngineCoreFinishReason> {
+        let request_id = request.request_id;
+        let client_index = request.client_index;
+        let prompt_len = request.prompt_token_ids.as_ref().map(Vec::len).unwrap_or_default();
+
+        let Some(sampling_params) = request.sampling_params else {
+            warn!(
+                request_id,
+                "request has no sampling params; returning engine error"
+            );
+            return Err(EngineCoreFinishReason::Error);
+        };
+        let max_tokens = sampling_params.max_tokens as usize;
+
+        if opt.log_requests {
+            info!(
+                request_id,
+                prompt_len,
+                max_tokens,
+                chunk_size = opt.output_token_chunk_size,
+                "mock request started"
+            );
+        }
+
+        if max_tokens == 0 {
+            return Err(EngineCoreFinishReason::Length);
+        }
+
+        Ok(ActiveRequest {
+            rng: StdRng::seed_from_u64(request_seed(opt.seed, engine_index, &request_id)),
+            request_id,
+            client_index,
+            prompt_len,
+            max_tokens,
+            generated: 0,
+        })
+    }
+
+    /// Advance this request by one mock engine step.
+    fn step(&mut self, opt: &Opt) -> EngineCoreOutput {
+        let remaining = self.max_tokens - self.generated;
+        let chunk_len = remaining.min(opt.output_token_chunk_size);
+        let mut new_token_ids = Vec::with_capacity(chunk_len);
+        for _ in 0..chunk_len {
+            new_token_ids.push(self.rng.random_range(0..opt.vocab_size));
+        }
+        self.generated += chunk_len;
+
+        let finished = self.generated >= self.max_tokens;
+        request_output(
+            self.request_id.clone(),
+            new_token_ids,
+            finished.then_some(EngineCoreFinishReason::Length),
+        )
+    }
+}
+
+/// Internal state for one mock engine instance, owned by the engine loop task.
+struct Engine {
+    engine_index: u32,
+    opt: Opt,
+    active_requests: HashMap<String, ActiveRequest>,
+}
+
+impl Engine {
+    /// Drain one frontend request message received on the input DEALER socket.
+    fn handle_input(&mut self, input: EngineInput) -> Result<Vec<EngineOutput>> {
+        let mut outputs = Vec::new();
+
+        match input {
+            EngineInput::Request(request) => {
+                let request_id = request.request_id.clone();
+                let client_index = request.client_index;
+
+                if self.active_requests.contains_key(&request_id) {
+                    warn!(
+                        engine_index = self.engine_index,
+                        request_id, "duplicate mock request id"
+                    );
+                    return Ok(vec![EngineOutput {
+                        client_index,
+                        outputs: empty_finish_outputs(
+                            self.engine_index,
+                            request_id,
+                            EngineCoreFinishReason::Error,
+                        ),
+                    }]);
+                }
+
+                match ActiveRequest::new(self.engine_index, request, &self.opt) {
+                    Ok(request) => {
+                        self.active_requests.insert(request_id, request);
+                    }
+                    Err(finish_reason) => {
+                        return Ok(vec![EngineOutput {
+                            client_index,
+                            outputs: empty_finish_outputs(
+                                self.engine_index,
+                                request_id,
+                                finish_reason,
+                            ),
+                        }]);
+                    }
+                }
+            }
+
+            EngineInput::Abort(request_ids) => {
+                let mut outputs_by_client =
+                    BTreeMap::<u32, (Vec<EngineCoreOutput>, BTreeSet<String>)>::new();
+                for request_id in request_ids {
+                    if let Some(request) = self.active_requests.remove(&request_id) {
+                        let output = request_output(
+                            request_id.clone(),
+                            Vec::new(),
+                            Some(EngineCoreFinishReason::Abort),
+                        );
+                        let (outputs, finished_requests) = outputs_by_client
+                            .entry(request.client_index)
+                            .or_insert_with(|| (Vec::new(), BTreeSet::new()));
+                        outputs.push(output);
+                        finished_requests.insert(request_id.clone());
+                        if self.opt.log_requests {
+                            info!(request_id, finish_reason = "abort", "mock request aborted");
+                        }
+                    }
+                }
+                for (client_index, (client_outputs, finished_requests)) in outputs_by_client {
+                    outputs.push({
+                        let outputs = EngineCoreOutputs {
+                            engine_index: self.engine_index,
+                            outputs: client_outputs,
+                            timestamp: now_secs(),
+                            finished_requests: Some(finished_requests),
+                            ..Default::default()
+                        };
+                        EngineOutput {
+                            client_index,
+                            outputs,
+                        }
+                    });
+                }
+            }
+
+            EngineInput::Utility(request) => {
+                debug!(
+                    engine_index = self.engine_index,
+                    call_id = %request.call_id,
+                    method = request.method_name,
+                    "mock utility request"
+                );
+                let client_index = request.client_index;
+                outputs.push({
+                    let outputs = utility_response(self.engine_index, request)?;
+                    EngineOutput {
+                        client_index,
+                        outputs,
+                    }
+                });
+            }
+
+            EngineInput::StartDpWave => {
+                debug!(
+                    engine_index = self.engine_index,
+                    "ignoring START_DP_WAVE in mock engine"
+                );
+            }
+        }
+
+        Ok(outputs)
+    }
+
+    /// Advance active requests once and return one batched engine output.
+    fn step(&mut self) -> Vec<EngineOutput> {
+        if self.active_requests.is_empty() {
+            return Vec::new();
+        }
+
+        let mut outputs_by_client =
+            BTreeMap::<u32, (Vec<EngineCoreOutput>, BTreeSet<String>)>::new();
+        let mut all_finished_requests = BTreeSet::new();
+
+        for request in self.active_requests.values_mut() {
+            let client_index = request.client_index;
+            let output = request.step(&self.opt);
+            let request_id = request.request_id.clone();
+            let finished = output.finished();
+            if output.finished() {
+                all_finished_requests.insert(request_id.clone());
+                if self.opt.log_requests {
+                    info!(
+                        request_id,
+                        prompt_len = request.prompt_len,
+                        output_tokens = request.generated,
+                        finish_reason = "length",
+                        "mock request finished"
+                    );
+                }
+            }
+            let (outputs, finished_requests) = outputs_by_client
+                .entry(client_index)
+                .or_insert_with(|| (Vec::new(), BTreeSet::new()));
+            if finished {
+                finished_requests.insert(request_id.clone());
+            }
+            outputs.push(output);
+        }
+
+        for request_id in &all_finished_requests {
+            self.active_requests.remove(request_id);
+        }
+
+        outputs_by_client
+            .into_iter()
+            .filter_map(|(client_index, (outputs, finished_requests))| {
+                (!outputs.is_empty()).then(|| EngineOutput {
+                    client_index,
+                    outputs: EngineCoreOutputs {
+                        engine_index: self.engine_index,
+                        outputs,
+                        timestamp: now_secs(),
+                        finished_requests: (!finished_requests.is_empty())
+                            .then_some(finished_requests),
+                        ..Default::default()
+                    },
+                })
+            })
+            .collect()
+    }
+}
+
+/// Run the main loop for the mock engine, receiving `EngineInput` from `input_rx`
+/// and sending `EngineOutput` to `output_tx` until `shutdown` is cancelled.
+pub(crate) async fn run_engine_loop(
+    engine_index: u32,
+    opt: Opt,
+    mut input_rx: mpsc::UnboundedReceiver<EngineInput>,
+    output_tx: mpsc::Sender<EngineOutput>,
+    shutdown: CancellationToken,
+) -> Result<()> {
+    let mut engine = Engine {
+        engine_index,
+        opt,
+        active_requests: HashMap::new(),
+    };
+
+    loop {
+        let outputs = tokio::select! {
+            biased;
+            _ = shutdown.cancelled() => break,
+
+            input = input_rx.recv() => {
+                let input = input
+                    .ok_or_else(|| anyhow!("mock engine input channel closed"))?;
+                engine.handle_input(input)?
+            }
+
+            // If there are active requests, step them once after yielding to the scheduler to
+            // avoid blocking the engine loop while still making steady progress on request outputs.
+            _ = yield_now(), if !engine.active_requests.is_empty() => {
+                engine.step()
+            }
+        };
+
+        for output in outputs {
+            output_tx
+                .send(output)
+                .await
+                .map_err(|_| anyhow!("mock engine IO task shut down"))?;
+        }
+    }
+
+    Ok(())
+}

rust/src/mock-engine/src/io.rs

@@ -0,0 +1,121 @@
+use anyhow::{Context as _, Result, anyhow, bail};
+use futures::{Stream, StreamExt as _, stream};
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+use tracing::warn;
+use vllm_engine_core_client::mock_engine::MockEngineDataSockets;
+use vllm_engine_core_client::protocol::utility::EngineCoreUtilityRequest;
+use vllm_engine_core_client::protocol::{
+    EngineCoreRequest, EngineCoreRequestType, decode_msgpack, encode_msgpack,
+};
+use zeromq::{DealerSocket, PushSocket, SocketRecv as _, SocketSend as _, ZmqMessage};
+
+use crate::engine::{EngineInput, EngineOutput};
+
+/// Send one engine output batch to the client over the appropriate push socket.
+async fn send_engine_outputs_to_client(
+    push_sockets: &mut [PushSocket],
+    EngineOutput {
+        client_index,
+        outputs,
+    }: EngineOutput,
+) -> Result<()> {
+    let message = ZmqMessage::from(encode_msgpack(&outputs)?);
+    push_sockets[client_index as usize].send(message).await?;
+    Ok(())
+}
+
+/// Create a stream of `EngineInput` by continuously receiving messages from the given dealer socket
+/// and decoding them into `EngineInput`.
+fn dealer_input_stream(dealer: DealerSocket) -> impl Stream<Item = Result<EngineInput>> {
+    stream::unfold(dealer, |mut dealer| async {
+        let input = loop {
+            let message =
+                match dealer.recv().await.context("failed to receive message from dealer socket") {
+                    Ok(message) => message,
+                    Err(err) => break Err(err),
+                };
+
+            match decode_request(message) {
+                Ok(input) => break Ok(input),
+                Err(err) => {
+                    warn!(%err, "failed to decode engine request message; ignoring");
+                }
+            }
+        };
+
+        Some((input, dealer))
+    })
+}
+
+/// Decode a `ZmqMessage` into an `EngineInput`. Returns an error if the message is malformed or
+/// contains an unknown/unsupported request type.
+fn decode_request(message: ZmqMessage) -> Result<EngineInput> {
+    let frames = message.into_vec();
+    if frames.is_empty() {
+        bail!("empty engine request message");
+    }
+    if frames.len() != 2 {
+        bail!("invalid frame count for engine request: {}", frames.len());
+    }
+
+    let request_type_frame = frames[0].as_ref();
+    let Some(request_type) = EngineCoreRequestType::from_frame(request_type_frame) else {
+        bail!("unknown engine request type: {:?}", request_type_frame);
+    };
+
+    let input = match request_type {
+        EngineCoreRequestType::Add => {
+            let request: Box<EngineCoreRequest> = decode_msgpack(frames[1].as_ref())?;
+            EngineInput::Request(request)
+        }
+        EngineCoreRequestType::Abort => {
+            let request_ids: Vec<String> = decode_msgpack(frames[1].as_ref())?;
+            EngineInput::Abort(request_ids)
+        }
+        EngineCoreRequestType::Utility => {
+            let request: EngineCoreUtilityRequest = decode_msgpack(frames[1].as_ref())?;
+            EngineInput::Utility(request)
+        }
+        EngineCoreRequestType::StartDpWave => EngineInput::StartDpWave,
+    };
+
+    Ok(input)
+}
+
+/// Run the main IO loop for the mock engine, continuously receiving and decoding raw messages from
+/// the dealer sockets, sending them to the engine loop task via `input_tx`, and receiving
+/// `EngineOutput` from the engine loop task via `output_rx` and sending them to the client over the
+/// appropriate push socket, until `shutdown` is cancelled.
+pub(crate) async fn run_io_loop(
+    data_sockets: Vec<MockEngineDataSockets>,
+    input_tx: mpsc::UnboundedSender<EngineInput>,
+    mut output_rx: mpsc::Receiver<EngineOutput>,
+    shutdown: CancellationToken,
+) -> Result<()> {
+    let (dealers, mut push_sockets): (Vec<_>, Vec<_>) =
+        data_sockets.into_iter().map(|sockets| (sockets.dealer, sockets.push)).unzip();
+    let mut input_streams =
+        stream::select_all(dealers.into_iter().map(dealer_input_stream).map(Box::pin));
+
+    loop {
+        tokio::select! {
+            biased;
+            _ = shutdown.cancelled() => return Ok(()),
+
+            output = output_rx.recv() => {
+                let output = output
+                    .ok_or_else(|| anyhow!("mock engine output channel closed"))?;
+                send_engine_outputs_to_client(&mut push_sockets, output).await?;
+            }
+
+            input = input_streams.next() => {
+                let input = input
+                    .ok_or_else(|| anyhow!("mock engine input streams closed"))??;
+                input_tx
+                    .send(input)
+                    .map_err(|_| anyhow!("mock engine state task shut down"))?;
+            }
+        }
+    }
+}

rust/src/mock-engine/src/lib.rs

@@ -0,0 +1,138 @@
+use anyhow::{Context, Result, bail};
+use clap::Parser;
+use tokio::sync::mpsc;
+use tokio::task::JoinSet;
+use tokio_util::sync::CancellationToken;
+use tracing::{error, info};
+use vllm_engine_core_client::EngineId;
+use vllm_engine_core_client::mock_engine::{
+    MockEngineConfig, MockEngineSockets, connect_to_frontend,
+};
+
+pub mod engine;
+pub mod io;
+
+/// Standalone engine-core protocol emulator for frontend stress testing.
+#[derive(Debug, Clone, Parser)]
+#[command(
+    name = "vllm-mock-engine",
+    about = "Run a mock vLLM headless engine for Rust frontend stress testing."
+)]
+pub struct Opt {
+    /// Frontend-owned ZMQ handshake address.
+    #[arg(long, default_value = "tcp://127.0.0.1:29550")]
+    pub handshake_address: String,
+
+    /// Number of mock engine identities to register with the frontend.
+    #[arg(long, default_value_t = 1)]
+    pub engine_count: usize,
+
+    /// Number of accepted output tokens included in each EngineCoreOutput.
+    #[arg(long, default_value_t = 1)]
+    pub output_token_chunk_size: usize,
+
+    /// Random token IDs are sampled uniformly from 0..vocab_size.
+    #[arg(long, default_value_t = 32_000)]
+    pub vocab_size: u32,
+
+    /// Base seed for deterministic random token generation.
+    #[arg(long, default_value_t = 0)]
+    pub seed: u64,
+
+    /// Log a summary line for each request.
+    #[arg(long)]
+    pub log_requests: bool,
+}
+
+/// Run one mock engine until shutdown or transport failure.
+async fn run_engine(engine_index: u32, opt: Opt, shutdown: CancellationToken) -> Result<()> {
+    let MockEngineSockets { data_sockets, .. } = connect_to_frontend(
+        &opt.handshake_address,
+        EngineId::from_engine_index(engine_index),
+        MockEngineConfig::default(),
+    )
+    .await
+    .with_context(|| format!("mock engine {engine_index} failed to connect to frontend"))?;
+
+    info!(engine_index, "mock engine connected to frontend");
+
+    let (input_tx, input_rx) = mpsc::unbounded_channel();
+    let (output_tx, output_rx) = mpsc::channel(64);
+
+    // IO loop: dealer -> input_tx, output_rx -> push
+    let mut io_loop = tokio::spawn(io::run_io_loop(
+        data_sockets,
+        input_tx,
+        output_rx,
+        shutdown.clone(),
+    ));
+    // Engine loop: input_rx -> engine logic -> output_tx
+    let mut engine_loop = tokio::spawn(engine::run_engine_loop(
+        engine_index,
+        opt,
+        input_rx,
+        output_tx,
+        shutdown.clone(),
+    ));
+
+    tokio::select! {
+        biased;
+        _ = shutdown.cancelled() => {
+            io_loop.abort();
+            engine_loop.abort();
+            io_loop.await.ok();
+            engine_loop.await.ok();
+        }
+
+        result = &mut io_loop => {
+            error!(engine_index, "mock engine IO loop exited unexpectedly");
+            engine_loop.abort();
+            engine_loop.await.ok();
+            result??;
+        }
+        result = &mut engine_loop => {
+            error!(engine_index, "mock engine loop exited unexpectedly");
+            io_loop.abort();
+            io_loop.await.ok();
+            result??;
+        }
+    }
+
+    info!(engine_index, "mock engine shut down");
+    Ok(())
+}
+
+/// Run all requested mock engines until cancellation or one engine task fails.
+pub async fn run(opt: Opt, shutdown: CancellationToken) -> Result<()> {
+    info!(?opt, "starting mock engine");
+
+    let mut engines = JoinSet::new();
+    for engine_index in 0..opt.engine_count {
+        engines.spawn(run_engine(
+            engine_index as u32,
+            opt.clone(),
+            shutdown.clone(),
+        ));
+    }
+
+    tokio::select! {
+        biased;
+        _ = shutdown.cancelled() => {
+            engines.abort_all();
+            while engines.join_next().await.is_some() {}
+            Ok(())
+        }
+
+        joined = engines.join_next() => {
+            match joined {
+                Some(Ok(Ok(()))) => bail!("mock engine exited unexpectedly"),
+                Some(Ok(Err(error))) => Err(error),
+                Some(Err(error)) => Err(error).context("mock engine task join failed"),
+                None => Ok(()),
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests;

rust/src/mock-engine/src/main.rs

@@ -0,0 +1,38 @@
+use anyhow::{Context, Result};
+use clap::Parser as _;
+use tokio_util::sync::CancellationToken;
+use tracing::{Level, info};
+use vllm_mock_engine::Opt;
+
+fn init_tracing() {
+    tracing_subscriber::fmt().with_max_level(Level::INFO).init();
+}
+
+/// Create a cancellation token that is triggered by Ctrl-C.
+fn shutdown_signal() -> CancellationToken {
+    let token = CancellationToken::new();
+    let shutdown = token.clone();
+
+    tokio::spawn(async move {
+        tokio::signal::ctrl_c().await.expect("failed to install Ctrl-C signal handler");
+        info!("received shutdown signal (Ctrl-C), shutting down...");
+        shutdown.cancel();
+    });
+
+    token
+}
+
+fn main() -> Result<()> {
+    init_tracing();
+    let opt = Opt::parse();
+
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .context("failed to build Tokio runtime")?;
+
+    runtime.block_on(async move {
+        let shutdown = shutdown_signal();
+        vllm_mock_engine::run(opt, shutdown).await
+    })
+}

rust/src/mock-engine/src/tests.rs

@@ -0,0 +1,197 @@
+use std::net::TcpListener;
+use std::time::Duration;
+
+use anyhow::Result;
+use futures::StreamExt as _;
+use tokio::time::timeout;
+use tokio_util::sync::CancellationToken;
+use vllm_engine_core_client::protocol::{
+    EngineCoreFinishReason, EngineCoreRequest, EngineCoreSamplingParams,
+};
+use vllm_engine_core_client::test_utils::IpcNamespace;
+use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, TransportMode};
+
+use crate::{Opt, run};
+
+fn free_tcp_address() -> String {
+    let listener = TcpListener::bind("127.0.0.1:0").expect("bind free port");
+    let port = listener.local_addr().expect("local addr").port();
+    drop(listener);
+    format!("tcp://127.0.0.1:{port}")
+}
+
+fn client_config(handshake_address: String, engine_count: usize) -> EngineCoreClientConfig {
+    EngineCoreClientConfig {
+        transport_mode: TransportMode::HandshakeOwner {
+            handshake_address,
+            advertised_host: "127.0.0.1".to_string(),
+            engine_count,
+            ready_timeout: Duration::from_secs(5),
+            local_input_address: None,
+            local_output_address: None,
+        },
+        coordinator_mode: None,
+        model_name: "mock-model".to_string(),
+        client_index: 0,
+    }
+}
+
+async fn connect_with_mock(
+    handshake_address: String,
+    engine_count: usize,
+    output_token_chunk_size: usize,
+) -> (
+    EngineCoreClient,
+    CancellationToken,
+    tokio::task::JoinHandle<Result<()>>,
+) {
+    let shutdown = CancellationToken::new();
+    let task = tokio::spawn(run(
+        Opt {
+            handshake_address: handshake_address.clone(),
+            engine_count,
+            output_token_chunk_size,
+            vocab_size: 32_000,
+            seed: 0,
+            log_requests: false,
+        },
+        shutdown.clone(),
+    ));
+
+    let client = timeout(
+        Duration::from_secs(5),
+        EngineCoreClient::connect(client_config(handshake_address, engine_count)),
+    )
+    .await
+    .expect("client connect timeout")
+    .expect("connect client");
+
+    (client, shutdown, task)
+}
+
+fn sample_request(request_id: &str, max_tokens: u32) -> EngineCoreRequest {
+    EngineCoreRequest {
+        request_id: request_id.to_string(),
+        prompt_token_ids: Some(vec![1, 2, 3]),
+        sampling_params: Some(EngineCoreSamplingParams {
+            max_tokens,
+            ..EngineCoreSamplingParams::for_test()
+        }),
+        arrival_time: 0.0,
+        ..Default::default()
+    }
+}
+
+async fn shutdown_mock(
+    client: EngineCoreClient,
+    shutdown: CancellationToken,
+    task: tokio::task::JoinHandle<Result<()>>,
+) {
+    client.shutdown().await.expect("client shutdown");
+    shutdown.cancel();
+    task.await.expect("mock join").expect("mock run");
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn mock_engine_connects_over_tcp() {
+    let handshake_address = free_tcp_address();
+    let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
+    assert_eq!(client.engine_count(), 1);
+    assert_eq!(client.engine_identities()[0], &[0, 0]);
+    assert_eq!(client.max_model_len(), Some(1024 * 1024));
+    shutdown_mock(client, shutdown, task).await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn mock_engine_connects_over_ipc() {
+    let ipc = IpcNamespace::new().expect("ipc namespace");
+    let handshake_address = ipc.handshake_endpoint();
+    let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
+    assert_eq!(client.engine_count(), 1);
+    assert_eq!(client.engine_identities()[0], &[0, 0]);
+    shutdown_mock(client, shutdown, task).await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn mock_engine_registers_multiple_identities() {
+    let handshake_address = free_tcp_address();
+    let (client, shutdown, task) = connect_with_mock(handshake_address, 2, 1).await;
+    assert_eq!(client.engine_count(), 2);
+    assert_eq!(client.engine_identities(), vec![&[0, 0][..], &[1, 0][..]]);
+    shutdown_mock(client, shutdown, task).await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn chunk_size_one_outputs_one_token_per_update() {
+    let handshake_address = free_tcp_address();
+    let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
+    let mut stream = client.call(sample_request("req-1", 3)).await.expect("call");
+
+    let first = stream.next().await.expect("first").expect("first ok");
+    assert_eq!(first.new_token_ids.len(), 1);
+    assert_eq!(first.finish_reason, None);
+    let second = stream.next().await.expect("second").expect("second ok");
+    assert_eq!(second.new_token_ids.len(), 1);
+    assert_eq!(second.finish_reason, None);
+    let third = stream.next().await.expect("third").expect("third ok");
+    assert_eq!(third.new_token_ids.len(), 1);
+    assert_eq!(third.finish_reason, Some(EngineCoreFinishReason::Length));
+    assert!(stream.next().await.is_none());
+
+    shutdown_mock(client, shutdown, task).await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn chunk_size_clips_final_output_to_max_tokens() {
+    let handshake_address = free_tcp_address();
+    let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 4).await;
+    let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call");
+
+    let first = stream.next().await.expect("first").expect("first ok");
+    assert_eq!(first.new_token_ids.len(), 4);
+    assert_eq!(first.finish_reason, None);
+    let second = stream.next().await.expect("second").expect("second ok");
+    assert_eq!(second.new_token_ids.len(), 2);
+    assert_eq!(second.finish_reason, Some(EngineCoreFinishReason::Length));
+    assert!(stream.next().await.is_none());
+
+    shutdown_mock(client, shutdown, task).await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn abort_cancels_active_request_and_emits_terminal_output() {
+    let handshake_address = free_tcp_address();
+    let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
+    let mut stream = client.call(sample_request("req-abort", 1_000_000)).await.expect("call");
+    let first = stream.next().await.expect("first").expect("first ok");
+    assert_eq!(first.finish_reason, None);
+
+    client.abort(&["req-abort".to_string()]).await.expect("abort");
+
+    loop {
+        let output = timeout(Duration::from_secs(5), stream.next())
+            .await
+            .expect("stream timeout")
+            .expect("terminal output")
+            .expect("output ok");
+        if output.finish_reason.is_some() {
+            assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Abort));
+            break;
+        }
+    }
+
+    shutdown_mock(client, shutdown, task).await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn utility_requests_return_minimal_success_responses() {
+    let handshake_address = free_tcp_address();
+    let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
+
+    assert!(!client.is_sleeping().await.expect("is sleeping"));
+    assert!(client.reset_prefix_cache(false, false).await.expect("reset prefix cache"));
+    client.reset_mm_cache().await.expect("reset mm cache");
+    client.reset_encoder_cache().await.expect("reset encoder cache");
+
+    shutdown_mock(client, shutdown, task).await;
+}

rust/src/text/Cargo.toml

@@ -26,6 +26,7 @@ vllm-tokenizer.workspace = true
 [dev-dependencies]
 expect-test.workspace = true
 futures.workspace = true
+serial_test.workspace = true
 tempfile.workspace = true
 tokio.workspace = true
 vllm-llm = { workspace = true, features = ["test-util"] }

rust/src/text/src/backend/hf/model_files.rs

@@ -401,7 +401,7 @@ mod tests {
     }
 
     #[tokio::test]
-    #[ignore = "requires network access to Hugging Face and downloads the real Kimi K2.5 tokenizer"]
+    #[ignore = "too slow for CI and requires network access to Hugging Face"]
     async fn tiktoken_real_kimi_k25_tokenizer_files_load_and_handle_special_tokens() {
         let files = ResolvedModelFiles::new("moonshotai/Kimi-K2.5")
             .await

rust/src/text/src/lower.rs

@@ -235,6 +235,8 @@ fn merge_unique_token_ids(
 mod tests {
     use std::collections::BTreeSet;
 
+    use serial_test::file_serial;
+
     use super::*;
     use crate::backend::hf::HfTextBackend;
     use crate::backend::{SamplingHints, TextBackend as _};
@@ -386,7 +388,7 @@ mod tests {
     }
 
     #[tokio::test]
-    #[ignore = "requires network access to Hugging Face"]
+    #[file_serial(hf_qwen3)]
     async fn lower_text_request_uses_real_qwen_generation_defaults() {
         let backend = HfTextBackend::from_model("Qwen/Qwen3-0.6B")
             .await
@@ -410,12 +412,8 @@ mod tests {
                 default_top_k: Some(
                     20,
                 ),
-                default_min_p: Some(
-                    0.1,
-                ),
-                default_repetition_penalty: Some(
-                    1.2,
-                ),
+                default_min_p: None,
+                default_repetition_penalty: None,
                 default_max_tokens: None,
                 max_model_len: Some(
                     40960,
@@ -439,10 +437,10 @@ mod tests {
                 min_tokens: 0,
                 logprobs: None,
                 prompt_logprobs: None,
-                min_p: 0.1,
+                min_p: 0.0,
                 frequency_penalty: 0.0,
                 presence_penalty: 0.0,
-                repetition_penalty: 1.2,
+                repetition_penalty: 1.0,
                 stop_token_ids: [
                     151643,
                 ],
@@ -453,6 +451,13 @@ mod tests {
                     151643,
                     151645,
                 },
+                logit_bias: None,
+                allowed_token_ids: None,
+                bad_words_token_ids: None,
+                structured_outputs: None,
+                logprob_token_ids: None,
+                skip_reading_prefix_cache: None,
+                extra_args: None,
             }
         "#]]
         .assert_debug_eq(&params);

rust/src/tool-parser/Cargo.toml

@@ -8,6 +8,7 @@ license.workspace = true
 test-util = []
 
 [dependencies]
+easy-ext.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true

rust/src/tool-parser/benches/gemma4.rs

@@ -9,6 +9,7 @@ use utils::feed_parser;
 
 const CHUNK_CHARS: usize = 7;
 const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
+const LONG_TOOL_ARGUMENT_REPEATS: usize = 256;
 
 fn mixed_fixture() -> String {
     concat!(
@@ -48,6 +49,24 @@ fn long_normal_text_fixture() -> String {
     line.repeat(LONG_NORMAL_TEXT_REPEATS)
 }
 
+fn long_tool_argument_fixture() -> String {
+    let line =
+        "<section><p>Literal } and <tool_call|> marker-shaped text inside content.</p></section>\n";
+    format!(
+        concat!(
+            "I will write the file.\n",
+            "<|tool_call>",
+            "call:write_file{{",
+            "path:<|\"|>index.html<|\"|>,",
+            "content:<|\"|>{}<|\"|>",
+            "}}",
+            "<tool_call|>",
+            "Done."
+        ),
+        line.repeat(LONG_TOOL_ARGUMENT_REPEATS)
+    )
+}
+
 fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
     Gemma4ToolParser::create(tools).expect("Gemma4 parser should initialize")
 }
@@ -98,6 +117,7 @@ fn run_stream_group(
 fn bench_gemma4(c: &mut Criterion) {
     let tools = test_tools();
     let mixed_text = mixed_fixture();
+    let long_tool_argument = long_tool_argument_fixture();
     let long_normal_text = long_normal_text_fixture();
 
     run_stream_group(
@@ -110,6 +130,16 @@ fn bench_gemma4(c: &mut Criterion) {
         2,
     );
 
+    run_stream_group(
+        c,
+        "gemma4/long_tool_argument",
+        &tools,
+        &long_tool_argument,
+        CHUNK_CHARS,
+        "I will write the file.\nDone.",
+        1,
+    );
+
     run_stream_group(
         c,
         "gemma4/long_normal_text",

rust/src/tool-parser/src/deepseek_dsml/deepseek_v32.rs

@@ -1,5 +1,5 @@
 use super::{DeepSeekDsmlToolParser, DsmlTokens};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for DeepSeek V3.2 models.
 ///
@@ -33,7 +33,6 @@ impl DeepSeekV32ToolParser {
 }
 
 impl ToolParser for DeepSeekV32ToolParser {
-    /// Create a boxed DeepSeek V3.2 tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -41,20 +40,21 @@ impl ToolParser for DeepSeekV32ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Preserve DSML special tokens while decoding.
     fn preserve_special_tokens(&self) -> bool {
         true
     }
 
-    /// Push one decoded text chunk through the DSML parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.0.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.0.parse_into(chunk, output)
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.0.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.0.reset()
+    }
 }
 
 #[cfg(test)]
@@ -63,8 +63,8 @@ mod tests {
     use thiserror_ext::AsReport;
 
     use super::DeepSeekV32ToolParser;
-    use crate::ToolParser;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
+    use crate::{ToolParser, ToolParserTestExt as _};
 
     fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
         let params = params
@@ -84,27 +84,27 @@ mod tests {
     #[test]
     fn deepseek_v32_parse_complete_without_tool_call_keeps_text() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = parser.parse_complete("Hello, world!").unwrap();
+        let output = parser.parse_complete("Hello, world!").unwrap();
 
-        assert_eq!(result.normal_text, "Hello, world!");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "Hello, world!");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn deepseek_v32_parse_complete_extracts_single_tool_call() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "get_weather",
                 &[("location", "SF"), ("date", "2024-01-16")],
             ))
             .unwrap();
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "location": "SF",
                 "date": "2024-01-16"
@@ -119,16 +119,16 @@ mod tests {
             "Thinking... {}",
             build_tool_call("get_weather", &[("location", "NYC")])
         );
-        let result = parser.parse_complete(&output).unwrap();
+        let output = parser.parse_complete(&output).unwrap();
 
-        assert_eq!(result.normal_text, "Thinking... ");
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.normal_text, "Thinking... ");
+        assert_eq!(output.calls.len(), 1);
     }
 
     #[test]
     fn deepseek_v32_parse_complete_converts_schema_types() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(
                 "<｜DSML｜function_calls>\n\
                  <｜DSML｜invoke name=\"convert\">\n\
@@ -142,9 +142,9 @@ mod tests {
             )
             .unwrap();
 
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.calls.len(), 1);
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "whole": 5.0,
                 "flag": true,
@@ -158,7 +158,7 @@ mod tests {
     #[test]
     fn deepseek_v32_parse_complete_string_attr_overrides_schema_types() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(
                 "<｜DSML｜function_calls>\n\
                  <｜DSML｜invoke name=\"convert\">\n\
@@ -172,9 +172,9 @@ mod tests {
             )
             .unwrap();
 
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.calls.len(), 1);
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "whole": "5.0",
                 "flag": "true",
@@ -188,7 +188,7 @@ mod tests {
     #[test]
     fn deepseek_v32_parse_complete_unescapes_literal_closing_tags_in_parameter_value() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "get_weather",
                 &[
@@ -202,7 +202,7 @@ mod tests {
             .unwrap();
 
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "location": "Hangzhou </｜DSML｜parameter></｜DSML｜invoke></｜DSML｜function_calls>",
                 "date": "2026-05-08",
@@ -213,7 +213,7 @@ mod tests {
     #[test]
     fn deepseek_v32_streaming_extracts_single_tool_call() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "<｜DSML｜function_calls>\n",
@@ -224,11 +224,11 @@ mod tests {
             ],
         );
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({ "location": "SF" })
         );
     }
@@ -236,7 +236,7 @@ mod tests {
     #[test]
     fn deepseek_v32_streaming_preserves_prefix_text() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "Thinking... ",
@@ -248,23 +248,23 @@ mod tests {
             ],
         );
 
-        assert_eq!(result.normal_text, "Thinking... ");
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.normal_text, "Thinking... ");
+        assert_eq!(output.calls.len(), 1);
     }
 
     #[test]
     fn deepseek_v32_streaming_without_tool_call_emits_text_incrementally() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = collect_stream(&mut parser, &["Hello, ", "world!"]);
+        let output = collect_stream(&mut parser, &["Hello, ", "world!"]);
 
-        assert_eq!(result.normal_text, "Hello, world!");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "Hello, world!");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn deepseek_v32_streaming_extracts_multiple_tool_calls_in_order() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[&format!(
                 "{}\n{}",
@@ -274,17 +274,17 @@ mod tests {
             )],
         );
 
-        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("get_weather"));
-        assert_eq!(result.calls[0].tool_index, 0);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({ "location": "SF" })
         );
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
             json!({ "location": "NYC" })
         );
     }
@@ -294,11 +294,11 @@ mod tests {
         let text = build_tool_call("get_weather", &[("location", "SF")]);
         let chunks = split_by_chars(&text, 5);
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.calls.len(), 1);
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({ "location": "SF" })
         );
     }
@@ -306,7 +306,7 @@ mod tests {
     #[test]
     fn deepseek_v32_streaming_handles_bpe_chunked_dsml_opener() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "<｜DSML｜",
@@ -333,11 +333,11 @@ mod tests {
             ],
         );
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({ "location": "Beijing" })
         );
     }
@@ -345,12 +345,12 @@ mod tests {
     #[test]
     fn deepseek_v32_streaming_truncated_parameter_does_not_leak_eos() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        parser.push("<｜DSML｜function_calls>\n").unwrap();
-        parser.push("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
+        parser.parse_chunk("<｜DSML｜function_calls>\n").unwrap();
+        parser.parse_chunk("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
         parser
-            .push("<｜DSML｜parameter name=\"location\" string=\"true\">Tokyo")
+            .parse_chunk("<｜DSML｜parameter name=\"location\" string=\"true\">Tokyo")
             .unwrap();
-        parser.push("<｜end▁of▁sentence｜>").unwrap();
+        parser.parse_chunk("<｜end▁of▁sentence｜>").unwrap();
 
         let error = parser.finish().unwrap_err();
         assert!(error.to_report_string().contains("incomplete DeepSeek DSML tool call"));
@@ -358,7 +358,7 @@ mod tests {
     #[test]
     fn deepseek_v32_streaming_drops_eos_after_complete_tool_calls() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "<｜DSML｜function_calls>\n",
@@ -369,15 +369,15 @@ mod tests {
             ],
         );
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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]
     fn deepseek_v32_streaming_ignores_text_after_complete_tool_calls() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "<｜DSML｜function_calls>\n",
@@ -389,17 +389,19 @@ mod tests {
             ],
         );
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
+        assert!(output.normal_text.is_empty());
+        assert_eq!(output.calls.len(), 1);
     }
 
     #[test]
     fn deepseek_v32_streaming_does_not_emit_incomplete_invoke() {
         let mut parser = DeepSeekV32ToolParser::new(&test_tools());
-        parser.push("<｜DSML｜function_calls>\n").unwrap();
-        parser.push("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
+        parser.parse_chunk("<｜DSML｜function_calls>\n").unwrap();
+        parser.parse_chunk("<｜DSML｜invoke name=\"get_weather\">\n").unwrap();
         parser
-            .push("<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n")
+            .parse_chunk(
+                "<｜DSML｜parameter name=\"location\" string=\"true\">SF</｜DSML｜parameter>\n",
+            )
             .unwrap();
 
         let error = parser.finish().unwrap_err();

rust/src/tool-parser/src/deepseek_dsml/deepseek_v4.rs

@@ -1,5 +1,5 @@
 use super::{DeepSeekDsmlToolParser, DsmlTokens};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for DeepSeek V4 models.
 ///
@@ -36,7 +36,6 @@ impl DeepSeekV4ToolParser {
 }
 
 impl ToolParser for DeepSeekV4ToolParser {
-    /// Create a boxed DeepSeek V4 tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -44,27 +43,29 @@ impl ToolParser for DeepSeekV4ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Preserve DSML special tokens while decoding.
     fn preserve_special_tokens(&self) -> bool {
         true
     }
 
-    /// Push one decoded text chunk through the DSML parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.0.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.0.parse_into(chunk, output)
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.0.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.0.reset()
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use serde_json::{Value, json};
 
-    use super::{DeepSeekV4ToolParser, ToolParser};
+    use super::DeepSeekV4ToolParser;
+    use crate::ToolParserTestExt as _;
     use crate::test_utils::{collect_stream, test_tools};
 
     fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
@@ -85,18 +86,18 @@ mod tests {
     #[test]
     fn deepseek_v4_parse_complete_reuses_dsml_parser_with_tool_calls_token() {
         let mut parser = DeepSeekV4ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "get_weather",
                 &[("location", "SF"), ("date", "2024-01-16")],
             ))
             .unwrap();
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "location": "SF",
                 "date": "2024-01-16"
@@ -107,7 +108,7 @@ mod tests {
     #[test]
     fn deepseek_v4_streaming_handles_tool_calls_token_split_across_chunks() {
         let mut parser = DeepSeekV4ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "Thinking... ",
@@ -122,11 +123,11 @@ mod tests {
             ],
         );
 
-        assert_eq!(result.normal_text, "Thinking... ");
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({ "location": "Beijing" })
         );
     }

rust/src/tool-parser/src/deepseek_dsml/mod.rs

@@ -6,7 +6,7 @@ use winnow::token::{literal, rest, take_until};
 
 use super::parameters::ToolSchemas;
 use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
-use super::{Result, ToolCallDelta, ToolParseResult};
+use super::{Result, ToolCallDelta, ToolParserOutput};
 use crate::Tool;
 
 mod deepseek_v32;
@@ -89,10 +89,10 @@ impl DeepSeekDsmlToolParser {
     }
 
     /// Apply one parsed DSML event to parser state and output.
-    fn apply_event(&mut self, event: DsmlEvent, result: &mut ToolParseResult) -> Result<()> {
+    fn apply_event(&mut self, event: DsmlEvent, output: &mut ToolParserOutput) -> Result<()> {
         match event {
             DsmlEvent::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
             }
             DsmlEvent::ToolCallsStart => self.mode = DsmlMode::ToolBlock,
             DsmlEvent::Invoke { name, raw_params } => {
@@ -112,7 +112,7 @@ impl DeepSeekDsmlToolParser {
                 let arguments = serde_json::to_string(&arguments)
                     .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
 
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index: self.emitted_invoke_count,
                     name: Some(name),
                     arguments,
@@ -125,46 +125,41 @@ impl DeepSeekDsmlToolParser {
         Ok(())
     }
 
-    /// Reset all streaming state.
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.mode = DsmlMode::Text;
         self.emitted_invoke_count = 0;
+        std::mem::take(&mut self.buffer)
     }
 
-    /// Push one decoded text chunk through the DSML parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         // Extract tool calls from streaming model output.
         //
         // Uses a buffer-until-complete-invoke strategy: text is buffered until
         // a complete invoke block is available, then parsed and emitted in one
         // shot.
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
         while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
             parse_next_dsml_event(input, self.mode, self.tokens)
         })? {
-            self.apply_event(event, &mut result)?;
+            self.apply_event(event, output)?;
             self.buffer.drain(..consumed_len);
         }
 
-        Ok(result)
+        Ok(())
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
-        let mut result = ToolParseResult::default();
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
         match self.mode {
-            DsmlMode::Text => result.normal_text.push_str(&self.buffer),
+            DsmlMode::Text => output.normal_text.push_str(&self.buffer),
             DsmlMode::Done => {}
             DsmlMode::ToolBlock => {
-                self.reset();
                 return Err(parsing_failed!("incomplete DeepSeek DSML tool call"));
             }
         }
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
     }
 }
 

rust/src/tool-parser/src/deepseek_json/deepseek_v3.rs

@@ -1,5 +1,5 @@
 use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for DeepSeek V3 JSON-fenced tool calls.
 ///
@@ -25,7 +25,6 @@ impl DeepSeekV3ToolParser {
 }
 
 impl ToolParser for DeepSeekV3ToolParser {
-    /// Create a boxed DeepSeek V3 tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -33,15 +32,17 @@ impl ToolParser for DeepSeekV3ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the DeepSeek V3 parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.0.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.0.parse_into(chunk, output)
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.0.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.0.reset()
+    }
 }
 
 #[cfg(test)]
@@ -55,7 +56,7 @@ mod tests {
         V3_JSON_START,
     };
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
-    use crate::{ToolParseResult, ToolParser};
+    use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
 
     fn v3_tool_call(function_name: &str, arguments: &str) -> String {
         format!(
@@ -70,39 +71,39 @@ mod tests {
     #[test]
     fn deepseek_v3_parse_complete_without_tool_call_keeps_text() {
         let mut parser = DeepSeekV3ToolParser::new(&test_tools());
-        let result = parser.parse_complete("Hello, world!").unwrap();
+        let output = parser.parse_complete("Hello, world!").unwrap();
 
-        assert_eq!(result.normal_text, "Hello, world!");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "Hello, world!");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn deepseek_v3_parse_complete_extracts_raw_json_arguments() {
         let mut parser = DeepSeekV3ToolParser::new(&test_tools());
         let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
-        let result = parser
+        let output = parser
             .parse_complete(&format!(
                 "Let me check.\n{} trailing text",
                 tool_section(&[v3_tool_call("get_weather", arguments)])
             ))
             .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!(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]
     fn deepseek_v3_does_not_validate_or_normalize_arguments() {
         let mut parser = DeepSeekV3ToolParser::new(&test_tools());
         let arguments = r#"{"location":"Tokyo",}"#;
-        let result = parser
+        let output = parser
             .parse_complete(&tool_section(&[v3_tool_call("get_weather", arguments)]))
             .unwrap();
 
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
@@ -122,23 +123,23 @@ mod tests {
             TOOL_CALLS_END,
         ];
 
-        let mut result = ToolParseResult::default();
+        let mut output = ToolParserOutput::default();
         let mut observed_arguments = Vec::new();
         for chunk in chunks {
-            let next = parser.push(chunk).unwrap();
+            let next = parser.parse_chunk(chunk).unwrap();
             observed_arguments.extend(
                 next.calls
                     .iter()
                     .filter(|call| call.name.is_none())
                     .map(|call| call.arguments.clone()),
             );
-            result.append(next);
+            output.append(next);
         }
-        result.append(parser.finish().unwrap());
+        output.append(parser.finish().unwrap());
 
         assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
         assert_eq!(
-            result.coalesce_calls().calls[0].arguments,
+            output.coalesce_calls().calls[0].arguments,
             r#"{"location":"Beijing"}"#
         );
     }
@@ -152,11 +153,11 @@ mod tests {
         let chunks = split_by_chars(&input, 5);
         let mut parser = DeepSeekV3ToolParser::new(&test_tools());
 
-        let result = collect_stream(&mut parser, &chunks);
+        let output = 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!(output.normal_text, "hello ");
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
     }
 
     #[test]
@@ -165,10 +166,10 @@ mod tests {
         let arguments = format!("{{\"text\":\"literal {V3_ARGUMENT_END} inside\"}}");
         let input = tool_section(&[v3_tool_call("echo", &arguments)]);
 
-        let result = parser.parse_complete(&input).unwrap();
+        let output = parser.parse_complete(&input).unwrap();
 
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
@@ -180,10 +181,10 @@ mod tests {
         let chunks = split_by_chars(&input, 7);
         let mut parser = DeepSeekV3ToolParser::new(&test_tools());
 
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
         expect![[r#"
-            ToolParseResult {
+            ToolParserOutput {
                 normal_text: "",
                 calls: [
                     ToolCallDelta {
@@ -203,14 +204,14 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 
     #[test]
     fn deepseek_v3_finish_fails_incomplete_tool_call() {
         let mut parser = DeepSeekV3ToolParser::new(&test_tools());
         parser
-            .push(&format!(
+            .parse_chunk(&format!(
                 "{TOOL_CALLS_START}{TOOL_CALL_START}function{TOOL_CALL_SEPARATOR}get_weather{V3_JSON_START}{{\"location\""
             ))
             .unwrap();
@@ -228,7 +229,7 @@ mod tests {
             "{TOOL_CALLS_START}{TOOL_CALL_START}tool{TOOL_CALL_SEPARATOR}get_weather{V3_JSON_START}{{}}"
         );
 
-        let error = parser.push(&input).unwrap_err();
+        let error = parser.parse_chunk(&input).unwrap_err();
 
         expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
     }

rust/src/tool-parser/src/deepseek_json/deepseek_v31.rs

@@ -1,5 +1,5 @@
 use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for DeepSeek V3.1 raw JSON tool calls.
 ///
@@ -21,7 +21,6 @@ impl DeepSeekV31ToolParser {
 }
 
 impl ToolParser for DeepSeekV31ToolParser {
-    /// Create a boxed DeepSeek V3.1 tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -29,15 +28,17 @@ impl ToolParser for DeepSeekV31ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the DeepSeek V3.1 parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.0.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.0.parse_into(chunk, output)
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.0.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.0.reset()
+    }
 }
 
 #[cfg(test)]
@@ -50,7 +51,7 @@ mod tests {
         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::{ToolParseResult, ToolParser};
+    use crate::{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}")
@@ -63,39 +64,39 @@ mod tests {
     #[test]
     fn deepseek_v31_parse_complete_without_tool_call_keeps_text() {
         let mut parser = DeepSeekV31ToolParser::new(&test_tools());
-        let result = parser.parse_complete("Hello, world!").unwrap();
+        let output = parser.parse_complete("Hello, world!").unwrap();
 
-        assert_eq!(result.normal_text, "Hello, world!");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "Hello, world!");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn deepseek_v31_parse_complete_extracts_raw_json_arguments() {
         let mut parser = DeepSeekV31ToolParser::new(&test_tools());
         let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
-        let result = parser
+        let output = parser
             .parse_complete(&format!(
                 "Let me check.{} trailing text",
                 tool_section(&[v31_tool_call("get_weather", arguments)])
             ))
             .unwrap();
 
-        assert_eq!(result.normal_text, "Let me check.");
-        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!(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]
     fn deepseek_v31_does_not_validate_or_normalize_arguments() {
         let mut parser = DeepSeekV31ToolParser::new(&test_tools());
         let arguments = r#"{"location":"Tokyo",}"#;
-        let result = parser
+        let output = parser
             .parse_complete(&tool_section(&[v31_tool_call("get_weather", arguments)]))
             .unwrap();
 
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
@@ -113,23 +114,23 @@ mod tests {
             TOOL_CALLS_END,
         ];
 
-        let mut result = ToolParseResult::default();
+        let mut output = ToolParserOutput::default();
         let mut observed_arguments = Vec::new();
         for chunk in chunks {
-            let next = parser.push(chunk).unwrap();
+            let next = parser.parse_chunk(chunk).unwrap();
             observed_arguments.extend(
                 next.calls
                     .iter()
                     .filter(|call| call.name.is_none())
                     .map(|call| call.arguments.clone()),
             );
-            result.append(next);
+            output.append(next);
         }
-        result.append(parser.finish().unwrap());
+        output.append(parser.finish().unwrap());
 
         assert_eq!(observed_arguments, ["{\"location\":", "\"Beijing\"", "}"]);
         assert_eq!(
-            result.coalesce_calls().calls[0].arguments,
+            output.coalesce_calls().calls[0].arguments,
             r#"{"location":"Beijing"}"#
         );
     }
@@ -143,11 +144,11 @@ mod tests {
         let chunks = split_by_chars(&input, 5);
         let mut parser = DeepSeekV31ToolParser::new(&test_tools());
 
-        let result = collect_stream(&mut parser, &chunks);
+        let output = 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!(output.normal_text, "hello ");
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].arguments, r#"{"location":"Tokyo"}"#);
     }
 
     #[test]
@@ -156,10 +157,10 @@ mod tests {
         let arguments = format!(r#"{{"text":"literal {TOOL_CALL_END} inside"}}"#);
         let input = tool_section(&[v31_tool_call("echo", &arguments)]);
 
-        let result = parser.parse_complete(&input).unwrap();
+        let output = parser.parse_complete(&input).unwrap();
 
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
@@ -171,10 +172,10 @@ mod tests {
         let chunks = split_by_chars(&input, 7);
         let mut parser = DeepSeekV31ToolParser::new(&test_tools());
 
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
         expect![[r#"
-            ToolParseResult {
+            ToolParserOutput {
                 normal_text: "",
                 calls: [
                     ToolCallDelta {
@@ -194,7 +195,7 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 
     #[test]
@@ -205,18 +206,18 @@ mod tests {
         );
         let mut parser = DeepSeekV31ToolParser::new(&test_tools());
 
-        let result = collect_stream(&mut parser, &[&input]);
+        let output = collect_stream(&mut parser, &[&input]);
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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]
     fn deepseek_v31_finish_fails_incomplete_tool_call() {
         let mut parser = DeepSeekV31ToolParser::new(&test_tools());
         parser
-            .push(&format!(
+            .parse_chunk(&format!(
                 "{TOOL_CALLS_START}{TOOL_CALL_START}get_weather{TOOL_CALL_SEPARATOR}{{\"location\""
             ))
             .unwrap();
@@ -232,7 +233,7 @@ mod tests {
         let mut parser = DeepSeekV31ToolParser::new(&test_tools());
         let input = format!("{TOOL_CALLS_START}{TOOL_CALL_START}{TOOL_CALL_SEPARATOR}{{}}");
 
-        let error = parser.push(&input).unwrap_err();
+        let error = parser.parse_chunk(&input).unwrap_err();
 
         expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
     }

rust/src/tool-parser/src/deepseek_json/mod.rs

@@ -10,7 +10,7 @@ use winnow::stream::Partial;
 use winnow::token::{literal, rest, take_until};
 
 use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
-use super::{Result, ToolCallDelta, ToolParseResult};
+use super::{Result, ToolCallDelta, ToolParserOutput};
 
 pub(super) const TOOL_CALLS_START: &str = "<｜tool▁calls▁begin｜>";
 pub(super) const TOOL_CALLS_END: &str = "<｜tool▁calls▁end｜>";
@@ -92,11 +92,11 @@ impl DeepSeekJsonToolParser {
     fn apply_event(
         &mut self,
         event: DeepSeekJsonEvent,
-        result: &mut ToolParseResult,
+        output: &mut ToolParserOutput,
     ) -> Result<()> {
         match event {
             DeepSeekJsonEvent::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&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(),
                 };
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index,
                     name: Some(function_name),
                     arguments: String::new(),
@@ -120,7 +120,7 @@ impl DeepSeekJsonToolParser {
                         self.format.parser_name()
                     ));
                 };
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index,
                     name: None,
                     arguments: self.buffer[..consumed_len].to_string(),
@@ -139,26 +139,23 @@ impl DeepSeekJsonToolParser {
         Ok(())
     }
 
-    /// Push one decoded text chunk through the DeepSeek JSON parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
         while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
             parse_next_deepseek_json_event(input, &mut self.mode, self.format)
         })? {
-            self.apply_event(event, &mut result)?;
+            self.apply_event(event, output)?;
             self.buffer.drain(..consumed_len);
         }
 
-        Ok(result)
+        Ok(())
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
-        let mut result = ToolParseResult::default();
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
         match &self.mode {
-            DeepSeekJsonMode::Text => result.normal_text.push_str(&self.buffer),
+            DeepSeekJsonMode::Text => output.normal_text.push_str(&self.buffer),
             DeepSeekJsonMode::ToolBlock | DeepSeekJsonMode::Done => {}
             DeepSeekJsonMode::Header | DeepSeekJsonMode::Arguments { .. } => {
                 return Err(parsing_failed!(
@@ -167,16 +164,15 @@ impl DeepSeekJsonToolParser {
                 ));
             }
         }
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
     }
 
-    /// Reset all streaming state.
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.mode = DeepSeekJsonMode::Text;
         self.active_tool_index = None;
         self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
     }
 }
 

rust/src/tool-parser/src/gemma4.rs

@@ -1,13 +1,13 @@
 use serde_json::{Map, Number, Value};
 use winnow::ascii::multispace0 as ws0;
-use winnow::combinator::{alt, delimited, opt, separated, seq, terminated};
+use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated};
 use winnow::error::{ContextError, ErrMode, ModalResult};
 use winnow::prelude::*;
-use winnow::stream::Partial;
+use winnow::stream::{Partial, Stream};
 use winnow::token::{literal, take_till, take_until};
 
-use super::utils::{parse_buffered_event, safe_text_len};
-use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
+use super::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len};
+use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
 use crate::Tool;
 
 const TOOL_CALL_START: &str = "<|tool_call>";
@@ -19,12 +19,26 @@ type Gemma4Input<'i> = Partial<&'i str>;
 
 #[derive(Debug, Clone, PartialEq)]
 enum Gemma4Event {
-    Text {
-        len: usize,
-    },
+    Text { len: usize },
+    ToolCallStart,
+    ToolCallHeader { name: String },
+    ToolCall { args: Map<String, Value> },
+}
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+struct Gemma4ArgsScanState {
+    scanned_len: usize,
+    in_string: bool,
+}
+
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+enum Gemma4Mode {
+    #[default]
+    Text,
+    Header,
     ToolCall {
         name: String,
-        args: Map<String, Value>,
+        args_scan: Gemma4ArgsScanState,
     },
 }
 
@@ -40,6 +54,7 @@ enum Gemma4Event {
 /// Arguments are emitted only after a full Gemma4 tool call is parsed.
 pub struct Gemma4ToolParser {
     buffer: String,
+    mode: Gemma4Mode,
     emitted_tool_count: usize,
 }
 
@@ -47,20 +62,34 @@ impl Gemma4ToolParser {
     fn new(_tools: &[Tool]) -> Self {
         Self {
             buffer: String::new(),
+            mode: Gemma4Mode::default(),
             emitted_tool_count: 0,
         }
     }
 
-    fn apply_event(&mut self, event: Gemma4Event, result: &mut ToolParseResult) -> Result<()> {
+    fn apply_event(&mut self, event: Gemma4Event, output: &mut ToolParserOutput) -> Result<()> {
         match event {
             Gemma4Event::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
             }
-            Gemma4Event::ToolCall { name, args } => {
+            Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header,
+            Gemma4Event::ToolCallHeader { name } => {
+                self.mode = Gemma4Mode::ToolCall {
+                    name,
+                    args_scan: Gemma4ArgsScanState::default(),
+                };
+            }
+            Gemma4Event::ToolCall { args } => {
+                let mode = std::mem::replace(&mut self.mode, Gemma4Mode::Text);
+                let Gemma4Mode::ToolCall { name, .. } = mode else {
+                    return Err(parsing_failed!(
+                        "Gemma4 arguments without an active tool call"
+                    ));
+                };
                 let arguments = serde_json::to_string(&args)
                     .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
 
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index: self.emitted_tool_count,
                     name: Some(name),
                     arguments,
@@ -71,9 +100,25 @@ impl Gemma4ToolParser {
         Ok(())
     }
 
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
+        let raw = match std::mem::replace(&mut self.mode, Gemma4Mode::Text) {
+            Gemma4Mode::Text => std::mem::take(&mut self.buffer),
+            Gemma4Mode::Header => {
+                format!("{}{}", TOOL_CALL_START, std::mem::take(&mut self.buffer))
+            }
+            Gemma4Mode::ToolCall { name, .. } => {
+                format!(
+                    "{}{}{}{{{}",
+                    TOOL_CALL_START,
+                    CALL_PREFIX,
+                    name,
+                    std::mem::take(&mut self.buffer)
+                )
+            }
+        };
+        self.mode = Gemma4Mode::Text;
         self.emitted_tool_count = 0;
+        raw
     }
 }
 
@@ -89,56 +134,85 @@ impl ToolParser for Gemma4ToolParser {
         true
     }
 
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
-        while let Some((event, consumed_len)) =
-            parse_buffered_event(&self.buffer, parse_next_gemma4_event)?
-        {
-            self.apply_event(event, &mut result)?;
+        while let Some((event, consumed_len)) = {
+            parse_buffered_event(&self.buffer, |input| {
+                parse_next_gemma4_event(input, &mut self.mode)
+            })?
+        } {
+            self.apply_event(event, output)?;
             self.buffer.drain(..consumed_len);
         }
 
-        Ok(result)
+        Ok(())
     }
 
-    fn finish(&mut self) -> Result<ToolParseResult> {
-        let mut result = ToolParseResult::default();
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
 
-        if !self.buffer.is_empty() {
-            if self.buffer.starts_with(TOOL_CALL_START) {
-                self.reset();
+        match &self.mode {
+            Gemma4Mode::Text => output.normal_text.push_str(&self.buffer),
+            Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => {
                 return Err(parsing_failed!("incomplete Gemma4 tool call"));
             }
-            result.normal_text.push_str(&self.buffer);
         }
 
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
+    }
+
+    fn reset(&mut self) -> String {
+        Gemma4ToolParser::reset(self)
     }
 }
 
 /// Parse one Gemma4 event from buffered streaming input.
-fn parse_next_gemma4_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
-    alt((tool_call_event, safe_text_event)).parse_next(input)
+fn parse_next_gemma4_event(
+    input: &mut Gemma4Input<'_>,
+    mode: &mut Gemma4Mode,
+) -> ModalResult<Gemma4Event> {
+    match mode {
+        Gemma4Mode::Text => parse_text_event(input),
+        Gemma4Mode::Header => tool_call_header_event(input),
+        Gemma4Mode::ToolCall { args_scan, .. } => tool_call_args_event(input, args_scan),
+    }
 }
 
-/// Parse a complete Gemma4 tool call.
-// TODO: incremental parsing arguments to reduce scanning from O(n^2) to O(n).
-fn tool_call_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
-    let (name, args) = seq!(
-        _: literal(TOOL_CALL_START),
+/// Parse a Gemma4 text-mode event.
+fn parse_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
+    alt((tool_call_start_event, safe_text_event)).parse_next(input)
+}
+
+/// Parse a Gemma4 tool-call start marker.
+fn tool_call_start_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
+    literal(TOOL_CALL_START).value(Gemma4Event::ToolCallStart).parse_next(input)
+}
+
+/// Parse a Gemma4 tool-call header.
+fn tool_call_header_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
+    let (name,) = seq!(
         _: literal(CALL_PREFIX),
         gemma4_tool_name,
         _: literal("{"),
-        gemma4_args,
-        _: literal("}"),
-        _: literal(TOOL_CALL_END),
     )
     .parse_next(input)?;
+    Ok(Gemma4Event::ToolCallHeader { name })
+}
 
-    Ok(Gemma4Event::ToolCall { name, args })
+/// Parse complete Gemma4 tool-call arguments.
+fn tool_call_args_event(
+    input: &mut Gemma4Input<'_>,
+    args_scan: &mut Gemma4ArgsScanState,
+) -> ModalResult<Gemma4Event> {
+    let raw_args = gemma4_raw_args_until_tool_call_end(input, args_scan)?;
+    let Some(args_input) = raw_args.strip_suffix('}') else {
+        return Err(ErrMode::Cut(ContextError::new()));
+    };
+    let args = parse_gemma4_args(args_input)?;
+
+    Ok(Gemma4Event::ToolCall { args })
 }
 
 /// Parse a Gemma4 tool name.
@@ -155,8 +229,75 @@ fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
     safe_text_len(input, TOOL_CALL_START).map(|len| Gemma4Event::Text { len })
 }
 
+/// Parse raw Gemma4 arguments through the first end marker outside a Gemma string.
+fn gemma4_raw_args_until_tool_call_end<'i>(
+    input: &mut Gemma4Input<'i>,
+    state: &mut Gemma4ArgsScanState,
+) -> ModalResult<&'i str> {
+    let text = **input;
+    if state.scanned_len > text.len() {
+        return incomplete();
+    }
+
+    loop {
+        let rest = &text[state.scanned_len..];
+        if state.in_string {
+            let Some(string_delim) = rest.find(STRING_DELIM) else {
+                state.scanned_len = safe_scan_len(text, state.scanned_len, &[STRING_DELIM]);
+                return incomplete();
+            };
+
+            state.scanned_len += string_delim + STRING_DELIM.len();
+            state.in_string = false;
+            continue;
+        }
+
+        let next_string_delim = rest.find(STRING_DELIM);
+        let next_tool_call_end = rest.find(TOOL_CALL_END);
+        match (next_string_delim, next_tool_call_end) {
+            (Some(string_delim), Some(tool_call_end)) if tool_call_end < string_delim => {
+                let end = state.scanned_len + tool_call_end;
+                state.scanned_len = end + TOOL_CALL_END.len();
+                input.next_slice(state.scanned_len);
+                return Ok(&text[..end]);
+            }
+            (Some(string_delim), _) => {
+                state.scanned_len += string_delim + STRING_DELIM.len();
+                state.in_string = true;
+            }
+            (None, Some(tool_call_end)) => {
+                let end = state.scanned_len + tool_call_end;
+                state.scanned_len = end + TOOL_CALL_END.len();
+                input.next_slice(state.scanned_len);
+                return Ok(&text[..end]);
+            }
+            (None, None) => {
+                state.scanned_len =
+                    safe_scan_len(text, state.scanned_len, &[STRING_DELIM, TOOL_CALL_END]);
+                return incomplete();
+            }
+        }
+    }
+}
+
+/// Return the scan length while holding back a split marker prefix.
+fn safe_scan_len(text: &str, start: usize, markers: &[&str]) -> usize {
+    let max_partial = markers
+        .iter()
+        .map(|marker| partial_prefix_len(&text[start..], marker))
+        .max()
+        .unwrap_or(0);
+    text.len() - max_partial
+}
+
+/// Parse complete Gemma4 custom key-value arguments.
+fn parse_gemma4_args(args: &str) -> ModalResult<Map<String, Value>> {
+    let mut input = args;
+    terminated(gemma4_args, eof).parse_next(&mut input)
+}
+
 /// Parse Gemma4's custom key-value argument object content.
-fn gemma4_args(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
+fn gemma4_args(input: &mut &str) -> ModalResult<Map<String, Value>> {
     let pairs: Vec<(String, Value)> = delimited(
         ws0,
         terminated(
@@ -170,7 +311,7 @@ fn gemma4_args(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
 }
 
 /// Parse a Gemma4 key-value pair.
-fn gemma4_pair(input: &mut Gemma4Input<'_>) -> ModalResult<(String, Value)> {
+fn gemma4_pair(input: &mut &str) -> ModalResult<(String, Value)> {
     let (key, value) = seq!(
         _: ws0,
         gemma4_key,
@@ -184,7 +325,7 @@ fn gemma4_pair(input: &mut Gemma4Input<'_>) -> ModalResult<(String, Value)> {
 }
 
 /// Parse a Gemma4 bare key.
-fn gemma4_key(input: &mut Gemma4Input<'_>) -> ModalResult<String> {
+fn gemma4_key(input: &mut &str) -> ModalResult<String> {
     let key = take_till(1.., |char: char| char == ':').parse_next(input)?.trim();
     if key.is_empty() {
         return Err(ErrMode::Cut(ContextError::new()));
@@ -193,7 +334,7 @@ fn gemma4_key(input: &mut Gemma4Input<'_>) -> ModalResult<String> {
 }
 
 /// Parse a Gemma4 value.
-fn gemma4_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
+fn gemma4_value(input: &mut &str) -> ModalResult<Value> {
     alt((
         gemma4_string.map(|value: &str| Value::String(value.to_string())),
         gemma4_object.map(Value::Object),
@@ -204,7 +345,7 @@ fn gemma4_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
 }
 
 /// Parse a Gemma4 string delimited by `<|"|>`.
-fn gemma4_string<'i>(input: &mut Gemma4Input<'i>) -> ModalResult<&'i str> {
+fn gemma4_string<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
     delimited(
         literal(STRING_DELIM),
         take_until(0.., STRING_DELIM),
@@ -214,17 +355,17 @@ fn gemma4_string<'i>(input: &mut Gemma4Input<'i>) -> ModalResult<&'i str> {
 }
 
 /// Parse a nested Gemma4 object.
-fn gemma4_object(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
+fn gemma4_object(input: &mut &str) -> ModalResult<Map<String, Value>> {
     delimited(literal("{"), gemma4_args, literal("}")).parse_next(input)
 }
 
 /// Parse a Gemma4 array value.
-fn gemma4_array_value(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>> {
+fn gemma4_array_value(input: &mut &str) -> ModalResult<Vec<Value>> {
     delimited(literal("["), gemma4_array_content, literal("]")).parse_next(input)
 }
 
 /// Parse Gemma4 array content.
-fn gemma4_array_content(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>> {
+fn gemma4_array_content(input: &mut &str) -> ModalResult<Vec<Value>> {
     delimited(
         ws0,
         terminated(
@@ -237,14 +378,14 @@ fn gemma4_array_content(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>>
 }
 
 /// Parse a Gemma4 bare scalar.
-fn gemma4_bare_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
+fn gemma4_bare_value(input: &mut &str) -> ModalResult<Value> {
     take_till(1.., |char: char| matches!(char, ',' | '}' | ']'))
         .map(parse_gemma4_scalar)
         .parse_next(input)
 }
 
 /// Parse a Gemma4 comma separator.
-fn comma_separator(input: &mut Gemma4Input<'_>) -> ModalResult<()> {
+fn comma_separator(input: &mut &str) -> ModalResult<()> {
     delimited(ws0, literal(","), ws0).void().parse_next(input)
 }
 
@@ -282,29 +423,15 @@ mod tests {
     use winnow::combinator::{eof, terminated};
     use winnow::error::ErrMode;
     use winnow::prelude::*;
-    use winnow::stream::Partial;
 
     use super::{
-        Gemma4ToolParser, ToolCallDelta, ToolParseResult, ToolParser, gemma4_args,
-        gemma4_array_content,
+        Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_array_content,
+        parse_gemma4_args,
     };
-    use crate::Tool;
-
-    fn parse_gemma4_args(args: &str) -> super::Result<serde_json::Map<String, Value>> {
-        let mut input = Partial::new(args);
-        let _ = input.complete();
-        match terminated(gemma4_args, eof).parse_next(&mut input) {
-            Ok(value) => Ok(value),
-            Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 arguments")),
-            Err(ErrMode::Backtrack(error) | ErrMode::Cut(error)) => {
-                Err(parsing_failed!("{}", error))
-            }
-        }
-    }
+    use crate::{Tool, ToolParserTestExt as _};
 
     fn parse_gemma4_array(array: &str) -> super::Result<Vec<Value>> {
-        let mut input = Partial::new(array);
-        let _ = input.complete();
+        let mut input = array;
         match terminated(gemma4_array_content, eof).parse_next(&mut input) {
             Ok(value) => Ok(value),
             Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 array")),
@@ -367,18 +494,18 @@ mod tests {
         ]
     }
 
-    fn collect_stream(chunks: &[&str]) -> ToolParseResult {
+    fn collect_stream(chunks: &[&str]) -> ToolParserOutput {
         let mut parser = Gemma4ToolParser::new(&test_tools());
-        let mut result = ToolParseResult::default();
+        let mut output = ToolParserOutput::default();
         for chunk in chunks {
-            result.append(parser.push(chunk).unwrap());
+            output.append(parser.parse_chunk(chunk).unwrap());
         }
-        result.append(parser.finish().unwrap());
-        result.coalesce_calls()
+        output.append(parser.finish().unwrap());
+        output.coalesce_calls()
     }
 
-    fn first_call(result: &ToolParseResult) -> &ToolCallDelta {
-        result.calls.first().expect("expected one tool call")
+    fn first_call(output: &ToolParserOutput) -> &ToolCallDelta {
+        output.calls.first().expect("expected one tool call")
     }
 
     #[test]
@@ -416,15 +543,15 @@ mod tests {
     #[test]
     fn gemma4_parse_complete_extracts_single_tool_call() {
         let mut parser = Gemma4ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}<tool_call|>")
             .unwrap();
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
+        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::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "location": "London" })
         );
     }
@@ -441,7 +568,7 @@ mod tests {
 
     #[test]
     fn gemma4_streaming_basic_single_tool_call() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:get_weather{",
             "location:<|\"|>Paris",
@@ -450,17 +577,17 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
+        assert!(output.normal_text.is_empty());
+        assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "location": "Paris, France" })
         );
     }
 
     #[test]
     fn gemma4_streaming_text_before_and_after_tool_call() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "Let me check ",
             "the weather. ",
             "<|tool_call>",
@@ -470,10 +597,10 @@ mod tests {
             "div>",
         ]);
 
-        assert_eq!(result.normal_text, "Let me check the weather. <div>");
-        assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
+        assert_eq!(output.normal_text, "Let me check the weather. <div>");
+        assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "location": "London" })
         );
     }
@@ -481,68 +608,68 @@ mod tests {
     #[test]
     fn gemma4_streaming_waits_for_complete_tool_call() {
         let mut parser = Gemma4ToolParser::new(&test_tools());
-        let mut result = ToolParseResult::default();
+        let mut output = ToolParserOutput::default();
 
         for chunk in [
             "<|tool_call>",
             "call:get_weather{",
             "location:<|\"|>Paris<|\"|>}",
         ] {
-            result.append(parser.push(chunk).unwrap());
-            assert!(result.calls.is_empty());
+            output.append(parser.parse_chunk(chunk).unwrap());
+            assert!(output.calls.is_empty());
         }
 
-        result.append(parser.push("<tool_call|>").unwrap());
-        let result = result.coalesce_calls();
+        output.append(parser.parse_chunk("<tool_call|>").unwrap());
+        let output = output.coalesce_calls();
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "location": "Paris" })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_boolean_split_across_chunks() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:search{input:{all:tru",
             "e}}",
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("search"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("search"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "input": { "all": true } })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_false_split_across_chunks() {
-        let result = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", "<tool_call|>"]);
+        let output = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", "<tool_call|>"]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("set"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("set"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "flag": false })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_number_split_across_chunks() {
-        let result = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", "<tool_call|>"]);
+        let output = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", "<tool_call|>"]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("set"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("set"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "count": 42 })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_split_string_delimiter() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:todowrite{",
             "content:<|\"|>Buy milk<|",
@@ -550,17 +677,32 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("todowrite"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("todowrite"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "content": "Buy milk" })
         );
-        assert!(!first_call(&result).arguments.contains("<|"));
+        assert!(!first_call(&output).arguments.contains("<|"));
+    }
+
+    #[test]
+    fn gemma4_streaming_handles_split_tool_call_end_marker() {
+        let output = collect_stream(&[
+            "<|tool_call>",
+            "call:get_weather{location:<|\"|>Paris<|\"|>}<tool",
+            "_call|>",
+        ]);
+
+        assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
+        assert_eq!(
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
+            json!({ "location": "Paris" })
+        );
     }
 
     #[test]
     fn gemma4_streaming_handles_end_marker_literal_inside_string() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:todowrite{",
             "content:<|\"|>literal }<tool_call|> inside",
@@ -568,16 +710,16 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("todowrite"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("todowrite"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "content": "literal }<tool_call|> inside" })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_html_argument_without_duplication() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:write_file{",
             "path:<|\"|>index.html<|\"|>,",
@@ -590,9 +732,9 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("write_file"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("write_file"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({
                 "path": "index.html",
                 "content": "<!DOCTYPE html>\n<html lang=\"zh-CN\">\n<head>\n    <meta charset=\"UTF-8\">\n    <meta name=\"viewport\" content=\"width=device-width\">\n",
@@ -602,7 +744,7 @@ mod tests {
 
     #[test]
     fn gemma4_streaming_trailing_bare_bool_is_not_duplicated() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:Edit{",
             "file_path:<|\"|>src/env.py<|\"|>,",
@@ -613,9 +755,9 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("Edit"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("Edit"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({
                 "file_path": "src/env.py",
                 "old_string": "old_val",
@@ -624,7 +766,7 @@ mod tests {
             })
         );
         assert_eq!(
-            first_call(&result).arguments.matches("replace_all").count(),
+            first_call(&output).arguments.matches("replace_all").count(),
             1
         );
     }
@@ -632,22 +774,53 @@ mod tests {
     #[test]
     fn gemma4_finish_flushes_partial_start_marker_as_text() {
         let mut parser = Gemma4ToolParser::new(&test_tools());
-        let mut result = parser.push("<").unwrap();
-        result.append(parser.finish().unwrap());
+        let mut output = parser.parse_chunk("<").unwrap();
+        output.append(parser.finish().unwrap());
 
-        assert_eq!(result.normal_text, "<");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "<");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn gemma4_finish_rejects_complete_args_without_end_marker() {
         let mut parser = Gemma4ToolParser::new(&test_tools());
         for chunk in ["<|tool_call>", "call:get_status{}"] {
-            parser.push(chunk).unwrap();
+            parser.parse_chunk(chunk).unwrap();
         }
 
         let error = parser.finish().unwrap_err();
 
         assert!(error.to_report_string().contains("incomplete Gemma4 tool call"));
     }
+
+    #[test]
+    fn gemma4_reset_preserves_internally_buffered_arguments() {
+        let mut parser = Gemma4ToolParser::new(&test_tools());
+        for chunk in [
+            "<|tool_call>",
+            "call:write_file{",
+            "content:<|\"|>hello ",
+            "world<|\"|>",
+        ] {
+            parser.parse_chunk(chunk).unwrap();
+        }
+
+        let raw = parser.reset();
+
+        assert_eq!(
+            raw,
+            "<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>"
+        );
+    }
+
+    #[test]
+    fn gemma4_reset_preserves_completed_arguments_after_parse_error() {
+        let mut parser = Gemma4ToolParser::new(&test_tools());
+        let input = "<|tool_call>call:set{broken}<tool_call|>";
+
+        let _error = parser.parse_chunk(input).unwrap_err();
+        let raw = parser.reset();
+
+        assert_eq!(raw, input);
+    }
 }

rust/src/tool-parser/src/glm_xml/glm45_moe.rs

@@ -1,5 +1,5 @@
 use super::{GlmXmlToolParser, Separator};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for GLM-4.5/4.6 MoE XML-style tool calls.
 ///
@@ -23,7 +23,6 @@ impl Glm45MoeToolParser {
 }
 
 impl ToolParser for Glm45MoeToolParser {
-    /// Create a boxed GLM-4.5/4.6 MoE tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -31,13 +30,15 @@ impl ToolParser for Glm45MoeToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the GLM MoE parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.0.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.0.parse_into(chunk, output)
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.0.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.0.reset()
+    }
 }

rust/src/tool-parser/src/glm_xml/glm47_moe.rs

@@ -1,5 +1,5 @@
 use super::{GlmXmlToolParser, Separator};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 /// Tool parser for GLM-4.7 MoE XML-style tool calls.
 ///
@@ -22,20 +22,25 @@ impl ToolParser for Glm47MoeToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.0.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.0.parse_into(chunk, output)
     }
 
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.0.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.0.reset()
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use serde_json::{Value, json};
 
-    use super::{Glm47MoeToolParser, ToolParser};
+    use super::Glm47MoeToolParser;
+    use crate::ToolParserTestExt as _;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
 
     fn glm47_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
@@ -58,13 +63,13 @@ mod tests {
             )
         );
 
-        let result = parser.parse_complete(&output).unwrap();
+        let output = parser.parse_complete(&output).unwrap();
 
-        assert_eq!(result.normal_text, "Let me search for that.\n");
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({"city": "Beijing", "date": "2024-12-25"})
         );
     }
@@ -79,14 +84,14 @@ mod tests {
         );
 
         let chunks = split_by_chars(&output, 7);
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
-        assert_eq!(result.normal_text, "");
-        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!(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::<Value>(&result.calls[1].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
             json!({"x": 1, "y": 2})
         );
     }
@@ -94,7 +99,7 @@ mod tests {
     #[test]
     fn glm47_parse_complete_converts_schema_types() {
         let mut parser = Glm47MoeToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&glm47_tool_call(
                 "convert",
                 &[
@@ -108,7 +113,7 @@ mod tests {
             .unwrap();
 
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "whole": 42,
                 "flag": true,
@@ -123,12 +128,12 @@ mod tests {
     fn glm47_parse_complete_extracts_zero_argument_call() {
         let mut parser = Glm47MoeToolParser::new(&test_tools());
 
-        let result = parser.parse_complete("<tool_call>add</tool_call>").unwrap();
+        let output = parser.parse_complete("<tool_call>add</tool_call>").unwrap();
 
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({})
         );
     }

rust/src/tool-parser/src/glm_xml/mod.rs

@@ -6,7 +6,7 @@ use winnow::token::{literal, rest, take_until, take_while};
 
 use super::parameters::ToolSchemas;
 use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
-use super::{Result, ToolCallDelta, ToolParseResult};
+use super::{Result, ToolCallDelta, ToolParserOutput};
 use crate::Tool;
 
 mod glm45_moe;
@@ -76,10 +76,10 @@ impl GlmXmlToolParser {
     }
 
     /// Apply one parsed GLM event to parser state and output.
-    fn apply_event(&mut self, event: GlmEvent, result: &mut ToolParseResult) -> Result<()> {
+    fn apply_event(&mut self, event: GlmEvent, output: &mut ToolParserOutput) -> Result<()> {
         match event {
             GlmEvent::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
             }
             GlmEvent::ToolCallStart => self.mode = GlmMode::ToolCall,
             GlmEvent::ToolCall { name, raw_params } => {
@@ -88,7 +88,7 @@ impl GlmXmlToolParser {
                 let arguments = serde_json::to_string(&arguments)
                     .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
 
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index: self.emitted_tool_count,
                     name: Some(name),
                     arguments,
@@ -100,40 +100,36 @@ impl GlmXmlToolParser {
         Ok(())
     }
 
-    /// Reset all streaming state.
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.mode = GlmMode::Text;
         self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
     }
 
-    /// Push one decoded text chunk through the GLM MoE parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
         while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
             parse_next_glm_event(input, self.mode, self.separator)
         })? {
-            self.apply_event(event, &mut result)?;
+            self.apply_event(event, output)?;
             self.buffer.drain(..consumed_len);
         }
 
-        Ok(result)
+        Ok(())
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
-        let mut result = ToolParseResult::default();
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
         if !self.buffer.is_empty() {
             match self.mode {
-                GlmMode::Text => result.normal_text.push_str(&self.buffer),
+                GlmMode::Text => output.normal_text.push_str(&self.buffer),
                 GlmMode::ToolCall => return Err(parsing_failed!("incomplete GLM MoE tool call")),
                 GlmMode::AfterToolCall => {}
             }
         }
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
     }
 }
 
@@ -256,8 +252,8 @@ mod tests {
     use thiserror_ext::AsReport;
 
     use super::Glm45MoeToolParser;
-    use crate::ToolParser;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
+    use crate::{ToolParser, ToolParserTestExt as _};
 
     fn glm45_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
         let params = params
@@ -273,10 +269,10 @@ mod tests {
     #[test]
     fn glm45_parse_complete_without_tool_call_keeps_text() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
-        let result = parser.parse_complete("Hello, world!").unwrap();
+        let output = parser.parse_complete("Hello, world!").unwrap();
 
-        assert_eq!(result.normal_text, "Hello, world!");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "Hello, world!");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
@@ -290,13 +286,13 @@ mod tests {
             )
         );
 
-        let result = parser.parse_complete(&output).unwrap();
+        let output = parser.parse_complete(&output).unwrap();
 
-        assert_eq!(result.normal_text, "Let me search for that.\n");
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({"city": "Beijing", "date": "2024-12-25"})
         );
     }
@@ -311,14 +307,14 @@ mod tests {
         );
 
         let chunks = split_by_chars(&output, 11);
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
-        assert_eq!(result.normal_text, "");
-        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!(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::<Value>(&result.calls[1].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
             json!({"x": 1, "y": 2})
         );
     }
@@ -326,7 +322,7 @@ mod tests {
     #[test]
     fn glm45_parse_complete_unescapes_literal_closing_tags_in_arg_value() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&glm45_tool_call(
                 "get_weather",
                 &[
@@ -337,7 +333,7 @@ mod tests {
             .unwrap();
 
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "city": "Paris </arg_value></tool_call>",
                 "date": "2026-05-08",
@@ -349,17 +345,17 @@ mod tests {
     fn glm45_streaming_without_tool_call_emits_text_incrementally() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
 
-        let result = collect_stream(&mut parser, &["hello ", "world"]);
+        let output = collect_stream(&mut parser, &["hello ", "world"]);
 
-        assert_eq!(result.normal_text, "hello world");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "hello world");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn glm45_streaming_preserves_prefix_text() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
 
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "Prefix ",
@@ -367,14 +363,14 @@ mod tests {
             ],
         );
 
-        assert_eq!(result.normal_text, "Prefix ");
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.normal_text, "Prefix ");
+        assert_eq!(output.calls.len(), 1);
     }
 
     #[test]
     fn glm45_streaming_handles_start_token_split_across_chunks() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "hello <tool",
@@ -383,26 +379,26 @@ mod tests {
             ],
         );
 
-        assert_eq!(result.normal_text, "hello ");
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.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]
     fn glm45_streaming_does_not_emit_incomplete_tool_call() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
 
-        let result = parser.push("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
+        let output = parser.parse_chunk("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
 
-        assert_eq!(result.normal_text, "");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn glm45_finish_fails_incomplete_tool_call() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
 
-        parser.push("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
+        parser.parse_chunk("<tool_call>get_weather\n<arg_key>city</arg_key>").unwrap();
         let error = parser.finish().unwrap_err();
 
         assert!(error.as_report().to_string().contains("incomplete GLM MoE tool call"));
@@ -412,7 +408,7 @@ mod tests {
     fn glm45_malformed_tool_call_fails_fast() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
 
-        let error = parser.push("<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value></tool_call>").unwrap_err();
+        let error = parser.parse_chunk("<tool_call>get_weather<arg_key>city</arg_key><arg_value>Paris</arg_value></tool_call>").unwrap_err();
 
         assert!(error.as_report().to_string().contains("tool parser parsing failed"));
     }
@@ -421,7 +417,7 @@ mod tests {
     fn glm45_streaming_ignores_trailing_text_after_tool_calls() {
         let mut parser = Glm45MoeToolParser::new(&test_tools());
 
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[&format!(
                 "{}<|endoftext|>",
@@ -429,7 +425,7 @@ mod tests {
             )],
         );
 
-        assert_eq!(result.normal_text, "");
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.normal_text, "");
+        assert_eq!(output.calls.len(), 1);
     }
 }

rust/src/tool-parser/src/hy_v3.rs

@@ -0,0 +1,559 @@
+use winnow::ascii::multispace0 as ws0;
+use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated};
+use winnow::prelude::*;
+use winnow::stream::Partial;
+use winnow::token::{literal, rest, take_until};
+
+use super::parameters::ToolSchemas;
+use super::utils::{parse_buffered_event, safe_text_len};
+use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
+use crate::Tool;
+
+const TOOL_CALLS_START: &str = "<tool_calls>";
+const TOOL_CALLS_END: &str = "</tool_calls>";
+const TOOL_CALL_START: &str = "<tool_call>";
+const TOOL_CALL_END: &str = "</tool_call>";
+const TOOL_SEP: &str = "<tool_sep>";
+const ARG_KEY_START: &str = "<arg_key>";
+const ARG_KEY_END: &str = "</arg_key>";
+const ARG_VALUE_START: &str = "<arg_value>";
+const ARG_VALUE_END: &str = "</arg_value>";
+
+type HyV3Input<'i> = Partial<&'i str>;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum HyV3Mode {
+    Text,
+    ToolBlock,
+    Done,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum HyV3Event {
+    Text {
+        len: usize,
+    },
+    ToolBlockStart,
+    ToolCall {
+        name: String,
+        raw_params: Vec<(String, String)>,
+    },
+    ToolBlockEnd,
+    IgnoredRest,
+}
+
+/// Tool parser for HY3 XML-style tool calls.
+///
+/// Example tool call content:
+///
+/// ```text
+/// <tool_calls>
+/// <tool_call>get_weather<tool_sep>
+/// <arg_key>city</arg_key><arg_value>Beijing</arg_value>
+/// </tool_call>
+/// </tool_calls>
+/// ```
+///
+/// Arguments are emitted only after a full `<tool_call>` block is parsed.
+/// HY3 marker tokens are added-vocabulary tokens rather than tokenizer special
+/// tokens, so the default `preserve_special_tokens() == false` is sufficient.
+pub struct HyV3ToolParser {
+    buffer: String,
+    mode: HyV3Mode,
+    emitted_tool_count: usize,
+    tool_parameters: ToolSchemas,
+}
+
+impl HyV3ToolParser {
+    /// Create a HY3 tool parser.
+    fn new(tools: &[Tool]) -> Self {
+        Self {
+            buffer: String::new(),
+            mode: HyV3Mode::Text,
+            emitted_tool_count: 0,
+            tool_parameters: ToolSchemas::from_tools(tools),
+        }
+    }
+
+    /// Apply one parsed HY3 event to parser state and output.
+    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]);
+            }
+            HyV3Event::ToolBlockStart => self.mode = HyV3Mode::ToolBlock,
+            HyV3Event::ToolCall { name, raw_params } => {
+                let arguments = self.tool_parameters.convert_params_with_schema(&name, raw_params);
+                let arguments = serde_json::to_string(&arguments)
+                    .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
+
+                output.calls.push(ToolCallDelta {
+                    tool_index: self.emitted_tool_count,
+                    name: Some(name),
+                    arguments,
+                });
+                self.emitted_tool_count += 1;
+            }
+            HyV3Event::ToolBlockEnd => self.mode = HyV3Mode::Done,
+            HyV3Event::IgnoredRest => {}
+        }
+        Ok(())
+    }
+}
+
+impl ToolParser for HyV3ToolParser {
+    fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
+    where
+        Self: Sized + 'static,
+    {
+        Ok(Box::new(Self::new(tools)))
+    }
+
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.buffer.push_str(chunk);
+
+        while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
+            parse_next_hy_v3_event(input, self.mode)
+        })? {
+            self.apply_event(event, output)?;
+            self.buffer.drain(..consumed_len);
+        }
+
+        Ok(())
+    }
+
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
+        match self.mode {
+            HyV3Mode::Text => output.normal_text.push_str(&self.buffer),
+            HyV3Mode::ToolBlock => return Err(parsing_failed!("incomplete HY3 tool call")),
+            HyV3Mode::Done => {}
+        }
+        let _ = self.reset();
+        Ok(output)
+    }
+
+    fn reset(&mut self) -> String {
+        self.mode = HyV3Mode::Text;
+        self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
+    }
+}
+
+/// Parse a HY3 event for the current parser mode.
+fn parse_next_hy_v3_event(input: &mut HyV3Input<'_>, mode: HyV3Mode) -> ModalResult<HyV3Event> {
+    match mode {
+        HyV3Mode::Text => parse_text_event(input),
+        HyV3Mode::ToolBlock => parse_tool_block_event(input),
+        HyV3Mode::Done => ignored_rest_event(input),
+    }
+}
+
+/// Parse a text-mode HY3 event.
+fn parse_text_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
+    alt((tool_block_start_event, safe_text_event)).parse_next(input)
+}
+
+/// Parse a HY3 tool-block start marker.
+fn tool_block_start_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
+    literal(TOOL_CALLS_START).value(HyV3Event::ToolBlockStart).parse_next(input)
+}
+
+/// Parse a safe text run before the next HY3 marker.
+fn safe_text_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
+    safe_text_len(input, TOOL_CALLS_START).map(|len| HyV3Event::Text { len })
+}
+
+/// Parse one event inside a HY3 tool block.
+fn parse_tool_block_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
+    alt((tool_block_end_event, tool_call_event)).parse_next(input)
+}
+
+/// Parse a HY3 tool-block end marker.
+fn tool_block_end_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
+    (ws0, literal(TOOL_CALLS_END)).value(HyV3Event::ToolBlockEnd).parse_next(input)
+}
+
+/// Parse a complete HY3 tool-call block.
+fn tool_call_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
+    let (name, body) = seq!(
+        _: ws0,
+        _: literal(TOOL_CALL_START),
+        take_until(0.., TOOL_SEP),
+        _: literal(TOOL_SEP),
+        take_until(0.., TOOL_CALL_END),
+        _: literal(TOOL_CALL_END),
+    )
+    .parse_next(input)?;
+    let raw_params = parse_tool_call_params(body)?;
+
+    Ok(HyV3Event::ToolCall {
+        name: name.trim().to_string(),
+        raw_params,
+    })
+}
+
+/// Parse all parameter blocks inside a complete HY3 tool call.
+fn parse_tool_call_params(tool_call_body: &str) -> ModalResult<Vec<(String, String)>> {
+    let mut input = tool_call_body;
+    delimited(ws0, repeat(0.., terminated(parameter, ws0)), eof).parse_next(&mut input)
+}
+
+/// Parse a HY3 argument key/value block.
+fn parameter(input: &mut &str) -> ModalResult<(String, String)> {
+    let (name, value) = seq!(
+        _: literal(ARG_KEY_START),
+        take_until(0.., ARG_KEY_END),
+        _: literal(ARG_KEY_END),
+        _: ws0,
+        _: literal(ARG_VALUE_START),
+        take_until(0.., ARG_VALUE_END),
+        _: literal(ARG_VALUE_END),
+    )
+    .parse_next(input)?;
+
+    Ok((name.trim().to_string(), value.to_string()))
+}
+
+/// Parse ignored rest after the HY3 tool block ends.
+fn ignored_rest_event(input: &mut HyV3Input<'_>) -> ModalResult<HyV3Event> {
+    rest.value(HyV3Event::IgnoredRest).parse_next(input)
+}
+
+#[cfg(test)]
+mod tests {
+    use expect_test::expect;
+    use serde_json::{Value, json};
+    use thiserror_ext::AsReport;
+
+    use super::{HyV3ToolParser, ToolParser};
+    use crate::test_utils::{collect_stream, split_by_chars, test_tools};
+    use crate::{ToolParserOutput, ToolParserTestExt as _};
+
+    fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
+        let params = params
+            .iter()
+            .map(|(name, value)| format!("<arg_key>{name}</arg_key><arg_value>{value}</arg_value>"))
+            .collect::<Vec<_>>()
+            .join("\n");
+        format!("<tool_call>{function_name}<tool_sep>{params}</tool_call>")
+    }
+
+    fn build_tool_calls(tool_calls: &[String]) -> String {
+        format!("<tool_calls>\n{}\n</tool_calls>", tool_calls.join("\n"))
+    }
+
+    fn parsed_arguments(output: &ToolParserOutput, index: usize) -> Value {
+        serde_json::from_str(&output.calls[index].arguments).unwrap()
+    }
+
+    #[test]
+    fn hy_v3_does_not_preserve_special_tokens() {
+        let parser = HyV3ToolParser::new(&test_tools());
+
+        assert!(!parser.preserve_special_tokens());
+    }
+
+    #[test]
+    fn hy_v3_parse_complete_without_tool_call_keeps_text() {
+        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());
+    }
+
+    #[test]
+    fn hy_v3_parse_complete_extracts_zero_arg_inline_tool_call() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let output = parser
+            .parse_complete(
+                "<tool_calls><tool_call>get_current_date<tool_sep></tool_call></tool_calls>",
+            )
+            .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!(parsed_arguments(&output, 0), json!({}));
+    }
+
+    #[test]
+    fn hy_v3_parse_complete_extracts_zero_arg_newline_tool_call() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let output = parser
+            .parse_complete(
+                "<tool_calls>\n<tool_call>get_current_date<tool_sep>\n</tool_call>\n</tool_calls>",
+            )
+            .unwrap();
+
+        assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
+        assert_eq!(parsed_arguments(&output, 0), json!({}));
+    }
+
+    #[test]
+    fn hy_v3_parse_complete_extracts_arguments_on_same_line() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let output = parser
+            .parse_complete(
+                "<tool_calls><tool_call>get_weather<tool_sep><arg_key>city</arg_key><arg_value>Beijing</arg_value><arg_key>date</arg_key><arg_value>2026-03-30</arg_value></tool_call></tool_calls>",
+            )
+            .unwrap();
+
+        assert_eq!(
+            parsed_arguments(&output, 0),
+            json!({ "city": "Beijing", "date": "2026-03-30" })
+        );
+    }
+
+    #[test]
+    fn hy_v3_parse_complete_extracts_arguments_with_newlines() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let output = parser
+            .parse_complete(&build_tool_calls(&[build_tool_call(
+                "get_weather",
+                &[("city", "Beijing"), ("date", "2026-03-30")],
+            )]))
+            .unwrap();
+
+        assert_eq!(
+            parsed_arguments(&output, 0),
+            json!({ "city": "Beijing", "date": "2026-03-30" })
+        );
+    }
+
+    #[test]
+    fn hy_v3_parse_complete_preserves_prefix_and_ignores_trailing_text() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let output = parser
+            .parse_complete(&format!(
+                "Checking.{} trailing text",
+                build_tool_calls(&[build_tool_call("get_current_date", &[])])
+            ))
+            .unwrap();
+
+        assert_eq!(output.normal_text, "Checking.");
+        assert_eq!(output.calls[0].name.as_deref(), Some("get_current_date"));
+    }
+
+    #[test]
+    fn hy_v3_parse_complete_extracts_multiple_tool_calls_in_one_block() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let output = parser
+            .parse_complete(&build_tool_calls(&[
+                build_tool_call(
+                    "get_weather",
+                    &[("city", "Beijing"), ("date", "2026-03-30")],
+                ),
+                build_tool_call(
+                    "get_weather",
+                    &[("city", "Hangzhou"), ("date", "2026-03-30")],
+                ),
+            ]))
+            .unwrap();
+
+        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\"}",
+                    },
+                ],
+            }
+        "#]]
+        .assert_debug_eq(&output);
+    }
+
+    #[test]
+    fn hy_v3_parse_complete_converts_schema_types() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let output = parser
+            .parse_complete(&build_tool_calls(&[build_tool_call(
+                "convert",
+                &[
+                    ("whole", "5.3"),
+                    ("flag", "true"),
+                    ("payload", r#"{"k":1}"#),
+                    ("items", "[1,2]"),
+                ],
+            )]))
+            .unwrap();
+
+        assert_eq!(
+            parsed_arguments(&output, 0),
+            json!({
+                "whole": 5.3,
+                "flag": true,
+                "payload": { "k": 1 },
+                "items": [1, 2]
+            })
+        );
+    }
+
+    #[test]
+    fn hy_v3_streaming_without_tool_call_emits_text_incrementally() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let mut output = ToolParserOutput::default();
+
+        output.append(parser.parse_chunk("This is ").unwrap());
+        output.append(parser.parse_chunk("a plain ").unwrap());
+        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());
+    }
+
+    #[test]
+    fn hy_v3_streaming_extracts_zero_arg_tool_call() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let chunks = [
+            "<tool_calls>",
+            "\n<tool_call>",
+            "get_current_date",
+            "<tool_sep>",
+            "\n</tool_call>",
+            "\n</tool_calls>",
+        ];
+
+        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!(parsed_arguments(&output, 0), json!({}));
+    }
+
+    #[test]
+    fn hy_v3_streaming_extracts_arguments() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let chunks = [
+            "<tool_calls>",
+            "\n<tool_call>",
+            "get_weather",
+            "<tool_sep>",
+            "\n<arg_key>city</arg_key>",
+            "\n<arg_value>Beijing</arg_value>",
+            "\n<arg_key>date</arg_key>",
+            "\n<arg_value>2026-03-30</arg_value>",
+            "\n</tool_call>",
+            "\n</tool_calls>",
+        ];
+
+        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!(
+            parsed_arguments(&output, 0),
+            json!({ "city": "Beijing", "date": "2026-03-30" })
+        );
+    }
+
+    #[test]
+    fn hy_v3_streaming_preserves_prefix_text() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let chunks = [
+            "Checking.",
+            "<tool_calls>",
+            "\n<tool_call>",
+            "get_current_date",
+            "<tool_sep>",
+            "\n</tool_call>",
+            "\n</tool_calls>",
+        ];
+
+        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"));
+    }
+
+    #[test]
+    fn hy_v3_streaming_extracts_multiple_tool_calls_in_one_block() {
+        let input = build_tool_calls(&[
+            build_tool_call(
+                "get_weather",
+                &[("city", "Beijing"), ("date", "2026-03-30")],
+            ),
+            build_tool_call(
+                "get_weather",
+                &[("city", "Hangzhou"), ("date", "2026-03-30")],
+            ),
+        ]);
+        let chunks = split_by_chars(&input, 9);
+        let mut parser = HyV3ToolParser::new(&test_tools());
+
+        let output = collect_stream(&mut parser, &chunks);
+
+        assert_eq!(output.calls.len(), 2);
+        assert_eq!(parsed_arguments(&output, 0)["city"], json!("Beijing"));
+        assert_eq!(parsed_arguments(&output, 1)["city"], json!("Hangzhou"));
+    }
+
+    #[test]
+    fn hy_v3_streaming_handles_start_marker_split_across_chunks() {
+        let input = format!(
+            "hello {}",
+            build_tool_calls(&[build_tool_call("get_weather", &[("city", "Beijing")])])
+        );
+        let chunks = split_by_chars(&input, 5);
+        let mut parser = HyV3ToolParser::new(&test_tools());
+
+        let output = collect_stream(&mut parser, &chunks);
+
+        assert_eq!(output.normal_text, "hello ");
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(parsed_arguments(&output, 0), json!({ "city": "Beijing" }));
+    }
+
+    #[test]
+    fn hy_v3_streaming_does_not_emit_incomplete_tool_call() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let mut output = ToolParserOutput::default();
+
+        parser
+            .parse_into(
+                "<tool_calls><tool_call>get_weather<tool_sep><arg_key>city</arg_key><arg_value>Bei",
+                &mut output,
+            )
+            .unwrap();
+
+        assert_eq!(output.normal_text, "");
+        assert!(output.calls.is_empty());
+    }
+
+    #[test]
+    fn hy_v3_finish_fails_incomplete_tool_call() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        parser.parse_chunk("<tool_calls><tool_call>get_weather<tool_sep>").unwrap();
+
+        let error = parser.finish().unwrap_err();
+
+        expect!["tool parser parsing failed: incomplete HY3 tool call"]
+            .assert_eq(&error.to_report_string());
+    }
+
+    #[test]
+    fn hy_v3_malformed_tool_call_fails_fast() {
+        let mut parser = HyV3ToolParser::new(&test_tools());
+        let error = parser
+            .parse_complete(
+                "<tool_calls><tool_call>get_weather<tool_sep><arg_key>city</arg_key><arg_value>Beijing</tool_call></tool_calls>",
+            )
+            .unwrap_err();
+
+        assert!(error.to_report_string().starts_with("tool parser parsing failed:"));
+    }
+}