test: touch triton_reshape_and_cache_flash.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-amd.yaml

@@ -2703,19 +2703,35 @@ steps:
   optional: true
   working_dir: "/vllm-workspace/"
   source_file_dependencies:
+  - csrc/custom_quickreduce.cu
+  - csrc/ops.h
+  - csrc/torch_bindings.cpp
   - vllm/distributed/
-  - vllm/v1/distributed/
+  - vllm/model_executor/layers/
+  - vllm/entrypoints/llm.py
+  - vllm/config/parallel.py
   - vllm/model_executor/layers/fused_moe/
+  - vllm/v1/engine/
+  - vllm/v1/executor/
+  - vllm/v1/worker/
+  - vllm/v1/distributed/
   - vllm/v1/attention/backends/
   - vllm/v1/attention/selector.py
-  - tests/distributed/test_context_parallel.py
-  - tests/v1/distributed/test_dbo.py
-  - examples/features/data_parallel/data_parallel_offline.py
   - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
   - vllm/platforms/rocm.py
+  - vllm/envs.py
+  - examples/offline_inference/data_parallel.py
+  - tests/distributed/test_context_parallel.py
+  - tests/distributed/test_rocm_quick_reduce.py
+  - tests/distributed/test_quick_all_reduce.py
+  - tests/v1/distributed/test_dbo.py
+  - tests/utils.py
   commands:
   - pytest -v -s tests/distributed/test_context_parallel.py
   - pytest -v -s tests/v1/distributed/test_dbo.py
+  - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
+  - pytest -v -s tests/distributed/test_quick_all_reduce.py
 
 #--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#
 

.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

@@ -38,6 +38,28 @@ steps:
   commands:
     - pytest -v -s kernels/core/test_minimax_reduce_rms.py
 
+- label: Deepseek V4 Kernel Test (H100)
+  key: deepseek-v4-kernel-test-h100
+  timeout_in_minutes: 15
+  device: h100
+  source_file_dependencies:
+  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+  - vllm/models/deepseek_v4/common/ops/
+  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  commands:
+    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+
+- label: Deepseek V4 Kernel Test (B200)
+  key: deepseek-v4-kernel-test-b200
+  timeout_in_minutes: 15
+  device: b200-k8s
+  source_file_dependencies:
+  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+  - vllm/models/deepseek_v4/common/ops/
+  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  commands:
+    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+
 - label: Kernels Attention Test %N
   key: kernels-attention-test
   timeout_in_minutes: 35
@@ -64,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:

AGENTS.md

@@ -101,6 +101,8 @@ pre-commit run ruff-check --all-files
 pre-commit run mypy-3.10 --all-files --hook-stage manual
 ```
 
+The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check.
+
 ### Commit messages
 
 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:

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/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -122,18 +122,45 @@ __device__ __forceinline__ float warpSum(float val) {
 // Per-slot inner pipeline
 // ────────────────────────────────────────────────────────────────────────────
 // Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per
-// token
-template <typename scalar_t_in>
+// token.  Templated on `kNumHeadsQPadded` so the KV-sentinel comparison and
+// q_out stride fold to compile-time constants.
+//
+// Slot layout (per token):
+//   slot < num_heads_q                          → live-Q   (RMSNorm + RoPE,
+//                                                           read q_in →
+//                                                           write q_out)
+//   num_heads_q <= slot < kNumHeadsQPadded      → pad-Q    (zero-fill q_out;
+//                                                           v0/v1 unused)
+//   slot == kNumHeadsQPadded                    → KV       (RoPE + UE8M0 quant
+//                                                           + paged-cache
+//                                                           insert)
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __device__ __forceinline__ void processDeepseekV4Slot(
     uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
     int const dim_base, int const laneId, int const num_heads_q,
-    float const eps, scalar_t_in* __restrict__ q_inout,
+    float const eps, scalar_t_in* __restrict__ q_out,
     uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
     int64_t const* __restrict__ position_ids,
     float const* __restrict__ cos_sin_cache, int const cache_block_size,
     int const kv_block_stride) {
   using Converter = vllm::_typeConvert<scalar_t_in>;
-  bool const isKV = (slotIdx == num_heads_q);
+  bool const isKV = (slotIdx == kNumHeadsQPadded);
+  bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
+
+  // ── Pad-Q branch: write 32 B of zeros and exit. ─────────────────────────
+  // FlashMLA reads these slots; bf16 +0.0 is bit pattern 0x0000, so a uint4
+  // zero literal is correct.  Matches the live-Q branch's vectorized store.
+  if (isPadQ) {
+    scalar_t_in* dst =
+        q_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
+        dim_base;
+    uint4 const zero4 = {0u, 0u, 0u, 0u};
+    *reinterpret_cast<uint4*>(dst) = zero4;
+    *reinterpret_cast<uint4*>(dst + 8) = zero4;
+    return;
+  }
 
   // ── Decode the bf16 → 16 fp32 registers ─────────────────────────────
   float elements[kElemsPerLane];
@@ -207,7 +234,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
   // triggering and per-iteration buffer rotation.
   // ═══════════════════════════════════════════════════════════════════
   if (!isKV) {
-    // ── Q: cast back to bf16 and store. ────────────────────────────
+    // ── Live-Q: cast back to bf16 and store into the padded q_out. ─────
     uint4 out0, out1;
     typename Converter::packed_hip_type* po0 =
         reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
@@ -224,8 +251,9 @@ __device__ __forceinline__ void processDeepseekV4Slot(
           make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
     }
     scalar_t_in* dst =
-        q_inout +
-        (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
+        q_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
         dim_base;
     *reinterpret_cast<uint4*>(dst) = out0;
     *reinterpret_cast<uint4*>(dst + 8) = out1;
@@ -313,20 +341,32 @@ __device__ __forceinline__ void processDeepseekV4Slot(
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
 //
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
+// Grid: 1D, gridDim.x = ceil(num_tokens_full * (kNumHeadsQPadded + 1) /
 // warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
-// warp handles one (token, head_slot) pair. head_slot < num_heads_q          →
-// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q         → KV
-// branch (RoPE + UE8M0 quant + insert)
+// warp handles one (token, head_slot) pair.
+//   slot < num_heads_q                              → live-Q branch
+//                                                     (RMSNorm + RoPE,
+//                                                      read q_in → write q_out)
+//   num_heads_q <= slot < kNumHeadsQPadded          → pad-Q branch
+//                                                     (zero-fill q_out)
+//   slot == kNumHeadsQPadded                        → KV branch
+//                                                     (RoPE + UE8M0 quant +
+//                                                      paged-cache insert)
+//
+// `kNumHeadsQPadded` is a template parameter (compile-time constant) so the
+// divisions in the grid math and the KV-sentinel comparison fold to fast
+// constant operations.  The launch wrapper dispatches the runtime value to
+// the matching instantiation.
 //
 // With DP padding, q/kv/position_ids can have more rows than slot_mapping.
-// The Q branch covers all `num_tokens_full` rows (downstream attention uses
-// them).  The KV branch only inserts the first `num_tokens_insert` tokens
-// (= slot_mapping length) into the paged cache.
+// The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream
+// attention uses them).  The KV branch only inserts the first
+// `num_tokens_insert` tokens (= slot_mapping length) into the paged cache.
 //
-template <typename scalar_t_in>
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
-    scalar_t_in* __restrict__ q_inout,         // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ q_in,      // [N, num_heads_q,      512]
+    scalar_t_in* __restrict__ q_out,           // [N, kNumHeadsQPadded, 512]
     scalar_t_in const* __restrict__ kv_in,     // [N, 512] bf16
     uint8_t* __restrict__ k_cache,             // [num_blocks, block_stride]
     int64_t const* __restrict__ slot_mapping,  // [num_tokens_insert] i64
@@ -335,7 +375,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     float const eps,
     int const num_tokens_full,    // = q.size(0) = kv.size(0)
     int const num_tokens_insert,  // = slot_mapping.size(0), ≤ num_tokens_full
-    int const num_heads_q,        // H
+    int const num_heads_q,        // live Q heads (input layout)
     int const cache_block_size,   // tokens per paged-cache block
     int const kv_block_stride) {  // bytes per paged-cache block
 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
@@ -351,12 +391,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     int const laneId = threadIdx.x % 32;
     int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
 
-    int const total_slots_per_token = num_heads_q + 1;
-    int const tokenIdx = globalWarpIdx / total_slots_per_token;
-    int const slotIdx = globalWarpIdx % total_slots_per_token;
+    constexpr int kTotalSlotsPerToken = kNumHeadsQPadded + 1;
+    int const tokenIdx = globalWarpIdx / kTotalSlotsPerToken;
+    int const slotIdx = globalWarpIdx % kTotalSlotsPerToken;
     if (tokenIdx >= num_tokens_full) return;
 
-    bool const isKV = (slotIdx == num_heads_q);
+    bool const isKV = (slotIdx == kNumHeadsQPadded);
+    bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
     // KV branch: skip DP-padded tokens (no slot reserved for them).
     if (isKV && tokenIdx >= num_tokens_insert) return;
 
@@ -371,22 +412,26 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
     // Dim range this lane owns within the 512-wide head.
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
 
-    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
-    // type; the shared per-slot helper bitcasts to scalar_t_in packed pairs.
-    scalar_t_in const* src_ptr;
-    if (isKV) {
-      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
-    } else {
-      int64_t const q_row_offset =
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      src_ptr = q_inout + q_row_offset;
+    // Load only for live-Q and KV slots; pad-Q skips the read (q_in beyond
+    // num_heads_q is out of bounds) and the helper zero-fills its output.
+    uint4 v0, v1;
+    if (!isPadQ) {
+      scalar_t_in const* src_ptr;
+      if (isKV) {
+        src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+      } else {
+        int64_t const q_row_offset =
+            (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
+                kHeadDim +
+            dim_base;
+        src_ptr = q_in + q_row_offset;
+      }
+      v0 = *reinterpret_cast<uint4 const*>(src_ptr);
+      v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
     }
-    uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
 
-    processDeepseekV4Slot<scalar_t_in>(
-        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
+    processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
+        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
         k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
         kv_block_stride);
 
@@ -408,9 +453,9 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
 // Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
 // KV branch (RoPE + UE8M0 quant + insert)
 //
-template <typename scalar_t_in>
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
-    scalar_t_in* __restrict__ q_inout,  // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
     scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
     int64_t const* __restrict__ slot_mapping,
     int64_t const* __restrict__ position_ids,
@@ -435,25 +480,32 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 #endif
 
     int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-    int const slot_end =
-        (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1);
+    // Slot enumeration: live-Q + pad-Q + (KV if this token has a slot).
+    int const slot_end = (tokenIdx >= num_tokens_insert)
+                             ? kNumHeadsQPadded
+                             : (kNumHeadsQPadded + 1);
 
-    auto src_for_slot = [&](int s) -> scalar_t_in const* {
-      if (s == num_heads_q) {
-        return kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+    auto load_slot = [&](int s, uint4& va, uint4& vb) {
+      // pad-Q slots skip the load — q_in beyond num_heads_q is OOB.
+      if (s >= num_heads_q && s < kNumHeadsQPadded) return;
+      scalar_t_in const* src;
+      if (s == kNumHeadsQPadded) {
+        src = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+      } else {
+        src = q_in +
+              (static_cast<int64_t>(tokenIdx) * num_heads_q +
+               static_cast<int64_t>(s)) *
+                  kHeadDim +
+              dim_base;
       }
-      return q_inout +
-             (static_cast<int64_t>(tokenIdx) * num_heads_q +
-              static_cast<int64_t>(s)) *
-                 kHeadDim +
-             dim_base;
+      va = *reinterpret_cast<uint4 const*>(src);
+      vb = *reinterpret_cast<uint4 const*>(src + 8);
     };
 
     if (warpId < slot_end) {
       int curr_slot = warpId;
-      scalar_t_in const* src_curr = src_for_slot(curr_slot);
-      uint4 v0_curr = *reinterpret_cast<uint4 const*>(src_curr);
-      uint4 v1_curr = *reinterpret_cast<uint4 const*>(src_curr + 8);
+      uint4 v0_curr, v1_curr;
+      load_slot(curr_slot, v0_curr, v1_curr);
 
       while (curr_slot < slot_end) {
         int const next_slot = curr_slot + warpsPerBlock;
@@ -462,14 +514,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
         // Prefetch src for the next slot
         uint4 v0_next, v1_next;
         if (has_next) {
-          scalar_t_in const* src_next = src_for_slot(next_slot);
-          v0_next = *reinterpret_cast<uint4 const*>(src_next);
-          v1_next = *reinterpret_cast<uint4 const*>(src_next + 8);
+          load_slot(next_slot, v0_next, v1_next);
         }
 
-        processDeepseekV4Slot<scalar_t_in>(
+        processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
             v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
-            num_heads_q, eps, q_inout, k_cache, slot_mapping, position_ids,
+            num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
             cos_sin_cache, cache_block_size, kv_block_stride);
 
         // ── Buffer rotation: hand the prefetched LDGs to the next iter.
@@ -490,10 +540,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 // ────────────────────────────────────────────────────────────────────────────
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t_in>
-void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
-    scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
-    int64_t const* slot_mapping, int64_t const* position_ids,
+template <typename scalar_t_in, int kNumHeadsQPadded>
+static void launchFusedDeepseekV4Templated(
+    scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
+    uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
     float const* cos_sin_cache, float const eps, int const num_tokens_full,
     int const num_tokens_insert, int const num_heads_q,
     int const cache_block_size, int const kv_block_stride,
@@ -501,7 +551,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
   constexpr int kBlockSize = 256;
   constexpr int kWarpsPerBlock = kBlockSize / 32;
   int64_t const total_warps =
-      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
+      static_cast<int64_t>(num_tokens_full) * (kNumHeadsQPadded + 1);
   int const grid =
       static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
 
@@ -532,46 +582,85 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
 
   if (num_tokens_full < NUM_TOKEN_CUTOFF) {
     cudaLaunchKernelEx(
-        &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
-        q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-        num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        &config,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in,
+                                                        kNumHeadsQPadded>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
         kv_block_stride);
   } else {
     config.gridDim = dim3(num_tokens_full);
     cudaLaunchKernelEx(
         &config,
-        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<scalar_t_in>,
-        q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-        num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<
+            scalar_t_in, kNumHeadsQPadded>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
         kv_block_stride);
   }
-
 #else
   // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
   // clang-format off
-  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
+  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in, kNumHeadsQPadded>
       <<<grid, kBlockSize, 0, stream>>>(
-          q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
-          eps, num_tokens_full, num_tokens_insert, num_heads_q,
+          q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
+          cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
           cache_block_size, kv_block_stride);
 #endif
 }
 
+// Runtime dispatch into one of the precompiled `kNumHeadsQPadded`
+// instantiations.  Supported padded head counts: 8, 16, 32, 64, 128.
+template <typename scalar_t_in>
+void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
+    scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
+    uint8_t* k_cache, int64_t const* slot_mapping,
+    int64_t const* position_ids, float const* cos_sin_cache, float const eps,
+    int const num_tokens_full, int const num_tokens_insert,
+    int const num_heads_q, int const num_heads_q_padded,
+    int const cache_block_size, int const kv_block_stride,
+    cudaStream_t stream) {
+#define DISPATCH(N)                                                         \
+  case N:                                                                   \
+    launchFusedDeepseekV4Templated<scalar_t_in, N>(                         \
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,            \
+        cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, \
+        cache_block_size, kv_block_stride, stream);                         \
+    return;
+
+  switch (num_heads_q_padded) {
+    DISPATCH(8)
+    DISPATCH(16)
+    DISPATCH(32)
+    DISPATCH(64)
+    DISPATCH(128)
+    default:
+      TORCH_CHECK(false,
+                  "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
+                  "unsupported num_heads_q_padded=",
+                  num_heads_q_padded,
+                  " (compiled instantiations: 8, 16, 32, 64, 128).");
+  }
+#undef DISPATCH
+}
+
 }  // namespace deepseek_v4_fused_ops
 }  // namespace vllm
 
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q,                    // [N, H, 512] bf16, in place
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
     torch::Tensor const& kv,             // [N, 512] bf16 (read-only)
     torch::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
     torch::Tensor const& slot_mapping,   // [N] int64
     torch::Tensor const& position_ids,   // [N] int64
     torch::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
+    int64_t q_head_padded,               // padded Q head count for output
     double eps, int64_t cache_block_size) {
-  TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
+  TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(),
+              "q_in must be contiguous CUDA");
   TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
   TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
   TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
@@ -579,9 +668,12 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
               "position_ids must be int64 CUDA");
   TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
-  TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
+  TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
+              "q_in shape [N, num_heads_q, 512]");
   TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
+  TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match");
+  TORCH_CHECK(q_head_padded >= q_in.size(1),
+              "q_head_padded must be >= q_in.size(1) (num_heads_q)");
   TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
   TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
               "cos_sin_cache shape [max_pos, 64]");
@@ -591,32 +683,41 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
   // With DP padding, slot_mapping can be shorter than q/kv/positions.
   // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
   // KV quant+insert runs only on the first slot_mapping.size(0) rows.
-  int const num_tokens_full = static_cast<int>(q.size(0));
+  int const num_tokens_full = static_cast<int>(q_in.size(0));
   int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
   TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
                   static_cast<int>(position_ids.size(0)) == num_tokens_full,
               "q/kv/position_ids row counts must match");
   TORCH_CHECK(num_tokens_insert <= num_tokens_full,
               "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
+  int const num_heads_q = static_cast<int>(q_in.size(1));
+  int const num_heads_q_padded = static_cast<int>(q_head_padded);
   int const cache_block_size_i = static_cast<int>(cache_block_size);
   int const kv_block_stride = static_cast<int>(k_cache.stride(0));
 
-  at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+  at::cuda::OptionalCUDAGuard device_guard(device_of(q_in));
   auto stream = at::cuda::getCurrentCUDAStream();
 
+  // Allocate the padded q output.  The kernel writes every element (live
+  // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
+  torch::Tensor q_out = torch::empty(
+      {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());
+
   VLLM_DISPATCH_HALF_TYPES(
-      q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
+      q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
         using qkv_scalar_t = scalar_t;
         vllm::deepseek_v4_fused_ops::
             launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
-                reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
+                reinterpret_cast<qkv_scalar_t const*>(q_in.data_ptr()),
+                reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
                 reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
                 reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
                 reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
                 reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
                 cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
                 num_tokens_full, num_tokens_insert, num_heads_q,
-                cache_block_size_i, kv_block_stride, stream);
+                num_heads_q_padded, cache_block_size_i, kv_block_stride,
+                stream);
       });
-}
+  return q_out;
+}

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.
@@ -70,30 +63,11 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
                         torch::Tensor& weight, double epsilon);
 
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
     torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
-    torch::Tensor const& cos_sin_cache, 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);
+    torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
+    int64_t cache_block_size);
 
 void silu_and_mul_per_block_quant(torch::Tensor& out,
                                   torch::Tensor const& input,
@@ -149,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) -> ()");
@@ -99,37 +84,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // kernel launch.
   ops.def(
       "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor! q, Tensor kv, Tensor! k_cache, "
+      "Tensor q_in, Tensor kv, Tensor! k_cache, "
       "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "float eps, int cache_block_size) -> ()");
+      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
   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

@@ -220,7 +220,8 @@ COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
-        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
+        sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' requirements/cuda.txt; \
     fi \
     && if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
         echo "Installing torch nightly..." \
@@ -746,7 +747,8 @@ COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
 RUN --mount=type=cache,target=/opt/uv/cache \
     if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
-        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' /tmp/requirements-cuda.txt; \
     fi && \
     uv pip install --system -r /tmp/requirements-cuda.txt \
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
@@ -866,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
@@ -905,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

@@ -231,13 +231,21 @@ vllm bench serve \
 
 #### Custom Image Dataset
 
-If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
+If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and can use "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
 
 ```json
 {"prompt": "How many animals are present in the given image?", "image_files": ["/path/to/image/folder/horsepony.jpg"]}
 {"prompt": "What colour is the bird shown in the image?", "image_files": ["/path/to/image/folder/flycatcher.jpeg"]}
 ```
 
+Every image listed in "image_files" is added to the request in the listed order after the prompt text. To preserve an interleaved order of text and images, use a "content" field with OpenAI-compatible content parts:
+
+```json
+{"content": [{"type": "text", "text": "Compare "}, {"type": "image", "image": "/path/to/image/folder/chart_a.png"}, {"type": "text", "text": " with "}, {"type": "image_url", "image_url": {"url": "/path/to/image/folder/chart_b.png"}}]}
+```
+
+The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string.
+
 ```bash
 # need a model with vision capability here
 vllm serve Qwen/Qwen2-VL-7B-Instruct
@@ -910,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/disagg_prefill.md

@@ -19,25 +19,25 @@ Two main reasons:
 
 Please refer to [examples/disaggregated/disaggregated_prefill.sh](../../examples/disaggregated/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
 
-Now supports 6 types of connectors:
+Now supports 9 types of connectors:
 
 - **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
 - **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
-- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md).
+- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as:
+
+  ```bash
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
+  ```
+
 - **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
 - **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
+- **MoRIIOConnector** (ROCm only): see [MoRI-IO Usage Guide](moriio_connector_usage.md) for example usage and detailed documentation.
 - **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:
 
   ```bash
   --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both"},{"kv_connector":"ExampleConnector","kv_role":"kv_both","kv_connector_extra_config":{"shared_storage_path":"local_storage"}}]}}'
   ```
 
-For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
-
-  ```bash
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
-  ```
-
 - **OffloadingConnector**: enable offloading of KV data to CPU memory, customizing the CPU block size (in tokens) and total CPU memory bytes to allocate:
 
   ```bash

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/moriio_connector_usage.md

@@ -0,0 +1,264 @@
+# MoRIIOConnector Usage Guide
+
+`MoRIIOConnector` is a high-performance KV connector used for KV cache transfer in PD disaggregated deployments, built on ROCm's [MoRI-IO](https://github.com/rocm/mori) communication library for point-to-point communication with ultra-low overhead.
+
+## Prerequisites
+
+### Installation
+
+**Docker:** MoRI is shipped with the official ROCm vLLM image: `vllm/vllm-openai-rocm:nightly`.
+
+**Manual installation:** MoRI wheel can be installed with
+
+```bash
+pip install amd_mori
+```
+
+Refer to the [Dockerfile.rocm_base](../../docker/Dockerfile.rocm_base) for more information, or [official MoRI repository](https://github.com/rocm/mori) for instructions on how to build MoRI from source.
+
+For instructions on installing appropriate NIC userspace libraries, see [Installing NIC userspace libraries](#appendix-installing-nic-userspace-libraries).
+
+## Basic usage (single host)
+
+Start the proxy first; the producer and consumer instances will retry registration until the proxy is reachable.
+
+### Producer (prefiller) configuration
+
+Start a prefiller instance that produces KV caches
+
+```bash
+# Prefill instance (GPU 0-3) 
+export VLLM_ROCM_USE_AITER=1
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+export HIP_VISIBLE_DEVICES=0,1,2,3
+ 
+vllm serve Qwen/Qwen3-235B-A22B-FP8 \
+  -tp 4 \
+  --port 20005 \
+  --gpu-memory-utilization 0.9 \
+  --kv-transfer-config '{
+    "kv_connector": "MoRIIOConnector",
+    "kv_role": "kv_producer",
+    "kv_connector_extra_config": {
+      "proxy_ip": "127.0.0.1",
+      "proxy_ping_port": "36367",
+      "http_port": "20005",
+      "handshake_port": "6301",
+      "notify_port": "6105"
+    }
+  }'
+```
+
+### Consumer (decoder) configuration
+
+Start a decoder instance that consumes KV caches:
+
+```bash
+# Decode instance (GPU 4-7)
+export VLLM_ROCM_USE_AITER=1
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+export HIP_VISIBLE_DEVICES=4,5,6,7
+
+vllm serve Qwen/Qwen3-235B-A22B-FP8 \
+  -tp 4 \
+  --port 40005 \
+  --gpu-memory-utilization 0.9 \
+  --kv-transfer-config '{
+    "kv_connector": "MoRIIOConnector",
+    "kv_role": "kv_consumer",
+    "kv_connector_extra_config": {
+      "proxy_ip": "127.0.0.1",
+      "http_port": "40005",
+      "proxy_ping_port": "36367",
+      "handshake_port": "7301",
+      "notify_port": "7501"
+    }
+  }'
+```
+
+### Proxy server
+
+The proxy fronts the producer and consumer instances and routes incoming requests to them. `vllm-router` is the recommended proxy; it can be installed manually or run as a Docker container. Note that the port `36367` below is the `proxy_ping_port` configured on each vLLM instance.
+
+**Docker:**
+
+```bash
+docker run \
+  --network host \
+  vllm/vllm-router:nightly \
+  vllm-router \
+  --vllm-pd-disaggregation \
+  --kv-connector moriio \
+  --vllm-discovery-address "0.0.0.0:36367"
+```
+
+**Manual install:**
+
+```bash
+pip install vllm-router
+vllm-router \
+  --vllm-pd-disaggregation \
+  --kv-connector moriio \
+  --vllm-discovery-address "0.0.0.0:36367"
+```
+
+Alternatively, you can use the reference implementation proxy shipped with vLLM:
+
+```bash
+cd <path_to>/vllm
+pip install quart aiohttp msgpack
+python examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+```
+
+## Configuration
+
+The connector is configured at two levels: the application level and the transport level.
+
+### Application-level configuration
+
+**Modes:** MoRI has two modes of operation: WRITE and READ mode.
+
+- In WRITE mode, the producer actively pushes computed KV blocks after every layer into the consumer's memory.
+- In READ mode, the consumer pulls the KV blocks from the producer all at once, as soon as it has been notified those blocks are ready.
+
+WRITE mode is used by default. READ mode can be configured by setting `--kv-transfer-config.kv_connector_extra_config.read_mode true`.
+
+**Control-plane configuration:** MoRI moves KV bytes over RDMA/xGMI, but producers and consumers also need out-of-band TCP channels for handshake, block id exchange, liveness, and completion signaling. These keys live under `kv_connector_extra_config`:
+
+- `proxy_ip`: IP address of the disaggregation proxy/router that fronts the prefiller and decoder. Each vLLM instance uses it to register itself and to send heartbeats so the proxy knows where to route incoming requests.
+- `proxy_ping_port`: TCP port on `proxy_ip` where the proxy listens for instance heartbeats and registration messages. Used to detect dead vLLM instances and keep routing tables fresh.
+- `http_port`: HTTP port that this vLLM instance exposes its OpenAI-compatible API on. The proxy registers this port, and forwards user requests to this port once it has picked an instance.
+- `handshake_port`: TCP port used for the one-time MoRI engine handshake between a prefiller and a decoder. The two sides exchange RDMA engine descriptors here before any KV transfer can happen.
+- `notify_port`: TCP port used for control and synchronization messages between prefiller and decoder. Used differently in the two modes:
+    - WRITE mode: **Block allocation:** the decoder notifies the prefiller about its block ids, so the prefiller can push its computed KV blocks into the correct place on the decoder instance. **Completion:** once all blocks have been transferred, the prefiller notifies the decoder that it's safe to use its blocks.
+    - READ mode: **Completion:** once the decoder has read all blocks from the prefiller, it notifies the prefiller so it can free its KV cache blocks.
+
+!!! note
+    `notify_port` is used as a *base* port: each (DP rank, TP rank) pair within an instance uses `notify_port + offset` where the offset is based on the rank. Make sure the range starting at `notify_port` is free on the host.
+
+### Transport configuration
+
+MoRI has two transport backends: RDMA and xGMI. You can select backend using `--kv-transfer-config.kv_connector_extra_config.backend $BACKEND`, with `$BACKEND` being `rdma` or `xgmi`. RDMA is the default backend and should be used in multi-node deployments.
+
+The configuration options for each backend are as follows.
+
+#### RDMA backend
+
+- `qp_per_transfer`: number of RDMA Queue Pairs (QPs) used per transfer. More QPs let a single transfer be striped over multiple QPs to increase NIC concurrency, at the cost of more RDMA resources.
+- `post_batch_size`: how many RDMA Work Requests (WR) are batched into one `ibv_post_send` doorbell. Defaults to -1, meaning the backend default. Larger batches reduce the posting overhead per WR.
+- `num_workers`: number of worker threads MoRI uses to post and poll transfer completions.
+
+Advanced users can also configure MoRI itself using environment variables such as `MORI_IO_QP_MAX_SEND_WR`, `MORI_IO_QP_MAX_CQE`, etc. These are MoRI library variables and are separate from vLLM's own `VLLM_MORIIO_*` settings. Refer to the [MoRI repository](https://github.com/rocm/mori) for more information.
+
+#### xGMI backend
+
+Use xGMI when the prefiller and decoder run on the same physical host so transfers go over the AMD GPU fabric and skip the NIC entirely. Currently only configured using MoRI-specific environment variables; see the [MoRI repository](https://github.com/rocm/mori).
+
+## Multi-node deployment
+
+The example below shows how to run a 1P1D deployment on two nodes. We run the proxy on the same node as the prefill instance.
+
+### On both nodes
+
+```bash
+# Set on both nodes before running any command
+export PREFILL_IP=<node1-ip>
+export DECODE_IP=<node2-ip>
+```
+
+### On node 1
+
+Start the proxy first as described in [Proxy server](#proxy-server), then start the prefill instance:
+
+```bash
+docker run \
+  --name moriio-prefill \
+  --init --network host --ipc host --privileged \
+  --security-opt seccomp=unconfined \
+  --ulimit memlock=-1 --ulimit stack=67108864 --shm-size 256G \
+  --group-add video --group-add render \
+  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+  -e VLLM_ROCM_USE_AITER=1 \
+  vllm/vllm-openai-rocm:nightly \
+  deepseek-ai/DeepSeek-R1-0528 \
+    --port 8100 \
+    --tensor-parallel-size 8 \
+    --enable-expert-parallel \
+    --gpu-memory-utilization 0.8 \
+    --trust-remote-code \
+    --kv-transfer-config '{
+      "kv_connector": "MoRIIOConnector",
+      "kv_role": "kv_producer",
+      "kv_connector_extra_config": {
+        "proxy_ip": "'"${PREFILL_IP}"'",
+        "proxy_ping_port": "36367",
+        "http_port": "8100",
+        "handshake_port": "6301",
+        "notify_port": "61005"
+      }
+    }'
+```
+
+### On node 2
+
+Decode instance:
+
+```bash
+docker run \
+  --name moriio-decode \
+  --init --network host --ipc host --privileged \
+  --security-opt seccomp=unconfined \
+  --ulimit memlock=-1 --ulimit stack=67108864 --shm-size 256G \
+  --group-add video --group-add render \
+  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+  -e VLLM_ROCM_USE_AITER=1 \
+  vllm/vllm-openai-rocm:nightly \
+  deepseek-ai/DeepSeek-R1-0528 \
+    --port 8200 \
+    --tensor-parallel-size 8 \
+    --gpu-memory-utilization 0.8 \
+    --trust-remote-code \
+    --enable-expert-parallel \
+    --kv-transfer-config '{
+      "kv_connector": "MoRIIOConnector",
+      "kv_role": "kv_consumer",
+      "kv_connector_extra_config": {
+        "proxy_ip": "'"${PREFILL_IP}"'",
+        "proxy_ping_port": "36367",
+        "http_port": "8200",
+        "handshake_port": "6301",
+        "notify_port": "61005"
+      }
+    }'
+```
+
+## Troubleshooting
+
+### `availDevices.size() > 0` assertion failure
+
+**Problem:** vLLM fails to launch with the following log:
+
+```bash
+libibverbs: Warning: Driver bnxt_re does not support the kernel ABI of 6 (supports 1 to 1) for device /sys/class/infiniband/rdma4
+...
+ker: /app/mori/src/io/rdma/backend_impl.cpp: mori::io::RdmaManager::RdmaManager(const RdmaBackendConfig, application::RdmaContext *): Assertion `availDevices.size() > 0' failed.
+```
+
+**Fix:** The installed RDMA userspace libraries do not match the driver and firmware version installed on the host. You must install NIC userspace libraries corresponding to your RDMA kernel module and firmware version. See [Installing NIC userspace
+libraries](#appendix-installing-nic-userspace-libraries) for more information.
+
+## Appendix: installing NIC userspace libraries
+
+To run MoRI with RDMA, your environment must have the necessary RDMA userspace libraries installed that match the associated kernel module and firmware version.
+
+The official image `vllm/vllm-openai-rocm:nightly` comes pre-installed with userspace libraries for the following NICs and kernel module versions:
+
+- AINIC (AMD Pensando Pollara): version `1.117.3-hydra`, tested with `ioinic-dkms=25.11.1.001`
+- Thor2 (Broadcom): version `235.2.86.0`, tested with `bnxt-en-dkms=1.10.3.235.2.86.0`, `bnxt-re-dkms=235.2.86.0`
+
+Refer to [Dockerfile.rocm](../../docker/Dockerfile.rocm) for more details. For users with NICs, kernel modules, and/or FW other than those stated above we refer to
+the vendors' own installation instructions.
+
+## Further reading
+
+- [Next-Level Inference: Why Your Single-Node vLLM Setup Needs Prefill-Decode Disaggregation](https://vllm.ai/blog/2026-04-07-moriio-kv-connector).

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,8 +21,11 @@ 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
-tokenspeed-mla==0.1.2
+tokenspeed-mla==0.1.2
+
+# Humming kernels for quantization gemm
+humming-kernels[cu13]==0.1.2

requirements/kv_connectors.txt

@@ -1,3 +1,6 @@
 lmcache >= 0.3.9
+# CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0
+# until a fixed newer release is verified for runtime images.
+cupy-cuda13x < 14.1.0
 nixl >= 1.1.0 # Required for disaggregated prefill
 mooncake-transfer-engine >= 0.3.8

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());
     }