Update quack-kernels requirement from >=0.3.3 to >=0.4.1

Updates the requirements on quack-kernels to permit the latest version.

---
updated-dependencies:
- dependency-name: quack-kernels
  dependency-version: 0.4.1
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>

.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/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/entrypoints.yaml

@@ -29,6 +29,7 @@ steps:
   mirror:
     amd:
       device: mi300_1
+      soft_fail: true
       depends_on:
       - image-build-amd
 
@@ -46,6 +47,7 @@ steps:
   mirror:
     amd:
       device: mi300_1
+      soft_fail: true
       timeout_in_minutes: 80
       depends_on:
       - image-build-amd
@@ -64,6 +66,7 @@ steps:
   mirror:
     amd:
       device: mi300_1
+      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -83,6 +86,7 @@ steps:
   mirror:
     amd:
       device: mi300_1
+      soft_fail: true
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -105,6 +109,7 @@ steps:
   mirror:
     amd:
       device: mi300_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

.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:

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:

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/ops.h

@@ -70,10 +70,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);
+    torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
+    int64_t cache_block_size);
 
 void apply_repetition_penalties_(torch::Tensor& logits,
                                  const torch::Tensor& prompt_mask,

csrc/torch_bindings.cpp

@@ -99,9 +99,9 @@ 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);
 

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 \

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

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

pyproject.toml

@@ -181,6 +181,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/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
-quack-kernels>=0.3.3
+nvidia-cutlass-dsl[cu13]==4.5.2
+quack-kernels>=0.4.1
 
 # 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/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

rust/Cargo.lock

@@ -2591,6 +2591,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 +5609,7 @@ dependencies = [
  "minijinja",
  "minijinja-contrib",
  "openai-harmony",
+ "paste",
  "reqwest",
  "rmp-serde",
  "serde",
@@ -5810,6 +5812,7 @@ dependencies = [
  "serde",
  "serde_json",
  "serde_with",
+ "serial_test",
  "tempfile",
  "thiserror 2.0.18",
  "thiserror-ext",
@@ -5847,6 +5850,7 @@ name = "vllm-tool-parser"
 version = "0.1.0"
 dependencies = [
  "criterion",
+ "easy-ext",
  "expect-test",
  "futures",
  "openai-protocol",

rust/Cargo.toml

@@ -45,13 +45,14 @@ 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"] }
+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"
@@ -65,11 +66,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/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

@@ -6,7 +6,7 @@ pub use vllm_tool_parser::{
     DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
     Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, KimiK2ToolParser,
     Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser, Qwen3CoderToolParser,
-    Qwen3XmlToolParser, ToolCallDelta, ToolParseResult, ToolParser, ToolParserError,
+    Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError, ToolParserOutput,
 };
 
 use crate::parser::ParserFactory;

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

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/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/src/deepseek_dsml/deepseek_v32.rs

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

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

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

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

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

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

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

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

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

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

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

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

@@ -7,7 +7,7 @@ use winnow::stream::Partial;
 use winnow::token::{literal, take_till, take_until};
 
 use super::utils::{parse_buffered_event, safe_text_len};
-use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
+use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
 use crate::Tool;
 
 const TOOL_CALL_START: &str = "<|tool_call>";
@@ -51,16 +51,16 @@ impl Gemma4ToolParser {
         }
     }
 
-    fn apply_event(&mut self, event: Gemma4Event, result: &mut ToolParseResult) -> Result<()> {
+    fn apply_event(&mut self, event: Gemma4Event, output: &mut ToolParserOutput) -> Result<()> {
         match event {
             Gemma4Event::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
             }
             Gemma4Event::ToolCall { name, args } => {
                 let arguments = serde_json::to_string(&args)
                     .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
 
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index: self.emitted_tool_count,
                     name: Some(name),
                     arguments,
@@ -71,9 +71,9 @@ impl Gemma4ToolParser {
         Ok(())
     }
 
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
     }
 }
 
@@ -89,33 +89,35 @@ impl ToolParser for Gemma4ToolParser {
         true
     }
 
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
         while let Some((event, consumed_len)) =
             parse_buffered_event(&self.buffer, parse_next_gemma4_event)?
         {
-            self.apply_event(event, &mut result)?;
+            self.apply_event(event, output)?;
             self.buffer.drain(..consumed_len);
         }
 
-        Ok(result)
+        Ok(())
     }
 
-    fn finish(&mut self) -> Result<ToolParseResult> {
-        let mut result = ToolParseResult::default();
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
 
         if !self.buffer.is_empty() {
             if self.buffer.starts_with(TOOL_CALL_START) {
-                self.reset();
                 return Err(parsing_failed!("incomplete Gemma4 tool call"));
             }
-            result.normal_text.push_str(&self.buffer);
+            output.normal_text.push_str(&self.buffer);
         }
 
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
+    }
+
+    fn reset(&mut self) -> String {
+        Gemma4ToolParser::reset(self)
     }
 }
 
@@ -285,10 +287,10 @@ mod tests {
     use winnow::stream::Partial;
 
     use super::{
-        Gemma4ToolParser, ToolCallDelta, ToolParseResult, ToolParser, gemma4_args,
+        Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_args,
         gemma4_array_content,
     };
-    use crate::Tool;
+    use crate::{Tool, ToolParserTestExt as _};
 
     fn parse_gemma4_args(args: &str) -> super::Result<serde_json::Map<String, Value>> {
         let mut input = Partial::new(args);
@@ -367,18 +369,18 @@ mod tests {
         ]
     }
 
-    fn collect_stream(chunks: &[&str]) -> ToolParseResult {
+    fn collect_stream(chunks: &[&str]) -> ToolParserOutput {
         let mut parser = Gemma4ToolParser::new(&test_tools());
-        let mut result = ToolParseResult::default();
+        let mut output = ToolParserOutput::default();
         for chunk in chunks {
-            result.append(parser.push(chunk).unwrap());
+            output.append(parser.parse_chunk(chunk).unwrap());
         }
-        result.append(parser.finish().unwrap());
-        result.coalesce_calls()
+        output.append(parser.finish().unwrap());
+        output.coalesce_calls()
     }
 
-    fn first_call(result: &ToolParseResult) -> &ToolCallDelta {
-        result.calls.first().expect("expected one tool call")
+    fn first_call(output: &ToolParserOutput) -> &ToolCallDelta {
+        output.calls.first().expect("expected one tool call")
     }
 
     #[test]
@@ -416,15 +418,15 @@ mod tests {
     #[test]
     fn gemma4_parse_complete_extracts_single_tool_call() {
         let mut parser = Gemma4ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete("<|tool_call>call:get_weather{location:<|\"|>London<|\"|>}<tool_call|>")
             .unwrap();
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
+        assert!(output.normal_text.is_empty());
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "location": "London" })
         );
     }
@@ -441,7 +443,7 @@ mod tests {
 
     #[test]
     fn gemma4_streaming_basic_single_tool_call() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:get_weather{",
             "location:<|\"|>Paris",
@@ -450,17 +452,17 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert!(result.normal_text.is_empty());
-        assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
+        assert!(output.normal_text.is_empty());
+        assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "location": "Paris, France" })
         );
     }
 
     #[test]
     fn gemma4_streaming_text_before_and_after_tool_call() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "Let me check ",
             "the weather. ",
             "<|tool_call>",
@@ -470,10 +472,10 @@ mod tests {
             "div>",
         ]);
 
-        assert_eq!(result.normal_text, "Let me check the weather. <div>");
-        assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
+        assert_eq!(output.normal_text, "Let me check the weather. <div>");
+        assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "location": "London" })
         );
     }
@@ -481,68 +483,68 @@ mod tests {
     #[test]
     fn gemma4_streaming_waits_for_complete_tool_call() {
         let mut parser = Gemma4ToolParser::new(&test_tools());
-        let mut result = ToolParseResult::default();
+        let mut output = ToolParserOutput::default();
 
         for chunk in [
             "<|tool_call>",
             "call:get_weather{",
             "location:<|\"|>Paris<|\"|>}",
         ] {
-            result.append(parser.push(chunk).unwrap());
-            assert!(result.calls.is_empty());
+            output.append(parser.parse_chunk(chunk).unwrap());
+            assert!(output.calls.is_empty());
         }
 
-        result.append(parser.push("<tool_call|>").unwrap());
-        let result = result.coalesce_calls();
+        output.append(parser.parse_chunk("<tool_call|>").unwrap());
+        let output = output.coalesce_calls();
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("get_weather"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "location": "Paris" })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_boolean_split_across_chunks() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:search{input:{all:tru",
             "e}}",
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("search"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("search"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "input": { "all": true } })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_false_split_across_chunks() {
-        let result = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", "<tool_call|>"]);
+        let output = collect_stream(&["<|tool_call>", "call:set{flag:fals", "e}", "<tool_call|>"]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("set"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("set"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "flag": false })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_number_split_across_chunks() {
-        let result = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", "<tool_call|>"]);
+        let output = collect_stream(&["<|tool_call>", "call:set{count:4", "2}", "<tool_call|>"]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("set"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("set"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "count": 42 })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_split_string_delimiter() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:todowrite{",
             "content:<|\"|>Buy milk<|",
@@ -550,17 +552,17 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("todowrite"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("todowrite"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "content": "Buy milk" })
         );
-        assert!(!first_call(&result).arguments.contains("<|"));
+        assert!(!first_call(&output).arguments.contains("<|"));
     }
 
     #[test]
     fn gemma4_streaming_handles_end_marker_literal_inside_string() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:todowrite{",
             "content:<|\"|>literal }<tool_call|> inside",
@@ -568,16 +570,16 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("todowrite"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("todowrite"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({ "content": "literal }<tool_call|> inside" })
         );
     }
 
     #[test]
     fn gemma4_streaming_handles_html_argument_without_duplication() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:write_file{",
             "path:<|\"|>index.html<|\"|>,",
@@ -590,9 +592,9 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("write_file"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("write_file"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({
                 "path": "index.html",
                 "content": "<!DOCTYPE html>\n<html lang=\"zh-CN\">\n<head>\n    <meta charset=\"UTF-8\">\n    <meta name=\"viewport\" content=\"width=device-width\">\n",
@@ -602,7 +604,7 @@ mod tests {
 
     #[test]
     fn gemma4_streaming_trailing_bare_bool_is_not_duplicated() {
-        let result = collect_stream(&[
+        let output = collect_stream(&[
             "<|tool_call>",
             "call:Edit{",
             "file_path:<|\"|>src/env.py<|\"|>,",
@@ -613,9 +615,9 @@ mod tests {
             "<tool_call|>",
         ]);
 
-        assert_eq!(first_call(&result).name.as_deref(), Some("Edit"));
+        assert_eq!(first_call(&output).name.as_deref(), Some("Edit"));
         assert_eq!(
-            serde_json::from_str::<Value>(&first_call(&result).arguments).unwrap(),
+            serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
             json!({
                 "file_path": "src/env.py",
                 "old_string": "old_val",
@@ -624,7 +626,7 @@ mod tests {
             })
         );
         assert_eq!(
-            first_call(&result).arguments.matches("replace_all").count(),
+            first_call(&output).arguments.matches("replace_all").count(),
             1
         );
     }
@@ -632,18 +634,18 @@ mod tests {
     #[test]
     fn gemma4_finish_flushes_partial_start_marker_as_text() {
         let mut parser = Gemma4ToolParser::new(&test_tools());
-        let mut result = parser.push("<").unwrap();
-        result.append(parser.finish().unwrap());
+        let mut output = parser.parse_chunk("<").unwrap();
+        output.append(parser.finish().unwrap());
 
-        assert_eq!(result.normal_text, "<");
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, "<");
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn gemma4_finish_rejects_complete_args_without_end_marker() {
         let mut parser = Gemma4ToolParser::new(&test_tools());
         for chunk in ["<|tool_call>", "call:get_status{}"] {
-            parser.push(chunk).unwrap();
+            parser.parse_chunk(chunk).unwrap();
         }
 
         let error = parser.finish().unwrap_err();

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

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

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

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

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

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

rust/src/tool-parser/src/json/hermes.rs

@@ -1,5 +1,5 @@
 use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
     parser_name: "Hermes",
@@ -38,7 +38,6 @@ impl HermesToolParser {
 }
 
 impl ToolParser for HermesToolParser {
-    /// Create a boxed Hermes tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -46,15 +45,17 @@ impl ToolParser for HermesToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the Hermes parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.inner.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.inner.parse_into(chunk, output)
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.inner.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.inner.reset()
+    }
 }
 
 #[cfg(test)]
@@ -64,7 +65,7 @@ mod tests {
 
     use super::HermesToolParser;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
-    use crate::{ToolParseResult, ToolParser};
+    use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
 
     fn build_tool_call(function_name: &str, arguments: &str) -> String {
         format!(r#"<tool_call>{{"name":"{function_name}","arguments":{arguments}}}</tool_call>"#)
@@ -73,51 +74,51 @@ mod tests {
     #[test]
     fn hermes_parse_complete_without_tool_call_keeps_text() {
         let mut parser = HermesToolParser::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 hermes_parse_complete_extracts_raw_json_arguments() {
         let mut parser = HermesToolParser::new(&test_tools());
         let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
-        let result = parser
+        let output = parser
             .parse_complete(&format!(
                 "Let me check.\n{}",
                 build_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 hermes_accepts_newline_after_tool_call_start() {
         let mut parser = HermesToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(
                 r#"<tool_call>
 {"name":"get_weather","arguments":{}}</tool_call>"#,
             )
             .unwrap();
 
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
     }
 
     #[test]
     fn hermes_does_not_validate_or_normalize_arguments() {
         let mut parser = HermesToolParser::new(&test_tools());
         let arguments = r#"{"location":"Tokyo",}"#;
-        let result = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
+        let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
 
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
@@ -132,24 +133,24 @@ mod tests {
             "}</tool_call> suffix",
         ];
 
-        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.normal_text, "preface  suffix");
+        assert_eq!(output.normal_text, "preface  suffix");
         assert_eq!(
-            result.coalesce_calls().calls[0].arguments,
+            output.coalesce_calls().calls[0].arguments,
             r#"{"location":"Beijing"}"#
         );
     }
@@ -163,11 +164,11 @@ mod tests {
         let chunks = split_by_chars(&input, 5);
         let mut parser = HermesToolParser::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]
@@ -180,10 +181,10 @@ mod tests {
         let chunks = split_by_chars(&input, 7);
         let mut parser = HermesToolParser::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 hermes_finish_fails_incomplete_tool_call() {
         let mut parser = HermesToolParser::new(&test_tools());
         parser
-            .push(r#"<tool_call>{"name":"get_weather","arguments":{"location""#)
+            .parse_chunk(r#"<tool_call>{"name":"get_weather","arguments":{"location""#)
             .unwrap();
 
         let error = parser.finish().unwrap_err();

rust/src/tool-parser/src/json/llama.rs

@@ -9,7 +9,7 @@ use super::{
     argument_delta_event, tool_call_header_event,
 };
 use crate::utils::{JsonObjectScanState, parse_buffered_event};
-use crate::{Result, Tool, ToolCallDelta, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 enum LlamaJsonMode {
@@ -78,7 +78,7 @@ impl Llama3JsonToolParser {
     }
 
     /// Apply one parsed Llama JSON event to parser state and output.
-    fn apply_event(&mut self, event: LlamaJsonEvent, result: &mut ToolParseResult) -> Result<()> {
+    fn apply_event(&mut self, event: LlamaJsonEvent, output: &mut ToolParserOutput) -> Result<()> {
         match event {
             LlamaJsonEvent::ToolCallHeader { function_name } => {
                 let tool_index = self.emitted_tool_count;
@@ -87,7 +87,7 @@ impl Llama3JsonToolParser {
                 self.mode = LlamaJsonMode::Arguments {
                     json_scan: JsonObjectScanState::default(),
                 };
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index,
                     name: Some(function_name),
                     arguments: String::new(),
@@ -99,7 +99,7 @@ impl Llama3JsonToolParser {
                         "Llama JSON arguments without an active tool call"
                     ));
                 };
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index,
                     name: None,
                     arguments: self.buffer[..consumed_len].to_string(),
@@ -117,17 +117,15 @@ impl Llama3JsonToolParser {
         Ok(())
     }
 
-    /// Reset all streaming state.
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.mode = LlamaJsonMode::Start;
         self.active_tool_index = None;
         self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
     }
 }
 
 impl ToolParser for Llama3JsonToolParser {
-    /// Create a boxed Llama JSON tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -135,37 +133,34 @@ impl ToolParser for Llama3JsonToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the Llama JSON parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
         if !self.commit_start() {
-            return Ok(result);
+            return Ok(());
         }
 
         if matches!(self.mode, LlamaJsonMode::Passthrough) {
-            result.normal_text.push_str(&self.buffer);
+            output.normal_text.push_str(&self.buffer);
             self.buffer.clear();
-            return Ok(result);
+            return Ok(());
         }
 
         while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
             parse_next_llama_json_event(input, &mut self.mode)
         })? {
-            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 {
             LlamaJsonMode::Start | LlamaJsonMode::Passthrough => {
-                result.normal_text.push_str(&self.buffer);
+                output.normal_text.push_str(&self.buffer);
             }
             LlamaJsonMode::AfterCall if self.buffer.trim().is_empty() => {}
             LlamaJsonMode::Header | LlamaJsonMode::Arguments { .. } => {
@@ -175,8 +170,12 @@ impl ToolParser for Llama3JsonToolParser {
                 return Err(parsing_failed!("invalid Llama JSON"));
             }
         }
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
+    }
+
+    fn reset(&mut self) -> String {
+        Llama3JsonToolParser::reset(self)
     }
 }
 
@@ -254,7 +253,7 @@ mod tests {
 
     use super::Llama3JsonToolParser;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
-    use crate::{ToolParseResult, ToolParser};
+    use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
 
     fn build_tool_call(function_name: &str, parameters: &str) -> String {
         format!(r#"{{"name":"{function_name}","parameters":{parameters}}}"#)
@@ -263,26 +262,28 @@ mod tests {
     #[test]
     fn llama_json_parse_complete_without_tool_call_keeps_text() {
         let mut parser = Llama3JsonToolParser::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 llama_json_passthrough_never_reenters_tool_parsing() {
         let mut parser = Llama3JsonToolParser::new(&test_tools());
-        let mut result = parser.push("plain text first ").unwrap();
-        result.append(
-            parser.push(&build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)).unwrap(),
+        let mut output = parser.parse_chunk("plain text first ").unwrap();
+        output.append(
+            parser
+                .parse_chunk(&build_tool_call("get_weather", r#"{"location":"Tokyo"}"#))
+                .unwrap(),
         );
-        result.append(parser.finish().unwrap());
+        output.append(parser.finish().unwrap());
 
         assert_eq!(
-            result.normal_text,
+            output.normal_text,
             r#"plain text first {"name":"get_weather","parameters":{"location":"Tokyo"}}"#
         );
-        assert!(result.calls.is_empty());
+        assert!(output.calls.is_empty());
     }
 
     #[test]
@@ -292,10 +293,10 @@ mod tests {
             "<|python_tag|>{}",
             build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
         );
-        let result = parser.parse_complete(&input).unwrap();
+        let output = parser.parse_complete(&input).unwrap();
 
-        assert_eq!(result.normal_text, input);
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, input);
+        assert!(output.calls.is_empty());
     }
 
     #[test]
@@ -305,22 +306,22 @@ mod tests {
             "\n  {}",
             build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
         );
-        let result = parser.parse_complete(&input).unwrap();
+        let output = parser.parse_complete(&input).unwrap();
 
-        assert_eq!(result.normal_text, input);
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, input);
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn llama_json_extracts_raw_parameters_object() {
         let mut parser = Llama3JsonToolParser::new(&test_tools());
         let arguments = r#"{ "location": "Tokyo", "days": 3 }"#;
-        let result = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
+        let output = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
 
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.calls[0].tool_index, 0);
-        assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(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]
@@ -344,10 +345,10 @@ mod tests {
             build_tool_call("get_weather", r#"{"location":"Shanghai"}"#),
             build_tool_call("add", r#"{"x":1,"y":2}"#),
         );
-        let result = parser.parse_complete(&input).unwrap();
+        let output = parser.parse_complete(&input).unwrap();
 
         expect![[r#"
-            ToolParseResult {
+            ToolParserOutput {
                 normal_text: "",
                 calls: [
                     ToolCallDelta {
@@ -367,7 +368,7 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 
     #[test]
@@ -380,23 +381,23 @@ mod tests {
             "}}",
         ];
 
-        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"}"#
         );
     }
@@ -411,16 +412,16 @@ mod tests {
         let chunks = split_by_chars(&input, 6);
         let mut parser = Llama3JsonToolParser::new(&test_tools());
 
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
-        assert_eq!(result.normal_text, "");
-        assert_eq!(result.calls.len(), 2);
+        assert_eq!(output.normal_text, "");
+        assert_eq!(output.calls.len(), 2);
         assert_eq!(
-            result.calls[0].arguments,
+            output.calls[0].arguments,
             r#"{"location":"Dallas","state":"TX"}"#
         );
-        assert_eq!(result.calls[1].name.as_deref(), Some("add"));
-        assert_eq!(result.calls[1].arguments, r#"{"x":4,"y":5}"#);
+        assert_eq!(output.calls[1].name.as_deref(), Some("add"));
+        assert_eq!(output.calls[1].arguments, r#"{"x":4,"y":5}"#);
     }
 
     #[test]
@@ -430,29 +431,29 @@ mod tests {
   "payload": {"items": [1, {"value": "literal { brace } and \"quote\""}]},
   "flag": true
 }"#;
-        let result = parser.parse_complete(&build_tool_call("convert", arguments)).unwrap();
+        let output = parser.parse_complete(&build_tool_call("convert", arguments)).unwrap();
 
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
     fn llama_json_keeps_trailing_whitespace_after_tool_call() {
         let mut parser = Llama3JsonToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&format!(
                 "{}\n\t ",
                 build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
             ))
             .unwrap();
 
-        assert_eq!(result.normal_text, "");
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.normal_text, "");
+        assert_eq!(output.calls.len(), 1);
     }
 
     #[test]
     fn llama_json_finish_fails_incomplete_tool_call() {
         let mut parser = Llama3JsonToolParser::new(&test_tools());
-        parser.push(r#"{"name":"get_weather","parameters":{"location""#).unwrap();
+        parser.parse_chunk(r#"{"name":"get_weather","parameters":{"location""#).unwrap();
 
         let error = parser.finish().unwrap_err();
 
@@ -463,7 +464,7 @@ mod tests {
     #[test]
     fn llama_json_malformed_field_order_fails_fast() {
         let mut parser = Llama3JsonToolParser::new(&test_tools());
-        let error = parser.push(r#"{"parameters":{},"name":"get_weather"}"#).unwrap_err();
+        let error = parser.parse_chunk(r#"{"parameters":{},"name":"get_weather"}"#).unwrap_err();
 
         expect![[r#"
             tool parser parsing failed: invalid Llama JSON
@@ -475,7 +476,7 @@ mod tests {
     fn llama_json_trailing_non_separator_content_errors() {
         let mut parser = Llama3JsonToolParser::new(&test_tools());
         let error = parser
-            .push(&format!(
+            .parse_chunk(&format!(
                 "{} trailing",
                 build_tool_call("get_weather", r#"{"location":"Tokyo"}"#)
             ))

rust/src/tool-parser/src/json/mistral.rs

@@ -1,5 +1,5 @@
 use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 const MISTRAL_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
     parser_name: "Mistral",
@@ -35,7 +35,6 @@ impl MistralToolParser {
 }
 
 impl ToolParser for MistralToolParser {
-    /// Create a boxed Mistral tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -43,15 +42,17 @@ impl ToolParser for MistralToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the Mistral parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.inner.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.inner.parse_into(chunk, output)
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.inner.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.inner.reset()
+    }
 }
 
 #[cfg(test)]
@@ -61,7 +62,7 @@ mod tests {
 
     use super::MistralToolParser;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
-    use crate::{ToolParseResult, ToolParser};
+    use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
 
     fn build_tool_call(function_name: &str, arguments: &str) -> String {
         format!(r#"{{"name":"{function_name}","arguments":{arguments}}}"#)
@@ -74,34 +75,34 @@ mod tests {
     #[test]
     fn mistral_parse_complete_without_tool_call_keeps_text() {
         let mut parser = MistralToolParser::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 mistral_parse_complete_extracts_raw_json_arguments() {
         let mut parser = MistralToolParser::new(&test_tools());
         let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
-        let result = parser
+        let output = parser
             .parse_complete(&format!(
                 "Let me check.\n{}",
                 build_tool_calls(&[build_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 mistral_parse_complete_extracts_pretty_multiple_tool_calls() {
         let mut parser = MistralToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(
                 r#"I'll help.
 [TOOL_CALLS] [
@@ -113,7 +114,7 @@ mod tests {
             .unwrap();
 
         expect![[r#"
-            ToolParseResult {
+            ToolParserOutput {
                 normal_text: "I'll help.\n",
                 calls: [
                     ToolCallDelta {
@@ -133,21 +134,21 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 
     #[test]
     fn mistral_does_not_validate_or_normalize_arguments() {
         let mut parser = MistralToolParser::new(&test_tools());
         let arguments = r#"{"location":"Tokyo",}"#;
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_calls(&[build_tool_call(
                 "get_weather",
                 arguments,
             )]))
             .unwrap();
 
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
@@ -162,24 +163,24 @@ mod tests {
             "}] suffix",
         ];
 
-        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.normal_text, "preface  suffix");
+        assert_eq!(output.normal_text, "preface  suffix");
         assert_eq!(
-            result.coalesce_calls().calls[0].arguments,
+            output.coalesce_calls().calls[0].arguments,
             r#"{"location":"Beijing"}"#
         );
     }
@@ -193,30 +194,30 @@ mod tests {
         let chunks = split_by_chars(&input, 5);
         let mut parser = MistralToolParser::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]
     fn mistral_keeps_array_bracket_literal_inside_json_string() {
         let mut parser = MistralToolParser::new(&test_tools());
         let arguments = r#"{"text":"Array notation: arr[0] = value[1]"}"#;
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_calls(&[build_tool_call("echo", arguments)]))
             .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]
     fn mistral_finish_fails_incomplete_tool_call() {
         let mut parser = MistralToolParser::new(&test_tools());
         parser
-            .push(r#"[TOOL_CALLS] [{"name":"get_weather","arguments":{"location""#)
+            .parse_chunk(r#"[TOOL_CALLS] [{"name":"get_weather","arguments":{"location""#)
             .unwrap();
 
         let error = parser.finish().unwrap_err();
@@ -229,7 +230,7 @@ mod tests {
     fn mistral_malformed_field_order_fails_fast() {
         let mut parser = MistralToolParser::new(&test_tools());
         let error = parser
-            .push(r#"[TOOL_CALLS] [{"arguments":{},"name":"get_weather"}]"#)
+            .parse_chunk(r#"[TOOL_CALLS] [{"arguments":{},"name":"get_weather"}]"#)
             .unwrap_err();
 
         expect![[r#"

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

@@ -20,7 +20,7 @@ use winnow::token::literal;
 use super::utils::{
     JsonObjectScanState, json_str, parse_buffered_event, safe_text_len, take_json_object,
 };
-use super::{Result, ToolCallDelta, ToolParseResult};
+use super::{Result, ToolCallDelta, ToolParserOutput};
 
 type JsonToolInput<'i> = Partial<&'i str>;
 
@@ -80,27 +80,24 @@ impl JsonToolCallParser {
         }
     }
 
-    /// Push one decoded text chunk through the JSON tool-call parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
         let config = self.config;
 
         while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
             parse_next_json_tool_call_event(input, &mut self.mode, config)
         })? {
-            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 {
-            JsonToolCallMode::Text => result.normal_text.push_str(&self.buffer),
+            JsonToolCallMode::Text => output.normal_text.push_str(&self.buffer),
             JsonToolCallMode::Header | JsonToolCallMode::Arguments { .. } => {
                 return Err(parsing_failed!(
                     "incomplete {} tool call",
@@ -108,19 +105,19 @@ impl JsonToolCallParser {
                 ));
             }
         }
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
     }
 
     /// Apply one parsed JSON tool-call event to parser state and output.
     fn apply_event(
         &mut self,
         event: JsonToolCallEvent,
-        result: &mut ToolParseResult,
+        output: &mut ToolParserOutput,
     ) -> Result<()> {
         match event {
             JsonToolCallEvent::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
             }
             JsonToolCallEvent::ToolCallStart => self.mode = JsonToolCallMode::Header,
             JsonToolCallEvent::ToolCallHeader { function_name } => {
@@ -130,7 +127,7 @@ impl JsonToolCallParser {
                 self.mode = JsonToolCallMode::Arguments {
                     json_scan: JsonObjectScanState::default(),
                 };
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index,
                     name: Some(function_name),
                     arguments: String::new(),
@@ -143,7 +140,7 @@ impl JsonToolCallParser {
                         self.config.parser_name
                     ));
                 };
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index,
                     name: None,
                     arguments: self.buffer[..consumed_len].to_string(),
@@ -161,12 +158,11 @@ impl JsonToolCallParser {
         Ok(())
     }
 
-    /// Reset all streaming state.
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.mode = JsonToolCallMode::Text;
         self.active_tool_index = None;
         self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
     }
 }
 
@@ -336,7 +332,7 @@ mod tests {
     use expect_test::expect;
 
     use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
-    use crate::ToolParseResult;
+    use crate::ToolParserOutput;
 
     const DELIMITED_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
         parser_name: "Delimited JSON",
@@ -356,13 +352,13 @@ mod tests {
         format!("<tool_calls>{}</tool_calls>", tool_calls.join(" <\n"))
     }
 
-    fn collect_chunks(parser: &mut JsonToolCallParser, chunks: &[&str]) -> ToolParseResult {
-        let mut result = ToolParseResult::default();
+    fn collect_chunks(parser: &mut JsonToolCallParser, chunks: &[&str]) -> ToolParserOutput {
+        let mut output = ToolParserOutput::default();
         for chunk in chunks {
-            result.append(parser.push(chunk).unwrap());
+            parser.parse_into(chunk, &mut output).unwrap();
         }
-        result.append(parser.finish().unwrap());
-        result.coalesce_calls()
+        output.append(parser.finish().unwrap());
+        output.coalesce_calls()
     }
 
     #[test]
@@ -373,10 +369,10 @@ mod tests {
         ]);
         let mut parser = JsonToolCallParser::new(DELIMITED_CONFIG);
 
-        let result = collect_chunks(&mut parser, &[&input]);
+        let output = collect_chunks(&mut parser, &[&input]);
 
         expect![[r#"
-            ToolParseResult {
+            ToolParserOutput {
                 normal_text: "",
                 calls: [
                     ToolCallDelta {
@@ -396,7 +392,7 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 
     #[test]
@@ -409,10 +405,10 @@ mod tests {
             "</tool_calls>",
         ];
 
-        let result = collect_chunks(&mut parser, &chunks);
+        let output = collect_chunks(&mut parser, &chunks);
 
         expect![[r#"
-            ToolParseResult {
+            ToolParserOutput {
                 normal_text: "",
                 calls: [
                     ToolCallDelta {
@@ -432,7 +428,7 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 
     #[test]
@@ -444,10 +440,10 @@ mod tests {
             "</tool_calls> trailing text",
         ];
 
-        let result = collect_chunks(&mut parser, &chunks);
+        let output = collect_chunks(&mut parser, &chunks);
 
         expect![[r#"
-            ToolParseResult {
+            ToolParserOutput {
                 normal_text: " trailing text",
                 calls: [
                     ToolCallDelta {
@@ -460,6 +456,6 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 }

rust/src/tool-parser/src/json/qwen.rs

@@ -1,5 +1,5 @@
 use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
-use crate::{Result, Tool, ToolParseResult, ToolParser};
+use crate::{Result, Tool, ToolParser, ToolParserOutput};
 
 const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
     parser_name: "Qwen XML",
@@ -40,7 +40,6 @@ impl Qwen3XmlToolParser {
 }
 
 impl ToolParser for Qwen3XmlToolParser {
-    /// Create a boxed Qwen XML tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -48,15 +47,17 @@ impl ToolParser for Qwen3XmlToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the Qwen XML parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
-        self.inner.push(chunk)
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+        self.inner.parse_into(chunk, output)
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
+    fn finish(&mut self) -> Result<ToolParserOutput> {
         self.inner.finish()
     }
+
+    fn reset(&mut self) -> String {
+        self.inner.reset()
+    }
 }
 
 #[cfg(test)]
@@ -66,7 +67,7 @@ mod tests {
 
     use super::Qwen3XmlToolParser;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
-    use crate::{ToolParseResult, ToolParser};
+    use crate::{ToolParser, ToolParserOutput, ToolParserTestExt as _};
 
     fn build_tool_call(function_name: &str, arguments: &str) -> String {
         format!(
@@ -77,37 +78,37 @@ mod tests {
     #[test]
     fn qwen_xml_parse_complete_without_tool_call_keeps_text() {
         let mut parser = Qwen3XmlToolParser::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 qwen_xml_parse_complete_extracts_raw_json_arguments() {
         let mut parser = Qwen3XmlToolParser::new(&test_tools());
         let arguments = r#"{ "location": "Tokyo", "days": "3" }"#;
-        let result = parser
+        let output = parser
             .parse_complete(&format!(
                 "Let me check.\n{}",
                 build_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 qwen_xml_does_not_validate_or_normalize_arguments() {
         let mut parser = Qwen3XmlToolParser::new(&test_tools());
         let arguments = r#"{"location":"Tokyo",}"#;
-        let result = parser.parse_complete(&build_tool_call("get_weather", arguments)).unwrap();
+        let output = parser.parse_complete(&build_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 {
             "}\n</tool_call>",
         ];
 
-        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,27 +153,27 @@ mod tests {
         let chunks = split_by_chars(&input, 5);
         let mut parser = Qwen3XmlToolParser::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]
     fn qwen_xml_keeps_end_marker_literal_inside_json_string() {
         let mut parser = Qwen3XmlToolParser::new(&test_tools());
         let arguments = r#"{"text":"literal </tool_call> inside"}"#;
-        let result = parser.parse_complete(&build_tool_call("echo", arguments)).unwrap();
+        let output = parser.parse_complete(&build_tool_call("echo", arguments)).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]
     fn qwen_xml_decodes_escaped_function_name() {
         let mut parser = Qwen3XmlToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(
                 r#"<tool_call>
 {"name":"say_\"hi","arguments":{}}
@@ -180,7 +181,7 @@ mod tests {
             )
             .unwrap();
 
-        assert_eq!(result.calls[0].name.as_deref(), Some("say_\"hi"));
+        assert_eq!(output.calls[0].name.as_deref(), Some("say_\"hi"));
     }
 
     #[test]
@@ -189,10 +190,10 @@ mod tests {
         let input = r#"<tool_call>{"name":"get_weather","arguments":{}}
 </tool_call>"#;
 
-        let result = parser.parse_complete(input).unwrap();
+        let output = parser.parse_complete(input).unwrap();
 
-        assert_eq!(result.normal_text, input);
-        assert!(result.calls.is_empty());
+        assert_eq!(output.normal_text, input);
+        assert!(output.calls.is_empty());
     }
 
     #[test]
@@ -218,10 +219,10 @@ mod tests {
         let chunks = split_by_chars(&input, 7);
         let mut parser = Qwen3XmlToolParser::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 {
@@ -241,14 +242,14 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 
     #[test]
     fn qwen_xml_finish_fails_incomplete_tool_call() {
         let mut parser = Qwen3XmlToolParser::new(&test_tools());
         parser
-            .push(
+            .parse_chunk(
                 r#"<tool_call>
 {"name":"get_weather","arguments":{"location""#,
             )
@@ -264,7 +265,7 @@ mod tests {
     fn qwen_xml_malformed_field_order_fails_fast() {
         let mut parser = Qwen3XmlToolParser::new(&test_tools());
         let error = parser
-            .push(
+            .parse_chunk(
                 r#"<tool_call>
 {"arguments":{},"name":"get_weather"}
 </tool_call>"#,

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

@@ -5,7 +5,7 @@ use winnow::stream::Partial;
 use winnow::token::{literal, rest, take_until, take_while};
 
 use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
-use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
+use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
 use crate::Tool;
 
 const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
@@ -73,10 +73,10 @@ impl KimiK2ToolParser {
     }
 
     /// Apply one parsed Kimi K2 event to parser state and output.
-    fn apply_event(&mut self, event: KimiK2Event, result: &mut ToolParseResult) -> Result<()> {
+    fn apply_event(&mut self, event: KimiK2Event, output: &mut ToolParserOutput) -> Result<()> {
         match event {
             KimiK2Event::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
             }
             KimiK2Event::ToolCallsStart => self.mode = KimiK2Mode::ToolBlock,
             KimiK2Event::ToolCallStart => self.mode = KimiK2Mode::Header,
@@ -89,7 +89,7 @@ impl KimiK2ToolParser {
                 self.mode = KimiK2Mode::Arguments {
                     json_scan: JsonObjectScanState::default(),
                 };
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index,
                     name: Some(function_name),
                     arguments: String::new(),
@@ -101,7 +101,7 @@ impl KimiK2ToolParser {
                         "Kimi K2 arguments without an active tool call"
                     ));
                 };
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index,
                     name: None,
                     arguments: self.buffer[..consumed_len].to_string(),
@@ -120,16 +120,14 @@ impl KimiK2ToolParser {
         Ok(())
     }
 
-    /// Reset all streaming state.
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.mode = KimiK2Mode::Text;
         self.active_tool_index = None;
+        std::mem::take(&mut self.buffer)
     }
 }
 
 impl ToolParser for KimiK2ToolParser {
-    /// Create a boxed Kimi K2 tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -137,38 +135,38 @@ impl ToolParser for KimiK2ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Preserve Kimi K2 special-token markers while decoding.
     fn preserve_special_tokens(&self) -> bool {
         true
     }
 
-    /// Push one decoded text chunk through the Kimi K2 parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
         while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
             parse_next_kimi_k2_event(input, &mut self.mode)
         })? {
-            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 {
-            KimiK2Mode::Text => result.normal_text.push_str(&self.buffer),
+            KimiK2Mode::Text => output.normal_text.push_str(&self.buffer),
             KimiK2Mode::ToolBlock | KimiK2Mode::Done => {}
             KimiK2Mode::Header | KimiK2Mode::Arguments { .. } => {
                 return Err(parsing_failed!("incomplete Kimi K2 tool call"));
             }
         }
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
+    }
+
+    fn reset(&mut self) -> String {
+        KimiK2ToolParser::reset(self)
     }
 }
 
@@ -323,8 +321,8 @@ mod tests {
         KimiK2ToolParser, TOOL_CALL_ARGUMENT_START, TOOL_CALL_END, TOOL_CALL_START, TOOL_CALLS_END,
         TOOL_CALLS_START, ToolParser, tool_header,
     };
-    use crate::ToolParseResult;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
+    use crate::{ToolParserOutput, ToolParserTestExt as _};
 
     fn build_tool_call(function_name: &str, index: usize, arguments: &str) -> String {
         format!(
@@ -339,35 +337,35 @@ mod tests {
     #[test]
     fn kimi_k2_parse_complete_without_tool_call_keeps_text() {
         let mut parser = KimiK2ToolParser::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 kimi_k2_parse_complete_extracts_raw_json_arguments() {
         let mut parser = KimiK2ToolParser::new(&test_tools());
         let arguments = r#"{ "location": "NYC", "days": "3" }"#;
-        let result = parser
+        let output = parser
             .parse_complete(&format!(
                 "Checking. {} trailing text",
                 build_tool_section(&[build_tool_call("get_weather", 0, arguments)])
             ))
             .unwrap();
 
-        assert_eq!(result.normal_text, "Checking. ");
-        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, "Checking. ");
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].tool_index, 0);
+        assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
     fn kimi_k2_does_not_validate_or_normalize_arguments() {
         let mut parser = KimiK2ToolParser::new(&test_tools());
         let arguments = r#"{"location":"NYC",}"#;
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_section(&[build_tool_call(
                 "get_weather",
                 0,
@@ -375,7 +373,7 @@ mod tests {
             )]))
             .unwrap();
 
-        assert_eq!(result.calls[0].arguments, arguments);
+        assert_eq!(output.calls[0].arguments, arguments);
     }
 
     #[test]
@@ -393,23 +391,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\":", "\"Paris\"", "}"]);
-        let result = result.coalesce_calls();
-        assert_eq!(result.calls[0].arguments, r#"{"location":"Paris"}"#);
+        let output = output.coalesce_calls();
+        assert_eq!(output.calls[0].arguments, r#"{"location":"Paris"}"#);
     }
 
     #[test]
@@ -427,11 +425,11 @@ mod tests {
             TOOL_CALLS_END,
         ];
 
-        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":"NYC"}"#);
+        assert_eq!(output.normal_text, "hello ");
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].arguments, r#"{"location":"NYC"}"#);
     }
 
     #[test]
@@ -440,10 +438,10 @@ mod tests {
         let arguments = format!(r#"{{"text":"literal {TOOL_CALL_END} inside"}}"#);
         let input = build_tool_section(&[build_tool_call("echo", 0, &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]
@@ -460,11 +458,11 @@ mod tests {
             TOOL_CALLS_END,
         ];
 
-        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!(
-            result.calls[0].arguments,
+            output.calls[0].arguments,
             r#"{"text":"literal <|tool_call_end|> inside"}"#
         );
     }
@@ -478,10 +476,10 @@ mod tests {
         ]);
 
         let chunks = split_by_chars(&input, 7);
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
         expect![[r#"
-            ToolParseResult {
+            ToolParserOutput {
                 normal_text: "",
                 calls: [
                     ToolCallDelta {
@@ -501,7 +499,7 @@ mod tests {
                 ],
             }
         "#]]
-        .assert_debug_eq(&result);
+        .assert_debug_eq(&output);
     }
 
     #[test]
@@ -511,11 +509,11 @@ mod tests {
             "{TOOL_CALLS_START}{TOOL_CALL_START}api.tools.search:42{TOOL_CALL_ARGUMENT_START}{{}}{TOOL_CALL_END}{TOOL_CALLS_END}"
         );
 
-        let result = parser.parse_complete(&input).unwrap();
+        let output = parser.parse_complete(&input).unwrap();
 
-        assert_eq!(result.calls[0].tool_index, 42);
-        assert_eq!(result.calls[0].name.as_deref(), Some("search"));
-        assert_eq!(result.calls[0].arguments, "{}");
+        assert_eq!(output.calls[0].tool_index, 42);
+        assert_eq!(output.calls[0].name.as_deref(), Some("search"));
+        assert_eq!(output.calls[0].arguments, "{}");
     }
 
     #[test]
@@ -536,7 +534,7 @@ mod tests {
     fn kimi_k2_finish_fails_incomplete_tool_call() {
         let mut parser = KimiK2ToolParser::new(&test_tools());
         parser
-            .push(&format!(
+            .parse_chunk(&format!(
                 "{TOOL_CALLS_START}{TOOL_CALL_START}functions.get_weather:0{TOOL_CALL_ARGUMENT_START}{{\"location\""
             ))
             .unwrap();
@@ -553,7 +551,7 @@ mod tests {
         let input =
             format!("{TOOL_CALLS_START}{TOOL_CALL_START}get_weather{TOOL_CALL_ARGUMENT_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/lib.rs

@@ -51,20 +51,20 @@ pub struct ToolCallDelta {
 
 /// Result of advancing tool parsing with one assistant-text input.
 #[derive(Debug, Clone, Default, PartialEq, Eq)]
-pub struct ToolParseResult {
+pub struct ToolParserOutput {
     /// Plain assistant text that is not part of any tool call.
     pub normal_text: String,
     /// Tool-call updates extracted from this input.
     pub calls: Vec<ToolCallDelta>,
 }
 
-impl ToolParseResult {
-    /// Append another parser result onto this one.
+impl ToolParserOutput {
+    /// Append another parser output onto this one.
     ///
     /// Note that this does not attempt to merge multiple deltas for the same
     /// tool call into one complete item. Call `coalesce_calls()` after if
     /// that behavior is desired.
-    pub(crate) fn append(&mut self, mut other: Self) {
+    pub fn append(&mut self, mut other: Self) {
         self.normal_text.push_str(&other.normal_text);
         self.calls.append(&mut other.calls);
     }
@@ -75,7 +75,7 @@ impl ToolParseResult {
     /// which delegates through the incremental parser lifecycle and then
     /// needs to collapse streaming-style argument fragments into one final
     /// tool call.
-    pub(crate) fn coalesce_calls(mut self) -> Self {
+    pub fn coalesce_calls(mut self) -> Self {
         let mut merged = BTreeMap::<usize, ToolCallDelta>::new();
         let mut order = Vec::new();
 
@@ -116,24 +116,55 @@ pub trait ToolParser: Send {
         false
     }
 
-    /// Feed one decoded text delta into the parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult>;
+    /// Feed one decoded text delta into the parser, appending committed output
+    /// into `output`.
+    ///
+    /// If this returns an error, any output already appended to `output`
+    /// remains committed parser output. The parser must keep its uncommitted
+    /// buffer intact so callers may recover it with `reset()`.
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()>;
 
     /// Flush any buffered partial state at end of stream.
-    fn finish(&mut self) -> Result<ToolParseResult> {
-        Ok(ToolParseResult::default())
+    ///
+    /// This operation is atomic: on error no partial output is returned and the
+    /// parser's buffered state is left intact.
+    fn finish(&mut self) -> Result<ToolParserOutput>;
+
+    /// Clear parser state and return currently uncommitted buffered text.
+    ///
+    /// Callers may use this to recover any text that failed to parse after an error
+    /// and output it as normal text.
+    fn reset(&mut self) -> String;
+}
+
+/// Extension methods for easily testing `ToolParser` implementations.
+///
+/// These helpers do not handle partial parsing or error recovery, so they are
+/// not intended for use in production code paths.
+#[cfg(any(test, feature = "test-util"))]
+#[easy_ext::ext(ToolParserTestExt)]
+impl<T: ToolParser + ?Sized> T {
+    /// Feed one decoded text delta and return only if the whole chunk parses.
+    ///
+    /// If parsing fails, partial committed output is discarded by this helper.
+    /// Prefer `parse_into` for more fine-grained control in error recovery.
+    pub fn parse_chunk(&mut self, chunk: &str) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
+        self.parse_into(chunk, &mut output)?;
+        Ok(output)
     }
 
     /// Parse complete tool calls from final output.
     ///
-    /// The default implementation reuses the incremental parser lifecycle by
-    /// feeding the full output through `push()` and then calling `finish()`.
-    /// This keeps one source of truth for robust parsers whose incremental
-    /// state machine is equivalent across arbitrary chunking.
-    fn parse_complete(&mut self, output: &str) -> Result<ToolParseResult> {
-        let mut result = self.push(output)?;
-        result.append(self.finish()?);
-        Ok(result.coalesce_calls())
+    /// This default implementation reuses the incremental parser lifecycle by
+    /// feeding the full output through `parse_chunk()` and then calling `finish()`.
+    ///
+    /// If parsing fails, partial committed output is discarded by this helper.
+    /// Prefer `parse_into` for more fine-grained control in error recovery.
+    pub fn parse_complete(&mut self, text: &str) -> Result<ToolParserOutput> {
+        let mut output = self.parse_chunk(text)?;
+        output.append(self.finish()?);
+        Ok(output.coalesce_calls())
     }
 }
 

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

@@ -1,12 +1,12 @@
 use winnow::ascii::{multispace0 as ws0, multispace1 as ws1};
-use winnow::combinator::{alt, delimited, repeat, seq, terminated};
+use winnow::combinator::{alt, delimited, eof, repeat, seq, terminated};
 use winnow::prelude::*;
 use winnow::stream::Partial;
 use winnow::token::{literal, rest, take_until};
 
 use super::parameters::ToolSchemas;
 use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
-use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
+use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
 use crate::Tool;
 
 const TOOL_CALL_START: &str = "<minimax:tool_call>";
@@ -69,10 +69,10 @@ impl MinimaxM2ToolParser {
     }
 
     /// Apply one parsed MiniMax M2 event to parser state and output.
-    fn apply_event(&mut self, event: MinimaxM2Event, result: &mut ToolParseResult) -> Result<()> {
+    fn apply_event(&mut self, event: MinimaxM2Event, output: &mut ToolParserOutput) -> Result<()> {
         match event {
             MinimaxM2Event::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
             }
             MinimaxM2Event::ToolBlockStart => self.mode = MinimaxM2Mode::ToolBlock,
             MinimaxM2Event::Invoke { name, raw_params } => {
@@ -80,7 +80,7 @@ impl MinimaxM2ToolParser {
                 let arguments = serde_json::to_string(&arguments)
                     .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
 
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index: self.emitted_tool_count,
                     name: Some(name),
                     arguments,
@@ -93,16 +93,14 @@ impl MinimaxM2ToolParser {
         Ok(())
     }
 
-    /// Reset all streaming state.
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.mode = MinimaxM2Mode::Text;
         self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
     }
 }
 
 impl ToolParser for MinimaxM2ToolParser {
-    /// Create a boxed MiniMax M2 tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -110,35 +108,36 @@ impl ToolParser for MinimaxM2ToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the MiniMax M2 parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
         while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
             parse_next_minimax_m2_event(input, self.mode)
         })? {
-            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 {
             MinimaxM2Mode::Text => {
-                result.normal_text.push_str(&self.buffer);
+                output.normal_text.push_str(&self.buffer);
             }
             MinimaxM2Mode::ToolBlock => {
                 return Err(parsing_failed!("incomplete MiniMax M2 tool call"));
             }
             MinimaxM2Mode::Done => {}
         }
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
+    }
+
+    fn reset(&mut self) -> String {
+        MinimaxM2ToolParser::reset(self)
     }
 }
 
@@ -183,16 +182,17 @@ fn tool_block_end_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2
 
 /// Parse a complete MiniMax M2 invoke block.
 fn invoke_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
-    let (name, raw_params) = seq!(
+    let (name, body) = seq!(
         _: ws0,
         _: literal(INVOKE_START),
         _: (ws1, literal("name=")),
-        attr_value,
+        partial_attr_value,
         _: literal(">"),
-        repeat(0.., terminated(parameter, ws0)),
+        take_until(0.., INVOKE_END),
         _: literal(INVOKE_END),
     )
     .parse_next(input)?;
+    let raw_params = parse_invoke_params(body)?;
 
     Ok(MinimaxM2Event::Invoke {
         name: name.trim().to_string(),
@@ -200,8 +200,14 @@ fn invoke_event(input: &mut MinimaxM2Input<'_>) -> ModalResult<MinimaxM2Event> {
     })
 }
 
+/// Parse all parameter blocks inside a complete MiniMax M2 invoke body.
+fn parse_invoke_params(invoke_body: &str) -> ModalResult<Vec<(String, String)>> {
+    let mut input = invoke_body;
+    delimited(ws0, repeat(0.., terminated(parameter, ws0)), eof).parse_next(&mut input)
+}
+
 /// Parse a MiniMax M2 parameter block.
-fn parameter(input: &mut MinimaxM2Input<'_>) -> ModalResult<(String, String)> {
+fn parameter(input: &mut &str) -> ModalResult<(String, String)> {
     let (name, value) = seq!(
         _: literal(PARAMETER_START),
         _: (ws1, literal("name=")),
@@ -216,7 +222,17 @@ fn parameter(input: &mut MinimaxM2Input<'_>) -> ModalResult<(String, String)> {
 }
 
 /// Parse a quoted or unquoted XML attribute value.
-fn attr_value<'i>(input: &mut MinimaxM2Input<'i>) -> ModalResult<&'i str> {
+fn attr_value<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
+    alt((
+        delimited(literal("\""), take_until(1.., "\""), literal("\"")),
+        delimited(literal("'"), take_until(1.., "'"), literal("'")),
+        take_until(1.., ">"),
+    ))
+    .parse_next(input)
+}
+
+/// Parse a quoted or unquoted XML attribute value from partial streaming input.
+fn partial_attr_value<'i>(input: &mut MinimaxM2Input<'i>) -> ModalResult<&'i str> {
     alt((
         delimited(literal("\""), take_until(1.., "\""), literal("\"")),
         delimited(literal("'"), take_until(1.., "'"), literal("'")),
@@ -237,6 +253,7 @@ mod tests {
     use thiserror_ext::AsReport;
 
     use super::{MinimaxM2ToolParser, TOOL_CALL_END, TOOL_CALL_START, ToolParser};
+    use crate::ToolParserTestExt as _;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
 
     fn build_tool_block(invokes: &[(&str, Vec<(&str, &str)>)]) -> String {
@@ -257,27 +274,27 @@ mod tests {
     #[test]
     fn minimax_m2_parse_complete_without_tool_call_keeps_text() {
         let mut parser = MinimaxM2ToolParser::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 minimax_m2_parse_complete_extracts_single_tool_call() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_block(&[(
                 "get_weather",
                 vec![("city", "Seattle"), ("days", "5")],
             )]))
             .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!({ "city": "Seattle", "days": 5 })
         );
     }
@@ -289,31 +306,31 @@ mod tests {
             "Let me check. {} This trailing text is ignored.",
             build_tool_block(&[("get_weather", vec![("city", "Seattle")])])
         );
-        let result = parser.parse_complete(&output).unwrap();
+        let output = parser.parse_complete(&output).unwrap();
 
-        assert_eq!(result.normal_text, "Let me check. ");
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.normal_text, "Let me check. ");
+        assert_eq!(output.calls.len(), 1);
     }
 
     #[test]
     fn minimax_m2_parse_complete_extracts_multiple_invokes() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_block(&[
                 ("get_weather", vec![("city", "Seattle")]),
                 ("get_weather", vec![("city", "NYC")]),
             ]))
             .unwrap();
 
-        assert_eq!(result.calls.len(), 2);
-        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].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!({ "city": "Seattle" })
         );
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
             json!({ "city": "NYC" })
         );
     }
@@ -321,7 +338,7 @@ mod tests {
     #[test]
     fn minimax_m2_parse_complete_converts_schema_types() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_block(&[(
                 "convert",
                 vec![
@@ -335,7 +352,7 @@ mod tests {
             .unwrap();
 
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "whole": 5.0,
                 "flag": true,
@@ -349,7 +366,7 @@ mod tests {
     #[test]
     fn minimax_m2_parse_complete_unescapes_literal_closing_tags_in_parameter_value() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_block(&[(
                 "get_weather",
                 vec![
@@ -363,7 +380,7 @@ mod tests {
             .unwrap();
 
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "city": "Seattle </parameter></invoke></minimax:tool_call>",
                 "days": 5,
@@ -374,7 +391,7 @@ mod tests {
     #[test]
     fn minimax_m2_parse_complete_handles_multiline_parameters() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(
                 "<minimax:tool_call>\
                  <invoke name=\"calculate_area\">\
@@ -387,7 +404,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!({
                 "shape": "\nrectangle\n",
                 "dimensions": { "width": 10, "height": 20 },
@@ -399,7 +416,7 @@ mod tests {
     #[test]
     fn minimax_m2_streaming_extracts_single_tool_call() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "<minimax:tool_call>",
@@ -409,11 +426,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!({ "city": "Seattle" })
         );
     }
@@ -421,7 +438,7 @@ mod tests {
     #[test]
     fn minimax_m2_streaming_preserves_prefix_text() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "Let me check. ",
@@ -431,17 +448,17 @@ mod tests {
             ],
         );
 
-        assert_eq!(result.normal_text, "Let me check. ");
-        assert_eq!(result.calls.len(), 1);
+        assert_eq!(output.normal_text, "Let me check. ");
+        assert_eq!(output.calls.len(), 1);
     }
 
     #[test]
     fn minimax_m2_streaming_without_tool_call_emits_text_incrementally() {
         let mut parser = MinimaxM2ToolParser::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]
@@ -449,10 +466,10 @@ mod tests {
         let text = build_tool_block(&[("get_weather", vec![("city", "Seattle")])]);
         let chunks = split_by_chars(&text, 3);
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
-        assert_eq!(result.calls.len(), 1);
-        assert!(result.normal_text.is_empty());
+        assert_eq!(output.calls.len(), 1);
+        assert!(output.normal_text.is_empty());
     }
 
     #[test]
@@ -463,11 +480,36 @@ mod tests {
         ]);
         let chunks = split_by_chars(&text, 7);
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
+        let output = collect_stream(&mut parser, &chunks);
+
+        assert_eq!(output.calls.len(), 2);
+        assert_eq!(output.calls[0].tool_index, 0);
+        assert_eq!(output.calls[1].tool_index, 1);
+    }
+
+    #[test]
+    fn minimax_m2_streaming_handles_template_whitespace_and_split_parameters() {
+        let text = concat!(
+            "I will call the tools.\n",
+            "<minimax:tool_call>\n",
+            "<invoke name=\"get_weather\">\n",
+            "<parameter name=\"city\">Seattle</parameter>\n",
+            "</invoke>\n",
+            "<invoke name=\"get_weather\">\n",
+            "<parameter name=\"city\">NYC</parameter>\n",
+            "</invoke>\n",
+            "</minimax:tool_call>",
+        );
+        let chunks = split_by_chars(text, 7);
+        let mut parser = MinimaxM2ToolParser::new(&test_tools());
         let result = collect_stream(&mut parser, &chunks);
 
+        assert_eq!(result.normal_text, "I will call the tools.\n");
         assert_eq!(result.calls.len(), 2);
         assert_eq!(result.calls[0].tool_index, 0);
+        assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
         assert_eq!(result.calls[1].tool_index, 1);
+        assert_eq!(result.calls[1].name.as_deref(), Some("get_weather"));
     }
 
     #[test]
@@ -478,25 +520,26 @@ mod tests {
         );
         let chunks = split_by_chars(&text, 5);
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
-        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 minimax_m2_streaming_does_not_emit_incomplete_tool_call() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let result = parser.push(r#"<minimax:tool_call><invoke name="get_weather">"#).unwrap();
+        let output =
+            parser.parse_chunk(r#"<minimax:tool_call><invoke name="get_weather">"#).unwrap();
 
-        assert!(result.normal_text.is_empty());
-        assert!(result.calls.is_empty());
+        assert!(output.normal_text.is_empty());
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn minimax_m2_finish_fails_incomplete_tool_call() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        parser.push(r#"<minimax:tool_call><invoke name="get_weather">"#).unwrap();
+        parser.parse_chunk(r#"<minimax:tool_call><invoke name="get_weather">"#).unwrap();
 
         assert!(parser.finish().is_err());
     }
@@ -504,7 +547,7 @@ mod tests {
     #[test]
     fn minimax_m2_finish_fails_after_bare_tool_block_start() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        parser.push("<minimax:tool_call>").unwrap();
+        parser.parse_chunk("<minimax:tool_call>").unwrap();
 
         assert!(parser.finish().is_err());
     }
@@ -512,7 +555,7 @@ mod tests {
     #[test]
     fn minimax_m2_malformed_tool_call_fails_fast() {
         let mut parser = MinimaxM2ToolParser::new(&test_tools());
-        let error = parser.push("<minimax:tool_call><bad></minimax:tool_call>").unwrap_err();
+        let error = parser.parse_chunk("<minimax:tool_call><bad></minimax:tool_call>").unwrap_err();
 
         expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
     }

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

@@ -224,10 +224,11 @@ fn convert_value(param_type: &JsonParamType, value: &str) -> Option<Value> {
 
 /// Convert one raw string value to a JSON number.
 fn convert_number(value: &str) -> Option<Value> {
-    if let Ok(parsed) = value.parse::<i64>() {
-        return Some(Value::Number(Number::from(parsed)));
-    }
-    Number::from_f64(value.parse::<f64>().ok()?).map(Value::Number)
+    serde_json::from_str::<Number>(value)
+        .or_else(|_| value.parse::<i64>().map(Number::from))
+        .or_else(|_| value.parse::<f64>().ok().and_then(Number::from_f64).ok_or(()))
+        .ok()
+        .map(Value::Number)
 }
 
 /// Convert one raw string value to a boolean.
@@ -304,7 +305,7 @@ mod tests {
     }
 
     #[test]
-    fn number_conversion_parses_int_then_float() {
+    fn number_conversion_preserves_json_number_spelling_with_legacy_fallback() {
         let params = ToolSchema::from_schema(&json!({
             "type": "object",
             "properties": {
@@ -312,17 +313,23 @@ mod tests {
             }
         }));
 
-        assert_eq!(params.convert("value", "5"), json!(5));
-        assert_eq!(params.convert("value", "5.0"), json!(5.0));
-        assert_eq!(params.convert("value", "5."), json!(5.0));
-        assert_eq!(params.convert("value", "+1"), json!(1));
-        assert_eq!(params.convert("value", "+1.0"), json!(1.0));
+        assert_eq!(converted_number_text(&params, "5"), "5");
+        assert_eq!(converted_number_text(&params, "5.0"), "5.0");
+        assert_eq!(converted_number_text(&params, "5.00"), "5.00");
+        assert_eq!(converted_number_text(&params, "1e0"), "1e+0");
+        assert_eq!(converted_number_text(&params, "5."), "5.0");
+        assert_eq!(converted_number_text(&params, "+1"), "1");
+        assert_eq!(converted_number_text(&params, "+1.0"), "1.0");
         assert_eq!(
-            params.convert("value", "9223372036854775807.5"),
-            json!(9223372036854775808.0)
+            converted_number_text(&params, "9223372036854775807.5"),
+            "9223372036854775807.5"
         );
     }
 
+    fn converted_number_text(params: &ToolSchema, value: &str) -> String {
+        serde_json::to_string(&params.convert("value", value)).unwrap()
+    }
+
     #[test]
     fn converts_upstream_aliases() {
         let params = ToolSchema::from_schema(&json!({

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

@@ -6,7 +6,7 @@ use winnow::token::{literal, take_until};
 
 use super::parameters::ToolSchemas;
 use super::utils::{parse_buffered_event, safe_text_len, xml_unescape};
-use super::{Result, ToolCallDelta, ToolParseResult, ToolParser};
+use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
 use crate::Tool;
 
 const TOOL_CALL_START: &str = "<tool_call>";
@@ -71,10 +71,10 @@ impl Qwen3CoderToolParser {
     }
 
     /// Apply one parsed Qwen Coder event to parser state and output.
-    fn apply_event(&mut self, event: QwenCoderEvent, result: &mut ToolParseResult) -> Result<()> {
+    fn apply_event(&mut self, event: QwenCoderEvent, output: &mut ToolParserOutput) -> Result<()> {
         match event {
             QwenCoderEvent::Text { len: consumed_len } => {
-                result.normal_text.push_str(&self.buffer[..consumed_len]);
+                output.normal_text.push_str(&self.buffer[..consumed_len]);
             }
             QwenCoderEvent::ToolCallStart => self.mode = QwenCoderMode::ToolCall,
             QwenCoderEvent::ToolCall { name, raw_params } => {
@@ -83,7 +83,7 @@ impl Qwen3CoderToolParser {
                 let arguments = serde_json::to_string(&arguments)
                     .map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
 
-                result.calls.push(ToolCallDelta {
+                output.calls.push(ToolCallDelta {
                     tool_index: self.emitted_tool_count,
                     name: Some(name),
                     arguments,
@@ -94,16 +94,14 @@ impl Qwen3CoderToolParser {
         Ok(())
     }
 
-    /// Reset all streaming state.
-    fn reset(&mut self) {
-        self.buffer.clear();
+    fn reset(&mut self) -> String {
         self.mode = QwenCoderMode::Text;
         self.emitted_tool_count = 0;
+        std::mem::take(&mut self.buffer)
     }
 }
 
 impl ToolParser for Qwen3CoderToolParser {
-    /// Create a boxed Qwen Coder tool parser.
     fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
     where
         Self: Sized + 'static,
@@ -111,32 +109,33 @@ impl ToolParser for Qwen3CoderToolParser {
         Ok(Box::new(Self::new(tools)))
     }
 
-    /// Push one decoded text chunk through the Qwen Coder parser.
-    fn push(&mut self, chunk: &str) -> Result<ToolParseResult> {
+    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
         self.buffer.push_str(chunk);
-        let mut result = ToolParseResult::default();
 
         while let Some((event, consumed_len)) = parse_buffered_event(&self.buffer, |input| {
             parse_next_qwen_coder_event(input, self.mode)
         })? {
-            self.apply_event(event, &mut result)?;
+            self.apply_event(event, output)?;
             self.buffer.drain(..consumed_len);
         }
 
-        Ok(result)
+        Ok(())
     }
 
-    /// Flush buffered text and reset parser state.
-    fn finish(&mut self) -> Result<ToolParseResult> {
-        let mut result = ToolParseResult::default();
+    fn finish(&mut self) -> Result<ToolParserOutput> {
+        let mut output = ToolParserOutput::default();
         if !self.buffer.is_empty() {
             if self.mode == QwenCoderMode::ToolCall || self.buffer.starts_with(TOOL_CALL_START) {
                 return Err(parsing_failed!("incomplete Qwen Coder tool call"));
             }
-            result.normal_text.push_str(&self.buffer);
+            output.normal_text.push_str(&self.buffer);
         }
-        self.reset();
-        Ok(result)
+        let _ = self.reset();
+        Ok(output)
+    }
+
+    fn reset(&mut self) -> String {
+        Qwen3CoderToolParser::reset(self)
     }
 }
 
@@ -229,6 +228,7 @@ mod tests {
     use thiserror_ext::AsReport;
 
     use super::{Qwen3CoderToolParser, ToolParser};
+    use crate::ToolParserTestExt as _;
     use crate::test_utils::{collect_stream, split_by_chars, test_tools};
 
     fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
@@ -243,27 +243,27 @@ mod tests {
     #[test]
     fn qwen_coder_parse_complete_without_tool_call_keeps_text() {
         let mut parser = Qwen3CoderToolParser::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 qwen_coder_parse_complete_extracts_single_tool_call() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "get_weather",
                 &[("location", "SF"), ("date", "2026-04-29")],
             ))
             .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": "2026-04-29"
@@ -278,16 +278,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 qwen_coder_parse_complete_converts_schema_types() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "convert",
                 &[
@@ -300,9 +300,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,
@@ -316,12 +316,12 @@ mod tests {
     #[test]
     fn qwen_coder_parse_complete_extracts_empty_arguments() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser.parse_complete(&build_tool_call("get_weather", &[])).unwrap();
+        let output = parser.parse_complete(&build_tool_call("get_weather", &[])).unwrap();
 
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].name.as_deref(), Some("get_weather"));
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({})
         );
     }
@@ -329,7 +329,7 @@ mod tests {
     #[test]
     fn qwen_coder_parse_complete_handles_upstream_multiline_typed_params() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(
                 "<tool_call>\n\
                  <function=calculate_area>\n\
@@ -348,10 +348,10 @@ mod tests {
             )
             .unwrap();
 
-        assert_eq!(result.calls.len(), 1);
-        assert_eq!(result.calls[0].name.as_deref(), Some("calculate_area"));
+        assert_eq!(output.calls.len(), 1);
+        assert_eq!(output.calls[0].name.as_deref(), Some("calculate_area"));
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({
                 "shape": "rectangle",
                 "dimensions": { "width": 10, "height": 20 },
@@ -363,7 +363,7 @@ mod tests {
     #[test]
     fn qwen_coder_parse_complete_handles_nested_json_parameter() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "convert",
                 &[(
@@ -373,9 +373,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!({
                 "payload": {
                     "nested": {
@@ -390,7 +390,7 @@ mod tests {
     #[test]
     fn qwen_coder_parse_complete_preserves_xml_like_parameter_values() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "process",
                 &[
@@ -403,9 +403,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!({
                 "html_content": r#"<div class="test"><span>Hello</span></div>"#,
                 "xml_snippet": r#"<root><child attr="value"/></root>"#,
@@ -416,7 +416,7 @@ mod tests {
     #[test]
     fn qwen_coder_parse_complete_unescapes_literal_closing_tags_in_parameter_value() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "get_weather",
                 &[
@@ -429,9 +429,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!({
                 "location": "杭州 </parameter></function></tool_call>",
                 "date": "2026-05-08",
@@ -442,16 +442,16 @@ mod tests {
     #[test]
     fn qwen_coder_parse_complete_does_not_double_encode_anyof_object() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(&build_tool_call(
                 "update_record",
                 &[("data", r#"{"key":"value","count":42}"#)],
             ))
             .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!({
                 "data": { "key": "value", "count": 42 },
             })
@@ -461,7 +461,7 @@ mod tests {
     #[test]
     fn qwen_coder_streaming_extracts_single_tool_call() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "<tool_call>\n",
@@ -472,11 +472,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" })
         );
     }
@@ -484,7 +484,7 @@ mod tests {
     #[test]
     fn qwen_coder_streaming_preserves_prefix_text() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = collect_stream(
+        let output = collect_stream(
             &mut parser,
             &[
                 "Thinking... ",
@@ -496,17 +496,17 @@ 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 qwen_coder_streaming_without_tool_call_emits_text_incrementally() {
         let mut parser = Qwen3CoderToolParser::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]
@@ -517,19 +517,19 @@ mod tests {
             build_tool_call("get_weather", &[("location", "NYC")])
         );
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = collect_stream(&mut parser, &[&text]);
+        let output = collect_stream(&mut parser, &[&text]);
 
-        assert_eq!(result.calls.len(), 2);
-        assert_eq!(result.calls[0].name.as_deref(), Some("get_weather"));
-        assert_eq!(result.calls[1].name.as_deref(), Some("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" })
         );
     }
@@ -543,19 +543,19 @@ mod tests {
         );
         let chunks = split_by_chars(&text, 5);
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = collect_stream(&mut parser, &chunks);
+        let output = collect_stream(&mut parser, &chunks);
 
         assert_eq!(
-            result.normal_text,
+            output.normal_text,
             "I'll check two cities.Between calls.Done."
         );
-        assert_eq!(result.calls.len(), 2);
+        assert_eq!(output.calls.len(), 2);
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[0].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[0].arguments).unwrap(),
             json!({ "city": "Dallas", "state": "TX" })
         );
         assert_eq!(
-            serde_json::from_str::<Value>(&result.calls[1].arguments).unwrap(),
+            serde_json::from_str::<Value>(&output.calls[1].arguments).unwrap(),
             json!({ "city": "Orlando", "state": "FL" })
         );
     }
@@ -565,11 +565,11 @@ mod tests {
         let text = build_tool_call("get_weather", &[("location", "SF")]);
         let chunks = split_by_chars(&text, 3);
         let mut parser = Qwen3CoderToolParser::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" })
         );
     }
@@ -577,19 +577,19 @@ mod tests {
     #[test]
     fn qwen_coder_streaming_does_not_emit_incomplete_tool_call() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
-            .push("<tool_call>\n<function=get_weather>\n<parameter=location>SF</parameter>")
+        let output = parser
+            .parse_chunk("<tool_call>\n<function=get_weather>\n<parameter=location>SF</parameter>")
             .unwrap();
 
-        assert!(result.normal_text.is_empty());
-        assert!(result.calls.is_empty());
+        assert!(output.normal_text.is_empty());
+        assert!(output.calls.is_empty());
     }
 
     #[test]
     fn qwen_coder_finish_fails_incomplete_tool_call() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
         parser
-            .push("<tool_call>\n<function=get_weather>\n<parameter=location>SF</parameter>")
+            .parse_chunk("<tool_call>\n<function=get_weather>\n<parameter=location>SF</parameter>")
             .unwrap();
 
         assert!(parser.finish().is_err());
@@ -598,7 +598,7 @@ mod tests {
     #[test]
     fn qwen_coder_malformed_tool_call_fails_fast() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let error = parser.push("<tool_call>\n<bad>\n</tool_call>").unwrap_err();
+        let error = parser.parse_chunk("<tool_call>\n<bad>\n</tool_call>").unwrap_err();
 
         expect!["tool parser parsing failed: "].assert_eq(&error.to_report_string());
     }
@@ -607,7 +607,7 @@ mod tests {
     fn qwen_coder_missing_parameter_end_fails_fast_after_function_end() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
         let error = parser
-            .push(
+            .parse_chunk(
                 "<tool_call>\n<function=get_weather>\n<parameter=location>SF</function>\n</tool_call>",
             )
             .unwrap_err();
@@ -618,14 +618,14 @@ mod tests {
     #[test]
     fn qwen_coder_parse_function_body_trims_one_wrapping_newline() {
         let mut parser = Qwen3CoderToolParser::new(&test_tools());
-        let result = parser
+        let output = parser
             .parse_complete(
                 "<tool_call>\n<function=get_weather>\n<parameter=location>\nHangzhou\n</parameter>\n</function>\n</tool_call>",
             )
             .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" })
         );
     }

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

@@ -1,7 +1,7 @@
 use serde_json::json;
 
-use super::{ToolParseResult, ToolParser};
-use crate::Tool;
+use super::{ToolParser, ToolParserOutput};
+use crate::{Tool, ToolParserTestExt as _};
 
 /// Build a reusable set of function tools for parser unit tests.
 pub fn test_tools() -> Vec<Tool> {
@@ -82,13 +82,15 @@ pub fn test_tools() -> Vec<Tool> {
 }
 
 /// Push chunks through a streaming parser and coalesce its tool-call deltas.
-pub fn collect_stream<T: ToolParser + ?Sized>(parser: &mut T, chunks: &[&str]) -> ToolParseResult {
-    let mut result = ToolParseResult::default();
+///
+/// Panics if there are any parsing errors along the way.
+pub fn collect_stream<T: ToolParser + ?Sized>(parser: &mut T, chunks: &[&str]) -> ToolParserOutput {
+    let mut output = ToolParserOutput::default();
     for chunk in chunks {
-        result.append(parser.push(chunk).unwrap());
+        output.append(parser.parse_chunk(chunk).unwrap());
     }
-    result.append(parser.finish().unwrap());
-    result.coalesce_calls()
+    output.append(parser.finish().unwrap());
+    output.coalesce_calls()
 }
 
 /// Split text into chunks containing at most `chunk_chars` Unicode scalar

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

@@ -1,4 +1,5 @@
-use super::{Result, Tool, ToolCallDelta, ToolParseResult, ToolParser};
+use super::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
+use crate::ToolParserTestExt as _;
 
 struct DefaultParser;
 
@@ -10,8 +11,16 @@ impl ToolParser for DefaultParser {
         Ok(Box::new(Self))
     }
 
-    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()
     }
 }
 
@@ -23,7 +32,7 @@ fn tool_parser_does_not_preserve_special_tokens_by_default() {
 }
 
 #[test]
-fn default_parse_complete_delegates_through_push_and_finish() {
+fn default_parse_complete_delegates_through_parse_chunk_and_finish() {
     struct StreamingParser;
 
     impl ToolParser for StreamingParser {
@@ -34,31 +43,30 @@ fn default_parse_complete_delegates_through_push_and_finish() {
             Ok(Box::new(Self))
         }
 
-        fn push(&mut self, _chunk: &str) -> Result<ToolParseResult> {
-            Ok(ToolParseResult {
-                normal_text: "prefix ".to_string(),
-                calls: vec![
-                    ToolCallDelta {
-                        tool_index: 0,
-                        name: Some("weather".to_string()),
-                        arguments: "{\"location\":".to_string(),
-                    },
-                    ToolCallDelta {
-                        tool_index: 0,
-                        name: None,
-                        arguments: "\"Paris\"".to_string(),
-                    },
-                    ToolCallDelta {
-                        tool_index: 1,
-                        name: Some("time".to_string()),
-                        arguments: "{\"timezone\":".to_string(),
-                    },
-                ],
-            })
+        fn parse_into(&mut self, _chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
+            output.normal_text.push_str("prefix ");
+            output.calls.extend([
+                ToolCallDelta {
+                    tool_index: 0,
+                    name: Some("weather".to_string()),
+                    arguments: "{\"location\":".to_string(),
+                },
+                ToolCallDelta {
+                    tool_index: 0,
+                    name: None,
+                    arguments: "\"Paris\"".to_string(),
+                },
+                ToolCallDelta {
+                    tool_index: 1,
+                    name: Some("time".to_string()),
+                    arguments: "{\"timezone\":".to_string(),
+                },
+            ]);
+            Ok(())
         }
 
-        fn finish(&mut self) -> Result<ToolParseResult> {
-            Ok(ToolParseResult {
+        fn finish(&mut self) -> Result<ToolParserOutput> {
+            Ok(ToolParserOutput {
                 normal_text: "suffix".to_string(),
                 calls: vec![
                     ToolCallDelta {
@@ -74,13 +82,17 @@ fn default_parse_complete_delegates_through_push_and_finish() {
                 ],
             })
         }
+
+        fn reset(&mut self) -> String {
+            String::new()
+        }
     }
 
     let mut parser = StreamingParser;
-    let result = parser.parse_complete("ignored").unwrap();
-    assert_eq!(result.normal_text, "prefix suffix");
+    let output = parser.parse_complete("ignored").unwrap();
+    assert_eq!(output.normal_text, "prefix suffix");
     assert_eq!(
-        result.calls,
+        output.calls,
         vec![
             ToolCallDelta {
                 tool_index: 0,

setup.py

@@ -1017,6 +1017,8 @@ def get_requirements() -> list[str]:
             if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
                 # [cu13] extra is the default; strip it on CUDA 12 builds.
                 req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl")
+            if "humming-kernels[cu13]" in req and cuda_major == "12":
+                req = req.replace("humming-kernels[cu13]", "humming-kernels[cu12]")
             modified_requirements.append(req)
         requirements = modified_requirements
     elif _is_hip():

tests/benchmarks/test_custom_image_dataset.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from argparse import Namespace
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from vllm.benchmarks.datasets import CustomImageDataset, get_samples
+from vllm.benchmarks.lib.endpoint_request_func import (
+    RequestFuncInput,
+    _get_chat_content,
+    _get_chat_messages,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+class _TokenizedPrompt:
+    def __init__(self, prompt: str) -> None:
+        self.input_ids = prompt.split()
+
+
+class _Tokenizer:
+    def __call__(self, prompt: str) -> _TokenizedPrompt:
+        return _TokenizedPrompt(prompt)
+
+
+def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
+    with path.open("w") as f:
+        for row in rows:
+            f.write(json.dumps(row) + "\n")
+
+
+def _args_for_custom_image(dataset_path: Path) -> Namespace:
+    return Namespace(
+        dataset_name="custom_image",
+        dataset_path=str(dataset_path),
+        disable_shuffle=True,
+        seed=0,
+        num_prompts=2,
+        custom_output_len=32,
+        enable_multimodal_chat=False,
+        request_id_prefix="req-",
+        no_oversample=False,
+    )
+
+
+@pytest.mark.benchmark
+def test_get_samples_custom_image_cli_path_supports_multi_image_and_content(
+    tmp_path: Path,
+) -> None:
+    image_a = tmp_path / "chart_a.png"
+    image_b = tmp_path / "chart_b.png"
+    image_c = tmp_path / "chart_c.png"
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(
+        jsonl,
+        [
+            {
+                "prompt": "Compare the first two charts.",
+                "image_files": [str(image_a), str(image_b)],
+            },
+            {
+                "content": [
+                    {"type": "text", "text": "Now compare "},
+                    {"type": "image", "image": str(image_c)},
+                ],
+            },
+        ],
+    )
+
+    samples = get_samples(_args_for_custom_image(jsonl), _Tokenizer())
+
+    assert len(samples) == 2
+    assert samples[0].request_id == "req-0"
+    assert isinstance(samples[0].multi_modal_data, list)
+    assert [part["image_url"]["url"] for part in samples[0].multi_modal_data] == [
+        f"file://{image_a}",
+        f"file://{image_b}",
+    ]
+
+    assert samples[1].request_id == "req-1"
+    assert samples[1].multi_modal_data is None
+    assert isinstance(samples[1].prompt, list)
+    assert samples[1].prompt[0] == {"type": "text", "text": "Now compare "}
+    assert samples[1].prompt[1]["image_url"]["url"] == f"file://{image_c}"
+
+
+@pytest.mark.benchmark
+def test_custom_image_dataset_uses_all_image_files(tmp_path: Path) -> None:
+    image_a = tmp_path / "chart_a.png"
+    image_b = tmp_path / "chart_b.png"
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(
+        jsonl,
+        [
+            {
+                "prompt": "Compare the charts.",
+                "image_files": [str(image_a), str(image_b)],
+            }
+        ],
+    )
+
+    dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    samples = dataset.sample(
+        tokenizer=_Tokenizer(),
+        num_requests=1,
+        output_len=32,
+    )
+
+    assert len(samples) == 1
+    sample = samples[0]
+    assert sample.prompt == "Compare the charts."
+    assert sample.prompt_len == 3
+    assert isinstance(sample.multi_modal_data, list)
+    assert [part["image_url"]["url"] for part in sample.multi_modal_data] == [
+        f"file://{image_a}",
+        f"file://{image_b}",
+    ]
+
+
+@pytest.mark.benchmark
+def test_custom_image_dataset_preserves_interleaved_content_order(
+    tmp_path: Path,
+) -> None:
+    image_a = tmp_path / "chart_a.png"
+    image_b = tmp_path / "chart_b.png"
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(
+        jsonl,
+        [
+            {
+                "content": [
+                    {"type": "text", "text": "Compare "},
+                    {"type": "image", "image": str(image_a)},
+                    {"type": "text", "text": " with "},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": str(image_b),
+                            "detail": "low",
+                        },
+                    },
+                ],
+            }
+        ],
+    )
+
+    dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    samples = dataset.sample(
+        tokenizer=_Tokenizer(),
+        num_requests=1,
+        output_len=32,
+    )
+
+    assert len(samples) == 1
+    sample = samples[0]
+    assert sample.multi_modal_data is None
+    assert sample.prompt_len == 2
+    assert isinstance(sample.prompt, list)
+    assert [part["type"] for part in sample.prompt] == [
+        "text",
+        "image_url",
+        "text",
+        "image_url",
+    ]
+    assert sample.prompt[1]["image_url"]["url"] == f"file://{image_a}"
+    assert sample.prompt[3]["image_url"] == {
+        "url": f"file://{image_b}",
+        "detail": "low",
+    }
+
+    request_input = RequestFuncInput(
+        prompt=sample.prompt,
+        api_url="http://localhost:8000/v1/chat/completions",
+        prompt_len=sample.prompt_len,
+        output_len=32,
+        model="test-model",
+    )
+    assert _get_chat_content(request_input) == sample.prompt
+
+
+@pytest.mark.benchmark
+def test_custom_image_dataset_wraps_interleaved_content_for_multimodal_chat(
+    tmp_path: Path,
+) -> None:
+    image = tmp_path / "chart.png"
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(
+        jsonl,
+        [
+            {
+                "content": [
+                    {"type": "text", "text": "Describe "},
+                    {"type": "image", "image": str(image)},
+                ],
+            }
+        ],
+    )
+
+    dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    samples = dataset.sample(
+        tokenizer=_Tokenizer(),
+        num_requests=1,
+        output_len=32,
+        enable_multimodal_chat=True,
+    )
+
+    sample = samples[0]
+    assert sample.multi_modal_data is None
+    assert sample.prompt == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe "},
+                {"type": "image_url", "image_url": {"url": f"file://{image}"}},
+            ],
+        }
+    ]
+
+    request_input = RequestFuncInput(
+        prompt=sample.prompt,
+        api_url="http://localhost:8000/v1/chat/completions",
+        prompt_len=sample.prompt_len,
+        output_len=32,
+        model="test-model",
+    )
+    assert _get_chat_messages(request_input) == sample.prompt
+
+
+@pytest.mark.benchmark
+def test_custom_image_dataset_rejects_invalid_content_part(
+    tmp_path: Path,
+) -> None:
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(jsonl, [{"content": [{"type": "audio", "audio": "clip.wav"}]}])
+
+    dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    with pytest.raises(ValueError, match="type 'text', 'image', or 'image_url'"):
+        dataset.sample(
+            tokenizer=_Tokenizer(),
+            num_requests=1,
+            output_len=32,
+        )

tests/compile/test_aot_compile.py

@@ -465,7 +465,10 @@ def test_standalone_compile_correctness():
         common_args,
         common_args,
         env1={"VLLM_USE_STANDALONE_COMPILE": "1"},
-        env2={"VLLM_USE_STANDALONE_COMPILE": "0"},
+        env2={
+            "VLLM_USE_STANDALONE_COMPILE": "0",
+            "VLLM_USE_MEGA_AOT_ARTIFACT": "0",
+        },
     )
 
 

tests/distributed/test_quick_all_reduce.py

@@ -25,16 +25,34 @@ from ..utils import (
     ensure_model_parallel_initialized,
     init_test_distributed_environment,
     multi_process_parallel,
+    set_random_seed,
 )
 
-torch.manual_seed(42)
-random.seed(44)
+
+def on_gfx942() -> bool:
+    if current_platform.is_rocm():
+        from vllm.platforms.rocm import on_gfx942 as rocm_on_gfx942
+
+        return rocm_on_gfx942()
+    return False
+
+
+set_random_seed(42)
+_test_size_rng = random.Random(44)
 # Size over 8MB is sufficient for custom quick allreduce.
-test_sizes = [random.randint(8 * 1024 * 1024, 10 * 1024 * 1024) for _ in range(8)]
+test_sizes = [
+    _test_size_rng.randint(8 * 1024 * 1024, 10 * 1024 * 1024) for _ in range(8)
+]
 for i, v in enumerate(test_sizes):
     test_sizes[i] -= v % 8
 
 
+def _assert_quickreduce(fa, inp):
+    assert fa is not None
+    assert not fa.disabled
+    assert fa.should_quick_allreduce(inp)
+
+
 @pytest.fixture
 def envs_cache_disabled():
     disable_envs_cache()
@@ -216,11 +234,14 @@ def graph_quickreduce(
 ):
     with monkeypatch.context() as m:
         m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+        m.delenv("HIP_VISIBLE_DEVICES", raising=False)
+        m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
         device = torch.device(f"cuda:{rank}")
         torch.accelerator.set_device_index(device)
         init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
         ensure_model_parallel_initialized(tp_size, pp_size)
         group = get_tp_group().device_group
+        fa = get_tp_group().device_communicator.qr_comm
 
         # A small all_reduce for warmup.
         # this is needed because device communicators might be created lazily
@@ -246,6 +267,8 @@ def graph_quickreduce(
                     device_idx = torch.accelerator.current_device_index()
                     inp1 = torch.randint(1, 23, (sz,), dtype=dtype, device=device_idx)
                     inp2 = torch.randint(-23, 1, (sz,), dtype=dtype, device=device_idx)
+                    _assert_quickreduce(fa, inp1)
+                    _assert_quickreduce(fa, inp2)
 
                     torch.accelerator.synchronize()
                     graph = torch.cuda.CUDAGraph()
@@ -270,6 +293,8 @@ def eager_quickreduce(
 ):
     with monkeypatch.context() as m:
         m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+        m.delenv("HIP_VISIBLE_DEVICES", raising=False)
+        m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
         device = torch.device(f"cuda:{rank}")
         torch.accelerator.set_device_index(device)
 
@@ -281,12 +306,42 @@ def eager_quickreduce(
         inp = torch.tensor(
             [1.0 * ((i) % 23) for i in range(sz)], dtype=torch.float16, device=device
         )
+        _assert_quickreduce(fa, inp)
         out = fa.quick_all_reduce(inp)
         torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
 
         inp = torch.tensor(
             [1.0 * ((i) % 23) for i in range(sz)], dtype=torch.bfloat16, device=device
         )
+        _assert_quickreduce(fa, inp)
+        out = fa.quick_all_reduce(inp)
+        torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
+
+
+@ray.remote(num_gpus=1, max_calls=1)
+def bf16_cast_quickreduce(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size,
+    pp_size,
+    rank,
+    distributed_init_port,
+):
+    with monkeypatch.context() as m:
+        m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+        m.delenv("HIP_VISIBLE_DEVICES", raising=False)
+        m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
+        m.setenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1")
+        device = torch.device(f"cuda:{rank}")
+        torch.accelerator.set_device_index(device)
+        init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+
+        sz = 16 * 1024 * 1024
+        fa = get_tp_group().device_communicator.qr_comm
+        inp = torch.tensor(
+            [1.0 * (i % 23) for i in range(sz)], dtype=torch.bfloat16, device=device
+        )
+        _assert_quickreduce(fa, inp)
+        assert fa.use_fp16_kernels
         out = fa.quick_all_reduce(inp)
         torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
 
@@ -308,12 +363,27 @@ def test_custom_quick_allreduce(
     world_size = tp_size * pipeline_parallel_size
     if world_size > torch.accelerator.device_count():
         pytest.skip("Not enough GPUs to run the test.")
+    if test_target is graph_quickreduce and on_gfx942():
+        pytest.xfail(
+            "CUDA graph capture with quick reduce hits "
+            "hipErrorStreamCaptureInvalidated on gfx942"
+        )
 
     monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
 
     multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="only test quick allreduce for rocm"
+)
+def test_custom_quick_allreduce_bf16_cast(monkeypatch: pytest.MonkeyPatch):
+    if torch.accelerator.device_count() < 2:
+        pytest.skip("Not enough GPUs to run the test.")
+    monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "FP")
+    multi_process_parallel(monkeypatch, 2, 1, bf16_cast_quickreduce)
+
+
 def qr_variable_input(rank, world_size):
     """
     When the tensor parallelism is set to 4 or 8, frequent changes

tests/distributed/test_rocm_quick_reduce.py

@@ -0,0 +1,750 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+import importlib
+import multiprocessing as mp
+import os
+import queue
+import traceback
+from functools import lru_cache
+from types import SimpleNamespace
+from typing import Literal
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.distributed as dist
+from huggingface_hub import snapshot_download
+
+import vllm.envs as envs
+from vllm import LLM, SamplingParams
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.platforms import current_platform
+from vllm.utils.network_utils import get_open_port
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_rocm(),
+    reason="ROCm-only quick-reduce tests",
+)
+
+MB = 1024 * 1024
+WORLD_SIZE = 2
+QUANT_LEVELS = ["FP", "INT8", "INT6", "INT4"]
+
+
+def _log(message: str) -> None:
+    print(f"[rocm_quick_reduce] {message}", flush=True)
+
+
+def _reload_envs():
+    return importlib.reload(envs)
+
+
+def _make_quick_allreduce(
+    *,
+    disabled: bool = False,
+    world_size: int = 2,
+    quant_level: str = "FP",
+    use_fp16_kernels: bool = False,
+    qr_max_size: int = 64 * MB,
+):
+    from vllm.distributed.device_communicators.quick_all_reduce import (
+        QuickAllReduce,
+        QuickReduceRegime,
+    )
+
+    qar = QuickAllReduce.__new__(QuickAllReduce)
+    qar.disabled = disabled
+    qar.world_size = world_size
+    qar.use_fp16_kernels = use_fp16_kernels
+    qar.qr_quant_level = QuickReduceRegime[quant_level]
+    qar.qr_max_size = qr_max_size
+    return qar
+
+
+def _quick_allreduce_worker(
+    rank: int,
+    port: int,
+    quant_level: str,
+    dtype_name: str,
+    cast_bf16: bool,
+):
+    os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level
+    os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "1" if cast_bf16 else "0"
+    _log(
+        f"worker start: rank={rank} quant={quant_level} "
+        f"dtype={dtype_name} cast_bf16={cast_bf16}"
+    )
+
+    device = torch.device(f"cuda:{rank}")
+    torch.accelerator.set_device_index(device)
+    dist.init_process_group(
+        backend="gloo",
+        init_method=f"tcp://127.0.0.1:{port}",
+        rank=rank,
+        world_size=WORLD_SIZE,
+    )
+
+    qar = None
+    try:
+        from vllm.distributed.device_communicators.quick_all_reduce import (
+            QuickAllReduce,
+        )
+
+        qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank)
+        assert not qar.disabled
+
+        num_elements = 8 * MB if dtype_name == "float16" else 4 * MB
+
+        dtype = getattr(torch, dtype_name)
+        inp = torch.ones(num_elements, dtype=dtype, device=device)
+
+        assert qar.should_quick_allreduce(inp)
+        if cast_bf16:
+            assert qar.use_fp16_kernels
+
+        out = qar.quick_all_reduce(inp)
+        assert torch.allclose(out, inp * WORLD_SIZE, atol=2.5, rtol=0.1)
+        _log(
+            f"worker complete: rank={rank} quant={quant_level} "
+            f"dtype={dtype_name} num_elements={num_elements} "
+            f"use_fp16_kernels={qar.use_fp16_kernels}"
+        )
+    finally:
+        if qar is not None:
+            qar.close()
+        if dist.is_initialized():
+            dist.destroy_process_group()
+
+
+def _run_two_gpu_quick_allreduce_test(
+    *,
+    quant_level: str,
+    dtype_name: str,
+    cast_bf16: bool,
+):
+    _log(
+        f"launch 2-GPU case: quant={quant_level} "
+        f"dtype={dtype_name} cast_bf16={cast_bf16}"
+    )
+    ctx = mp.get_context("spawn")
+    port = get_open_port()
+    procs = []
+
+    for rank in range(WORLD_SIZE):
+        proc = ctx.Process(
+            target=_quick_allreduce_worker,
+            args=(rank, port, quant_level, dtype_name, cast_bf16),
+        )
+        proc.start()
+        procs.append(proc)
+
+    for proc in procs:
+        proc.join(timeout=60)
+        assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}"
+    _log(
+        f"finished 2-GPU case: quant={quant_level} "
+        f"dtype={dtype_name} cast_bf16={cast_bf16}"
+    )
+
+
+MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
+E2E_PREFILL_TOKENS = 1024
+E2E_MAX_MODEL_LEN = 1536
+E2E_GPU_MEMORY_UTILIZATION = 0.3
+E2E_KV_CACHE_MEMORY_BYTES = 2 << 30
+
+_BACKGROUND_LINE = (
+    "Background filler: this archived operations memo repeats a routine status "
+    "line so the distributed test uses a realistically long prefill."
+)
+_BACKGROUND_BLOCK = " ".join([_BACKGROUND_LINE] * 48)
+
+
+def _build_prompt(*, fact_block: str, question: str) -> str:
+    return (
+        "Read the archived operations memo below. Most of the memo is filler. "
+        "Use only the fact block near the end when answering.\n"
+        f"{_BACKGROUND_BLOCK}\n"
+        "Fact block:\n"
+        f"{fact_block}\n"
+        f"Question: {question}\n"
+        "Answer in one short sentence."
+    )
+
+
+E2E_PROMPTS = [
+    _build_prompt(
+        fact_block=(
+            "- Festival city: Oslo\n- Mascot animal: otter\n- Welcome drink: tea"
+        ),
+        question="Which city hosts the festival, and what animal is the mascot?",
+    ),
+    _build_prompt(
+        fact_block=(
+            "- Meeting day: Tuesday\n"
+            "- Planned snack: apricot cake\n"
+            "- Backup room: Cedar"
+        ),
+        question="What day is the meeting, and what snack is planned?",
+    ),
+]
+RECORDED_RESPONSE_TEXTS = (
+    " The city hosting the festival is Oslo, and the mascot is an otter.",
+    " The meeting is on Tuesday and the snack planned is apricot cake.",
+)
+REQUIRED_WORDS = (("oslo", "otter"), ("tuesday", "apricot"))
+
+
+def _log_prompt_summaries() -> None:
+    for i, prompt in enumerate(E2E_PROMPTS):
+        prompt_lines = prompt.splitlines()
+        fact_block = [line for line in prompt_lines if line.startswith("- ")]
+        fact_summary = "; ".join(line.removeprefix("- ") for line in fact_block)
+        _log(f"prompt {i} facts: {fact_summary}")
+
+
+@lru_cache(maxsize=1)
+def _get_model_path() -> str:
+    try:
+        path = snapshot_download(repo_id=MODEL_NAME, local_files_only=True)
+        _log(f"using cached model snapshot: {path}")
+        return path
+    except Exception:
+        path = snapshot_download(repo_id=MODEL_NAME)
+        _log(f"downloaded model snapshot: {path}")
+        return path
+
+
+def _get_hidden_size(model_config) -> int:
+    hidden_size = getattr(model_config, "hidden_size", None)
+    if hidden_size is None and hasattr(model_config, "text_config"):
+        hidden_size = getattr(model_config.text_config, "hidden_size", None)
+    assert isinstance(hidden_size, int)
+    return hidden_size
+
+
+def _check_tp_allreduce_uses_quick_reduce(
+    self,
+    num_tokens: int,
+    dtype_name: str = "float16",
+) -> dict[str, int | bool]:
+    from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+    from vllm.distributed.parallel_state import get_tp_group
+
+    assert self.device is not None
+    qr_comm = get_tp_group().device_communicator.qr_comm
+    assert qr_comm is not None
+    assert not qr_comm.disabled
+
+    hidden_size = _get_hidden_size(self.model_runner.model.config)
+    dtype = getattr(torch, dtype_name)
+    sample = torch.full(
+        (num_tokens, hidden_size),
+        fill_value=float(self.rank + 1),
+        dtype=dtype,
+        device=self.device,
+    )
+
+    assert qr_comm.should_quick_allreduce(sample)
+
+    expected = sample.clone()
+    reduced = tensor_model_parallel_all_reduce(sample)
+    dist.all_reduce(expected, group=get_tp_group().device_group)
+    torch.testing.assert_close(reduced, expected, atol=2.5, rtol=0.1)
+
+    stats = {
+        "rank": self.rank,
+        "hidden_size": hidden_size,
+        "num_tokens": num_tokens,
+        "use_fp16_kernels": qr_comm.use_fp16_kernels,
+    }
+    _log(
+        "worker quick-reduce check: "
+        f"rank={self.rank} hidden_size={hidden_size} "
+        f"num_tokens={num_tokens} use_fp16_kernels={qr_comm.use_fp16_kernels}"
+    )
+    return stats
+
+
+def _check_quick_reduce_disabled(self) -> int:
+    from vllm.distributed.parallel_state import get_tp_group
+
+    qr_comm = get_tp_group().device_communicator.qr_comm
+    assert qr_comm is not None
+    assert qr_comm.disabled
+    _log(f"worker confirmed quick reduce is disabled: rank={self.rank}")
+    return self.rank
+
+
+def _collect_generations(outputs) -> list[tuple[tuple[int, ...], str]]:
+    return [
+        (tuple(output.outputs[0].token_ids), output.outputs[0].text)
+        for output in outputs
+    ]
+
+
+def _shutdown_llm(llm: LLM | None) -> None:
+    if llm is None:
+        cleanup_dist_env_and_memory()
+        return
+
+    with contextlib.suppress(Exception):
+        llm.llm_engine.engine_core.shutdown()
+
+    del llm
+    cleanup_dist_env_and_memory()
+
+
+def _log_generations(
+    label: str,
+    generations: list[tuple[tuple[int, ...], str]],
+) -> None:
+    for i, (token_ids, text) in enumerate(generations):
+        _log(f"{label} prompt {i} token ids: {list(token_ids)}")
+        _log(f"{label} prompt {i} text: {text!r}")
+
+
+def _assert_required_words(
+    label: str,
+    generations: list[tuple[tuple[int, ...], str]],
+) -> None:
+    for i, (_, text) in enumerate(generations):
+        lowered = text.lower()
+        missing = [word for word in REQUIRED_WORDS[i] if word not in lowered]
+        assert not missing, (
+            f"{label} prompt {i} is missing required words {missing}. "
+            f"Observed text: {text!r}"
+        )
+
+
+def _collect_soft_mismatches(
+    baseline_generations: list[tuple[tuple[int, ...], str]],
+    quick_reduce_generations: list[tuple[tuple[int, ...], str]],
+) -> list[str]:
+    mismatches = []
+
+    for i, (_, text) in enumerate(baseline_generations):
+        expected = RECORDED_RESPONSE_TEXTS[i]
+        if text != expected:
+            mismatches.append(
+                f"baseline prompt {i} drifted from the recorded response.\n"
+                f"expected={expected!r}\nactual={text!r}"
+            )
+
+    for i, (_, text) in enumerate(quick_reduce_generations):
+        expected = RECORDED_RESPONSE_TEXTS[i]
+        if text != expected:
+            mismatches.append(
+                f"quick-reduce prompt {i} drifted from the recorded response.\n"
+                f"expected={expected!r}\nactual={text!r}"
+            )
+
+    for i, ((_, baseline_text), (_, quick_reduce_text)) in enumerate(
+        zip(baseline_generations, quick_reduce_generations)
+    ):
+        if baseline_text != quick_reduce_text:
+            mismatches.append(
+                f"baseline and quick-reduce responses differ for prompt {i}.\n"
+                f"baseline={baseline_text!r}\nquick_reduce={quick_reduce_text!r}"
+            )
+
+    return mismatches
+
+
+def _run_generation(
+    *,
+    backend: Literal["mp", "ray"],
+    quant_mode: str,
+    expect_quick_reduce: bool,
+) -> list[tuple[tuple[int, ...], str]]:
+    llm = None
+    monkeypatch = pytest.MonkeyPatch()
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+        m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
+        model_path = _get_model_path()
+        _log(
+            f"starting generation: backend={backend} quant={quant_mode} "
+            f"gpu_memory_utilization={E2E_GPU_MEMORY_UTILIZATION} "
+            f"kv_cache_bytes={E2E_KV_CACHE_MEMORY_BYTES} model={model_path}"
+        )
+
+        try:
+            llm = LLM(
+                model=model_path,
+                tokenizer=model_path,
+                tensor_parallel_size=2,
+                distributed_executor_backend=backend,
+                dtype="half",
+                enforce_eager=True,
+                max_model_len=E2E_MAX_MODEL_LEN,
+                max_num_seqs=len(E2E_PROMPTS),
+                gpu_memory_utilization=E2E_GPU_MEMORY_UTILIZATION,
+                kv_cache_memory_bytes=E2E_KV_CACHE_MEMORY_BYTES,
+                seed=0,
+            )
+
+            if not expect_quick_reduce:
+                assert llm.collective_rpc(_check_quick_reduce_disabled) == [0, 1]
+
+            if expect_quick_reduce:
+                worker_stats = llm.collective_rpc(
+                    _check_tp_allreduce_uses_quick_reduce,
+                    args=(E2E_PREFILL_TOKENS,),
+                )
+                assert [stat["rank"] for stat in worker_stats] == [0, 1]
+                worker_summary = "; ".join(
+                    "rank={rank} hidden_size={hidden_size} num_tokens={num_tokens} "
+                    "use_fp16_kernels={use_fp16_kernels}".format(**stat)
+                    for stat in worker_stats
+                )
+                _log(f"{backend} quick-reduce worker checks: {worker_summary}")
+
+            outputs = llm.generate(
+                E2E_PROMPTS,
+                SamplingParams(
+                    temperature=0.0,
+                    max_tokens=20,
+                    stop=["\nAnswer:", " Answer:"],
+                ),
+                use_tqdm=False,
+            )
+            generations = _collect_generations(outputs)
+            assert all(text.strip() for _, text in generations)
+            _log_generations(f"{backend} {quant_mode}", generations)
+            return generations
+        finally:
+            _shutdown_llm(llm)
+
+
+def _run_quick_reduce_llm_e2e_in_subprocess(
+    *,
+    backend: Literal["mp", "ray"],
+) -> str | None:
+    _log(f"running LLM e2e: backend={backend}")
+    _log_prompt_summaries()
+    baseline_outputs = _run_generation(
+        backend=backend,
+        quant_mode="NONE",
+        expect_quick_reduce=False,
+    )
+    quick_reduce_outputs = _run_generation(
+        backend=backend,
+        quant_mode="FP",
+        expect_quick_reduce=True,
+    )
+
+    _assert_required_words("baseline", baseline_outputs)
+    _assert_required_words("quick-reduce", quick_reduce_outputs)
+
+    mismatches = _collect_soft_mismatches(baseline_outputs, quick_reduce_outputs)
+    if mismatches:
+        details = "\n\n".join(mismatches)
+        _log(f"soft response mismatch:\n{details}")
+        return details
+
+    _log(f"LLM e2e backend={backend} matched the recorded responses exactly")
+    return None
+
+
+def _quick_reduce_llm_e2e_worker(
+    result_queue: mp.Queue,
+    backend: Literal["mp", "ray"],
+) -> None:
+    try:
+        xfail_reason = _run_quick_reduce_llm_e2e_in_subprocess(backend=backend)
+    except Exception:
+        result_queue.put({"status": "error", "reason": traceback.format_exc()})
+        raise
+    else:
+        if xfail_reason is not None:
+            result_queue.put({"status": "xfail", "reason": xfail_reason})
+        else:
+            result_queue.put({"status": "ok"})
+
+
+def run_quick_reduce_llm_e2e(
+    *,
+    backend: Literal["mp", "ray"],
+) -> None:
+    ctx = mp.get_context("spawn")
+    result_queue = ctx.Queue()
+    proc = ctx.Process(
+        target=_quick_reduce_llm_e2e_worker,
+        args=(result_queue, backend),
+    )
+    proc.start()
+    proc.join(timeout=600)
+
+    try:
+        result = result_queue.get(timeout=5)
+    except queue.Empty as exc:
+        if proc.exitcode != 0:
+            raise AssertionError(
+                f"quick-reduce llm e2e subprocess failed for backend={backend} "
+                f"with exit code {proc.exitcode} and produced no result"
+            ) from exc
+        raise AssertionError(
+            f"quick-reduce llm e2e subprocess produced no result for backend={backend}"
+        ) from exc
+
+    if result["status"] == "xfail":
+        pytest.xfail(result["reason"])
+    if result["status"] == "error":
+        raise AssertionError(
+            f"quick-reduce llm e2e subprocess failed for backend={backend}:\n"
+            f"{result['reason']}"
+        )
+
+    assert proc.exitcode == 0, (
+        f"quick-reduce llm e2e subprocess failed for backend={backend} "
+        f"with exit code {proc.exitcode}"
+    )
+
+
+def test_quick_reduce_regime_values():
+    from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
+
+    assert QuickReduceRegime.FP.value == 0
+    assert QuickReduceRegime.INT8.value == 1
+    assert QuickReduceRegime.INT6.value == 2
+    assert QuickReduceRegime.INT4.value == 3
+    assert QuickReduceRegime.NONE.value == 4
+
+
+def test_quick_reduce_regime_names():
+    from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
+
+    assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
+
+
+@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
+def test_quick_reduce_quantization_env_var(monkeypatch, quant_level):
+    monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_level)
+
+    reloaded_envs = _reload_envs()
+    assert quant_level == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION
+
+
+def test_quick_reduce_quantization_default(monkeypatch):
+    monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", raising=False)
+
+    reloaded_envs = _reload_envs()
+    assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION == "NONE"
+
+
+@pytest.mark.parametrize("cast_bf16", [True, False])
+def test_quick_reduce_cast_bf16_to_fp16_env_var(monkeypatch, cast_bf16):
+    monkeypatch.setenv(
+        "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1" if cast_bf16 else "0"
+    )
+
+    reloaded_envs = _reload_envs()
+    assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is cast_bf16
+
+
+def test_quick_reduce_cast_bf16_to_fp16_default(monkeypatch):
+    monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", raising=False)
+
+    reloaded_envs = _reload_envs()
+    assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is True
+
+
+@pytest.mark.parametrize("max_mb", [128, 512, 2048, None])
+def test_quick_reduce_max_size_env_var(monkeypatch, max_mb):
+    if max_mb is None:
+        monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False)
+    else:
+        monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", str(max_mb))
+
+    reloaded_envs = _reload_envs()
+    assert max_mb == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB
+
+
+def test_quick_reduce_max_size_default(monkeypatch):
+    monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False)
+
+    reloaded_envs = _reload_envs()
+    assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is None
+
+
+@pytest.mark.parametrize(
+    ("gcn_arch_name", "expected"),
+    [
+        ("gfx942", True),
+        ("gfx950", True),
+        ("gfx90a", False),
+        ("", False),
+    ],
+)
+def test_quick_allreduce_rocm_arch_available(gcn_arch_name, expected):
+    from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
+
+    qar = QuickAllReduce.__new__(QuickAllReduce)
+    qar.disabled = True
+
+    with (
+        patch(
+            "vllm.distributed.device_communicators.quick_all_reduce.current_platform."
+            "is_rocm",
+            return_value=True,
+        ),
+        patch(
+            "torch.cuda.get_device_properties",
+            return_value=SimpleNamespace(gcnArchName=gcn_arch_name),
+        ),
+    ):
+        assert qar._rocm_arch_available() is expected
+
+
+def test_quick_allreduce_rocm_arch_available_handles_probe_failure():
+    from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
+
+    qar = QuickAllReduce.__new__(QuickAllReduce)
+    qar.disabled = True
+
+    with (
+        patch(
+            "vllm.distributed.device_communicators.quick_all_reduce.current_platform."
+            "is_rocm",
+            return_value=True,
+        ),
+        patch("torch.cuda.get_device_properties", side_effect=RuntimeError),
+    ):
+        assert qar._rocm_arch_available() is False
+
+
+def test_quick_allreduce_rejects_disabled():
+    qar = _make_quick_allreduce(disabled=True)
+
+    inp = torch.zeros(1024, dtype=torch.float16)
+    assert qar.should_quick_allreduce(inp) is False
+
+
+def test_quick_allreduce_rejects_unsupported_dtype():
+    qar = _make_quick_allreduce()
+
+    inp = torch.zeros(1024 * 1024, dtype=torch.float32)
+    assert qar.should_quick_allreduce(inp) is False
+
+
+def test_quick_allreduce_rejects_non_aligned_input():
+    qar = _make_quick_allreduce()
+
+    inp = torch.zeros(5, dtype=torch.float16)
+    assert qar.should_quick_allreduce(inp) is False
+
+
+def test_quick_allreduce_rejects_non_contiguous_input():
+    qar = _make_quick_allreduce()
+
+    inp = torch.zeros((1024, 1024), dtype=torch.float16)[:, ::2]
+    assert qar.should_quick_allreduce(inp) is False
+
+
+def test_quick_allreduce_rejects_input_smaller_than_threshold():
+    qar = _make_quick_allreduce()
+
+    inp = torch.zeros((MB // 2) - 8, dtype=torch.float16)
+    assert qar.should_quick_allreduce(inp) is False
+
+
+def test_quick_allreduce_accepts_input_at_threshold():
+    qar = _make_quick_allreduce()
+
+    inp = torch.zeros(MB // 2, dtype=torch.float16)
+    assert qar.should_quick_allreduce(inp) is True
+
+
+def test_quick_allreduce_rejects_input_larger_than_max_size():
+    qar = _make_quick_allreduce(qr_max_size=1 * MB)
+
+    inp = torch.zeros(MB, dtype=torch.float16)
+    assert qar.should_quick_allreduce(inp) is False
+
+
+def test_quick_allreduce_bf16_uses_fp16_threshold_when_cast_enabled():
+    inp = torch.zeros(MB // 2, dtype=torch.bfloat16)
+
+    without_cast = _make_quick_allreduce(use_fp16_kernels=False)
+    with_cast = _make_quick_allreduce(use_fp16_kernels=True)
+
+    assert without_cast.should_quick_allreduce(inp) is False
+    assert with_cast.should_quick_allreduce(inp) is True
+
+
+def test_quick_allreduce_supported_world_sizes():
+    from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
+
+    assert QuickAllReduce._SUPPORTED_WORLD_SIZES == [2, 4, 8]
+
+
+def test_quick_allreduce_supported_dtypes():
+    from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
+
+    assert [torch.float16, torch.bfloat16] == QuickAllReduce._SUPPORTED_DTYPES
+
+
+def test_quick_allreduce_min_size_table():
+    from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
+
+    for dtype in [torch.float16, torch.bfloat16]:
+        for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
+            min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
+            assert len(min_sizes) == 4
+            assert all(size > 0 for size in min_sizes)
+
+
+def test_qr_max_size():
+    from vllm import _custom_ops as ops
+
+    max_size = ops.qr_max_size()
+    assert isinstance(max_size, int)
+    assert max_size > 0
+
+
+@pytest.mark.skipif(
+    current_platform.device_count() < WORLD_SIZE,
+    reason="requires 2 ROCm GPUs",
+)
+@pytest.mark.parametrize("quant_level", QUANT_LEVELS)
+def test_quick_allreduce_two_gpu_correctness(quant_level):
+    _log(f"two-GPU correctness case: quant={quant_level}")
+    _run_two_gpu_quick_allreduce_test(
+        quant_level=quant_level,
+        dtype_name="float16",
+        cast_bf16=False,
+    )
+
+
+@pytest.mark.skipif(
+    current_platform.device_count() < WORLD_SIZE,
+    reason="requires 2 ROCm GPUs",
+)
+def test_quick_allreduce_bf16_cast_mode():
+    _log("BF16 cast case")
+    _run_two_gpu_quick_allreduce_test(
+        quant_level="FP",
+        dtype_name="bfloat16",
+        cast_bf16=True,
+    )
+
+
+@pytest.mark.skipif(
+    current_platform.device_count() < WORLD_SIZE,
+    reason="requires 2 ROCm GPUs",
+)
+def test_quick_allreduce_llm_e2e():
+    _log("LLM e2e case: backend=mp")
+    run_quick_reduce_llm_e2e(backend="mp")
+
+
+@pytest.mark.skipif(
+    current_platform.device_count() < WORLD_SIZE,
+    reason="requires 2 ROCm GPUs",
+)
+def test_quick_allreduce_llm_e2e_ray():
+    _log("LLM e2e case: backend=ray")
+    run_quick_reduce_llm_e2e(backend="ray")

tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from pydantic import ValidationError
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+
+
+@pytest.mark.parametrize("raw_value", [-2, 0.6, 10.5])
+def test_chat_completion_request_rejects_invalid_thinking_token_budget(raw_value):
+    with pytest.raises(ValidationError, match="thinking_token_budget"):
+        ChatCompletionRequest.model_validate(
+            {
+                "model": "qwen",
+                "messages": [{"role": "user", "content": "hello"}],
+                "thinking_token_budget": raw_value,
+            }
+        )
+
+
+def test_chat_completion_request_accepts_valid_thinking_token_budget():
+    request = ChatCompletionRequest.model_validate(
+        {
+            "model": "qwen",
+            "messages": [{"role": "user", "content": "hello"}],
+            "thinking_token_budget": 10,
+        }
+    )
+    assert request.thinking_token_budget == 10
+
+
+def test_chat_completion_request_accepts_minus_one_as_unlimited():
+    request = ChatCompletionRequest.model_validate(
+        {
+            "model": "qwen",
+            "messages": [{"role": "user", "content": "hello"}],
+            "thinking_token_budget": -1,
+        }
+    )
+    assert request.thinking_token_budget is None
+
+
+@pytest.mark.parametrize("raw_value", [0.6, 3.14, -2])
+def test_completion_request_rejects_invalid_thinking_token_budget(raw_value):
+    with pytest.raises(ValidationError, match="thinking_token_budget"):
+        CompletionRequest.model_validate(
+            {
+                "model": "qwen",
+                "prompt": "hello",
+                "thinking_token_budget": raw_value,
+            }
+        )
+
+
+def test_completion_request_accepts_valid_thinking_token_budget():
+    request = CompletionRequest.model_validate(
+        {
+            "model": "qwen",
+            "prompt": "hello",
+            "thinking_token_budget": 5,
+        }
+    )
+    assert request.thinking_token_budget == 5
+
+
+def test_completion_request_accepts_minus_one_as_unlimited():
+    request = CompletionRequest.model_validate(
+        {
+            "model": "qwen",
+            "prompt": "hello",
+            "thinking_token_budget": -1,
+        }
+    )
+    assert request.thinking_token_budget is None

tests/entrypoints/openai/test_reasoning_enable_thinking.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for reasoning_effort -> enable_thinking mapping.
+
+Models like Gemma4 require enable_thinking=True in chat_template_kwargs to
+activate thinking mode. This mapping ensures that when a user requests
+reasoning (via reasoning_effort or reasoning.effort), the template kwarg
+is injected automatically.
+"""
+
+import pytest
+from openai.types.shared import Reasoning
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+
+
+def _build_chat_request(**kwargs) -> ChatCompletionRequest:
+    defaults = dict(
+        model="test-model",
+        messages=[{"role": "user", "content": "Hello"}],
+    )
+    defaults.update(kwargs)
+    return ChatCompletionRequest(**defaults)
+
+
+def _build_responses_request(**kwargs) -> ResponsesRequest:
+    defaults = dict(
+        model="test-model",
+        input=[{"role": "user", "content": "Hello"}],
+    )
+    defaults.update(kwargs)
+    return ResponsesRequest(**defaults)
+
+
+class TestChatCompletionReasoningEffort:
+    """Chat Completions: reasoning_effort -> enable_thinking."""
+
+    @pytest.mark.parametrize("effort", ["low", "medium", "high"])
+    def test_non_none_effort_injects_enable_thinking_true(self, effort):
+        request = _build_chat_request(reasoning_effort=effort)
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is True
+
+    def test_none_effort_injects_enable_thinking_false(self):
+        request = _build_chat_request(reasoning_effort="none")
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is False
+
+    def test_no_effort_does_not_inject(self):
+        request = _build_chat_request()
+        params = request.build_chat_params(None, "auto")
+        assert "enable_thinking" not in params.chat_template_kwargs
+
+    def test_explicit_user_kwarg_not_overridden(self):
+        request = _build_chat_request(
+            reasoning_effort="high",
+            chat_template_kwargs={"enable_thinking": False},
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is False
+
+    def test_reasoning_effort_still_in_kwargs(self):
+        request = _build_chat_request(reasoning_effort="high")
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["reasoning_effort"] == "high"
+
+
+class TestResponsesReasoningEffort:
+    """Responses API: reasoning.effort -> enable_thinking."""
+
+    @pytest.mark.parametrize("effort", ["low", "medium", "high"])
+    def test_non_none_effort_injects_enable_thinking_true(self, effort):
+        request = _build_responses_request(
+            reasoning=Reasoning(effort=effort),
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is True
+
+    def test_none_effort_injects_enable_thinking_false(self):
+        request = _build_responses_request(
+            reasoning=Reasoning(effort="none"),
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is False
+
+    def test_no_reasoning_does_not_inject(self):
+        request = _build_responses_request()
+        params = request.build_chat_params(None, "auto")
+        assert "enable_thinking" not in params.chat_template_kwargs
+
+    def test_explicit_user_kwarg_not_overridden(self):
+        request = _build_responses_request(
+            reasoning=Reasoning(effort="high"),
+            chat_template_kwargs={"enable_thinking": False},
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["enable_thinking"] is False
+
+    def test_reasoning_effort_still_in_kwargs(self):
+        request = _build_responses_request(
+            reasoning=Reasoning(effort="high"),
+        )
+        params = request.build_chat_params(None, "auto")
+        assert params.chat_template_kwargs["reasoning_effort"] == "high"

tests/entrypoints/test_api_server_process_manager.py

@@ -8,8 +8,14 @@ import time
 from unittest.mock import patch
 
 import pytest
+import zmq
 
-from vllm.v1.utils import APIServerProcessManager, wait_for_completion_or_failure
+from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
+from vllm.v1.utils import (
+    APIServerProcessManager,
+    get_engine_client_zmq_addr,
+    wait_for_completion_or_failure,
+)
 
 # Global variables to control worker behavior
 WORKER_RUNTIME_SECONDS = 0.5
@@ -23,6 +29,39 @@ def mock_run_api_server_worker(listen_address, sock, args, client_config=None):
     print("Mock worker completed successfully")
 
 
+# Module-level stub for the gather_actual_addresses test. Must be
+# importable by `multiprocessing.spawn` (no closures, no nesting).
+def defer_addresses_stub_worker(listen_address, sock, args, client_config):
+    """Bind ROUTER/PULL with a kernel-assigned port, report the actual
+    endpoints back via the pipe, then exit."""
+    ctx = zmq.Context()
+    try:
+        in_sock = make_zmq_socket(
+            ctx, client_config["input_address"], zmq.ROUTER, bind=True
+        )
+        out_sock = make_zmq_socket(
+            ctx, client_config["output_address"], zmq.PULL, bind=True
+        )
+        try:
+            pipe = client_config["actual_address_pipe"]
+            try:
+                pipe.send(
+                    {
+                        "input_address": in_sock.getsockopt(zmq.LAST_ENDPOINT).decode(),
+                        "output_address": out_sock.getsockopt(
+                            zmq.LAST_ENDPOINT
+                        ).decode(),
+                    }
+                )
+            finally:
+                pipe.close()
+        finally:
+            in_sock.close(linger=0)
+            out_sock.close(linger=0)
+    finally:
+        ctx.term()
+
+
 @pytest.fixture
 def api_server_args():
     """Fixture to provide arguments for APIServerProcessManager."""
@@ -268,3 +307,92 @@ def test_external_process_monitoring(api_server_args):
         manager.shutdown()
         mock_coordinator.shutdown()
         time.sleep(0.2)
+
+
+@pytest.mark.timeout(60)
+def test_gather_actual_addresses_end_to_end():
+    """Each child binds ROUTER/PULL with a kernel-picked port and reports
+    the bound endpoints back via its per-child pipe; the manager surfaces
+    them via :py:meth:`gather_actual_addresses`."""
+    host = "127.0.0.1"
+    num_servers = 4
+
+    placeholder_inputs = [
+        get_engine_client_zmq_addr(local_only=False, host=host)
+        for _ in range(num_servers)
+    ]
+    placeholder_outputs = [
+        get_engine_client_zmq_addr(local_only=False, host=host)
+        for _ in range(num_servers)
+    ]
+    for addr in placeholder_inputs + placeholder_outputs:
+        assert addr == f"tcp://{host}:0", addr
+
+    sock = socket.socket()
+    manager = APIServerProcessManager(
+        listen_address=f"tcp://{host}:0",
+        sock=sock,
+        args="test_args",
+        num_servers=num_servers,
+        input_addresses=placeholder_inputs,
+        output_addresses=placeholder_outputs,
+        target_server_fn=defer_addresses_stub_worker,
+    )
+
+    try:
+        assert len(manager.processes) == num_servers
+        actual_inputs, actual_outputs = manager.gather_actual_addresses(timeout=15.0)
+    finally:
+        manager.shutdown()
+        time.sleep(0.2)
+        sock.close()
+
+    assert len(actual_inputs) == num_servers
+    assert len(actual_outputs) == num_servers
+
+    for addr in actual_inputs + actual_outputs:
+        scheme, parsed_host, port = split_zmq_path(addr)
+        assert scheme == "tcp", addr
+        assert parsed_host == host, addr
+        assert port and int(port) > 0, addr
+
+    all_addrs = actual_inputs + actual_outputs
+    assert len(set(all_addrs)) == len(all_addrs), all_addrs
+
+
+@pytest.mark.timeout(30)
+def test_gather_actual_addresses_child_crash_before_report():
+    """A child that exits before sending its endpoints must surface a
+    clear ``RuntimeError`` rather than hang or return ``None`` slots."""
+    host = "127.0.0.1"
+    num_servers = 2
+    placeholder_inputs = [
+        get_engine_client_zmq_addr(local_only=False, host=host)
+        for _ in range(num_servers)
+    ]
+    placeholder_outputs = [
+        get_engine_client_zmq_addr(local_only=False, host=host)
+        for _ in range(num_servers)
+    ]
+
+    sock = socket.socket()
+    manager = APIServerProcessManager(
+        listen_address=f"tcp://{host}:0",
+        sock=sock,
+        args="test_args",
+        num_servers=num_servers,
+        input_addresses=placeholder_inputs,
+        output_addresses=placeholder_outputs,
+        # mock_run_api_server_worker exits without touching
+        # ``actual_address_pipe`` — simulates a child that dies before
+        # reporting its bound addresses.
+        target_server_fn=mock_run_api_server_worker,
+    )
+    try:
+        # Sentinel-first vs pipe-EOF-first both produce "reporting".
+        with pytest.raises(RuntimeError, match="reporting"):
+            manager.gather_actual_addresses(timeout=10.0)
+    finally:
+        manager.shutdown()
+        time.sleep(0.2)
+        sock.close()

tests/kernels/core/test_minimax_reduce_rms.py

@@ -10,7 +10,7 @@ from torch.multiprocessing import spawn
 from tests.kernels.utils import opcheck
 from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
 from vllm.distributed import cleanup_dist_env_and_memory
-from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01RMSNormTP
+from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
 from vllm.platforms import current_platform
 from vllm.utils.network_utils import get_open_port
 from vllm.utils.torch_utils import set_random_seed
@@ -59,7 +59,7 @@ def _worker_forward_qk(
 
     # Set up Lamport workspace.
     from vllm.distributed.parallel_state import get_tp_group
-    from vllm.model_executor.layers.mamba.lamport_workspace import (
+    from vllm.model_executor.layers.minimax_rms_norm.lamport_workspace import (
         get_allreduce_workspace,
     )
 

tests/kernels/moe/test_cpu_int4_moe.py

@@ -0,0 +1,296 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for CPU INT4 W4A8 dynamic quantized fused MoE kernel (CPUExpertsInt4)."""
+
+import sys
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+if not current_platform.is_cpu():
+    pytest.skip("skipping CPU-only tests", allow_module_level=True)
+
+# Check if the dynamic_4bit_int_moe op is available
+if not hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
+    pytest.skip("dynamic_4bit_int_moe op not available", allow_module_level=True)
+
+# Check if KleidiAI ops are available
+if not hasattr(torch.ops.aten, "_dyn_quant_pack_4bit_weight"):
+    pytest.skip("KleidiAI 4-bit ops not available", allow_module_level=True)
+
+
+# Tolerance for INT4 W4A8
+INT4_W4A8_ATOL = 2e-2
+INT4_W4A8_RTOL = 2e-2
+
+
+def _silu_and_mul(x: torch.Tensor) -> torch.Tensor:
+    """SwiGLU activation: SiLU(gate) * up."""
+    d = x.shape[-1] // 2
+    return F.silu(x[..., :d]) * x[..., d:]
+
+
+def _pack_int4_weight_to_kleidi(
+    int4_as_int8: torch.Tensor,
+    scales: torch.Tensor,
+    bias: torch.Tensor | None,
+    group_size: int,
+    in_features: int,
+    out_features: int,
+) -> torch.Tensor:
+    """Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.
+
+    Args:
+        int4_as_int8: [out, in] int8 tensor with values in [-8, 7]
+        scales: [out, in//group_size] or [out, 1] for channel-wise
+        bias: [out] optional bias
+        group_size: Quantization group size (-1 for channel-wise)
+        in_features: Input dimension
+        out_features: Output dimension
+
+    Returns:
+        Packed weight tensor in KleidiAI format
+    """
+    # Shift to unsigned nibble [0, 15]
+    tmp = int4_as_int8.add(8)
+    # Pack pairs along input dimension
+    uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(torch.uint8)
+
+    # Determine scale dtype based on group_size
+    scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
+    scales_typed = scales.to(scale_dtype)
+    bias_typed = None if bias is None else bias.to(torch.float32)
+
+    # Pack using KleidiAI op
+    actual_group_size = in_features if group_size == -1 else group_size
+    return torch.ops.aten._dyn_quant_pack_4bit_weight(
+        uint8_nibbles,
+        scales_typed,
+        bias_typed,
+        actual_group_size,
+        in_features,
+        out_features,
+    )
+
+
+def _make_int4_moe_weights(
+    E: int,
+    N: int,
+    K: int,
+    group_size: int,
+    has_bias: bool = False,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor | None,
+    torch.Tensor | None,
+]:
+    """Generate random INT4 MoE weights with random scales.
+
+    Args:
+        E: Number of experts
+        N: Intermediate size
+        K: Hidden size
+        group_size: Quantization group size (-1 for channel-wise)
+        has_bias: Whether to include bias
+
+    Returns:
+        (w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias)
+        where *_ref are the dequantized float reference weights
+    """
+    # Generate INT4 weights as int8 values in [-8, 7]
+    w13_int4 = torch.randint(-8, 8, (E, 2 * N, K), dtype=torch.int8)
+    w2_int4 = torch.randint(-8, 8, (E, K, N), dtype=torch.int8)
+
+    # Determine number of scale columns
+    def _n_scale_cols(in_features: int) -> int:
+        return 1 if group_size == -1 else (in_features // group_size)
+
+    # Generate random scales
+    scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
+    w13_scales = torch.rand(E, 2 * N, _n_scale_cols(K), dtype=scale_dtype) * 0.01
+    w2_scales = torch.rand(E, K, _n_scale_cols(N), dtype=scale_dtype) * 0.01
+
+    # Generate biases if needed
+    w13_bias = None
+    w2_bias = None
+    if has_bias:
+        w13_bias = torch.randn(E, 2 * N, dtype=torch.float32) * 0.01
+        w2_bias = torch.randn(E, K, dtype=torch.float32) * 0.01
+
+    # Pack weights for each expert
+    w13_packed_list = []
+    w2_packed_list = []
+
+    for e in range(E):
+        w13_packed_list.append(
+            _pack_int4_weight_to_kleidi(
+                w13_int4[e],
+                w13_scales[e],
+                w13_bias[e] if (has_bias and w13_bias is not None) else None,
+                group_size,
+                K,
+                2 * N,
+            )
+        )
+        w2_packed_list.append(
+            _pack_int4_weight_to_kleidi(
+                w2_int4[e],
+                w2_scales[e],
+                w2_bias[e] if (has_bias and w2_bias is not None) else None,
+                group_size,
+                N,
+                K,
+            )
+        )
+
+    w13_packed = torch.stack(w13_packed_list, dim=0)
+    w2_packed = torch.stack(w2_packed_list, dim=0)
+
+    # Create reference dequantized weights
+    w13_ref = torch.zeros(E, 2 * N, K, dtype=torch.float32)
+    w2_ref = torch.zeros(E, K, N, dtype=torch.float32)
+
+    for e in range(E):
+        # Dequantize w13
+        for i in range(2 * N):
+            for j in range(K):
+                group_idx = 0 if group_size == -1 else (j // group_size)
+                w13_ref[e, i, j] = (
+                    w13_int4[e, i, j].float() * w13_scales[e, i, group_idx].float()
+                )
+                if has_bias and w13_bias is not None:
+                    w13_ref[e, i, j] += w13_bias[e, i].float()
+
+        # Dequantize w2
+        for i in range(K):
+            for j in range(N):
+                group_idx = 0 if group_size == -1 else (j // group_size)
+                w2_ref[e, i, j] = (
+                    w2_int4[e, i, j].float() * w2_scales[e, i, group_idx].float()
+                )
+                if has_bias and w2_bias is not None:
+                    w2_ref[e, i, j] += w2_bias[e, i].float()
+
+    return w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias
+
+
+def ref_int4_moe(
+    a: torch.Tensor,
+    w13_ref: torch.Tensor,
+    w2_ref: torch.Tensor,
+    topk_weight: torch.Tensor,
+    topk_ids: torch.Tensor,
+) -> torch.Tensor:
+    """Reference INT4 W4A8 fused MoE using dequantized weights.
+
+    Steps:
+      1. Use dequantized float weights
+      2. For each expert: matmul → SwiGLU → matmul
+      3. Weighted sum across top-k experts
+    """
+    B, D = a.shape
+    topk = topk_ids.size(1)
+
+    a_exp = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float()
+    out = torch.zeros(B * topk, w2_ref.shape[1], dtype=torch.float32)
+
+    topk_weight_flat = topk_weight.view(-1)
+    topk_ids_flat = topk_ids.view(-1)
+
+    for i in range(w13_ref.shape[0]):
+        mask = topk_ids_flat == i
+        if mask.sum():
+            # w13: [2N, K], input: [B, K] -> output: [B, 2N]
+            gate_up = torch.matmul(a_exp[mask], w13_ref[i].transpose(0, 1))
+            # SwiGLU activation
+            hidden = _silu_and_mul(gate_up)
+            # w2: [K, N], hidden: [B, N] -> output: [B, K]
+            out[mask] = torch.matmul(hidden, w2_ref[i].transpose(0, 1))
+
+    return (
+        (out.view(B, -1, w2_ref.shape[1]) * topk_weight_flat.view(B, -1, 1))
+        .sum(dim=1)
+        .to(a.dtype)
+    )
+
+
+NUM_TOKENS = [1, 2, 64, 128]
+# (intermediate_size N, hidden_size K, num_experts E, topk, group_size)
+MoE_CONFIGS = [
+    (256, 512, 8, 2, 128),
+    (256, 512, 8, 2, 64),
+    (256, 512, 8, 2, -1),  # channel-wise
+    (512, 256, 8, 4, 128),
+    (512, 512, 8, 2, 128),
+    (768, 2048, 8, 2, 128),
+    (768, 2048, 16, 4, 64),
+]
+SEEDS = [0, 42]
+
+
+@pytest.mark.parametrize("M", NUM_TOKENS)
+@pytest.mark.parametrize("N,K,E,topk,group_size", MoE_CONFIGS)
+@pytest.mark.parametrize("seed", SEEDS)
+def test_cpu_int4_moe_kernel(M, N, K, E, topk, group_size, seed):
+    """Test dynamic_4bit_int_moe kernel against dequantized torch reference."""
+    set_random_seed(seed)
+
+    # Generate input activations
+    a = torch.randn(M, K, dtype=torch.bfloat16) / (K**0.5)
+
+    # Generate INT4 weights
+    w13_packed, w2_packed, w13_ref, w2_ref, w13_bias, w2_bias = _make_int4_moe_weights(
+        E, N, K, group_size, has_bias=False
+    )
+
+    # Generate router logits and topk
+    score = torch.randn(M, E, dtype=torch.bfloat16)
+    score = torch.softmax(score, dim=-1, dtype=torch.float32)
+    topk_weight, topk_ids = torch.topk(score, topk)
+    topk_ids = topk_ids.to(torch.long)
+
+    # Reference output using dequantized weights
+    ref_out = ref_int4_moe(
+        a,
+        w13_ref,
+        w2_ref,
+        topk_weight,
+        topk_ids,
+    )
+
+    # Test dynamic_4bit_int_moe kernel
+    # Activation kind: 1 = SwiGLU_Ug (SiLU(u)*g) for OAI-style
+    activation_kind = 1
+    apply_router_weight_on_input = False
+
+    out = torch.ops._C.dynamic_4bit_int_moe(
+        a,
+        topk_ids,
+        topk_weight,
+        w13_packed,
+        w2_packed,
+        K,  # H (hidden_size / w2_out_features)
+        N,  # I (intermediate_size / w2_in_features)
+        2 * N,  # I2 (2*intermediate_size / w13_out_features)
+        group_size,
+        apply_router_weight_on_input,
+        activation_kind,
+    )
+
+    torch.testing.assert_close(
+        ref_out.bfloat16(),
+        out,
+        atol=INT4_W4A8_ATOL,
+        rtol=INT4_W4A8_RTOL,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

@@ -120,9 +120,19 @@ pytestmark = pytest.mark.skipif(
 )
 
 
-def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs):
-    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-        q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+def _call_fused(
+    q_in, q_head_padded, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+):
+    return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        q_in,
+        kv,
+        k_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        q_head_padded,
+        eps,
+        bs,
     )
 
 
@@ -130,8 +140,23 @@ def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
 
 
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 64, 2048])
-@pytest.mark.parametrize("n_heads", [8, 64])
-def test_q_path_matches_reference(num_tokens: int, n_heads: int):
+@pytest.mark.parametrize(
+    "n_heads,padded_heads",
+    [
+        # Each supported padded_heads instantiation: padded (n_heads <
+        # padded_heads) and unpadded (n_heads == padded_heads).
+        (1, 8),
+        (8, 8),
+        (8, 16),
+        (16, 16),
+        (16, 32),
+        (32, 32),
+        (8, 64),
+        (64, 64),
+        (64, 128),
+    ],
+)
+def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: int):
     torch.manual_seed(0)
     device = "cuda"
     dtype = torch.bfloat16
@@ -156,10 +181,16 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int):
         num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device
     ).view(num_blocks, -1)
     slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
-    q_fused = q.clone()
-    _call_fused(q_fused, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
+    q_out = _call_fused(
+        q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+    )
 
-    torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
+    if n_heads < padded_heads:
+        pad_region = q_out[:, n_heads:padded_heads]
+        assert pad_region.abs().max().item() == 0.0, (
+            "padded head slots must be exact zero"
+        )
 
 
 # ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
@@ -201,11 +232,12 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):
         kv_ref, k_cache_ref, slot_mapping, block_size=block_size
     )
 
-    # ── Fused path (dummy q, single head) ──────────────────────────────────
+    # ── Fused path (dummy q, padded to FlashMLA's min head count 64) ───────
     k_cache_fused = torch.zeros_like(k_cache_ref)
     q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device)
-    _call_fused(
+    _ = _call_fused(
         q_dummy,
+        64,
         kv,
         k_cache_fused,
         slot_mapping,
@@ -298,8 +330,9 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
     # Fused: pass full-sized q/kv/positions, shorter slot_mapping.
     q_dummy = torch.zeros(total, 1, HEAD_DIM, dtype=dtype, device=device)
     k_cache_fused = torch.zeros_like(k_cache_ref)
-    _call_fused(
+    _ = _call_fused(
         q_dummy,
+        64,
         kv,
         k_cache_fused,
         slot_mapping,
@@ -316,9 +349,26 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
 
 
 @pytest.mark.parametrize("num_tokens", [1, 4, 17, 2048])
-@pytest.mark.parametrize("n_heads", [8, 64])
+@pytest.mark.parametrize(
+    "n_heads,padded_heads",
+    [
+        # Each supported padded_heads instantiation: padded (n_heads <
+        # padded_heads) and unpadded (n_heads == padded_heads).
+        (1, 8),
+        (8, 8),
+        (8, 16),
+        (16, 16),
+        (16, 32),
+        (32, 32),
+        (8, 64),
+        (64, 64),
+        (64, 128),
+    ],
+)
 @pytest.mark.parametrize("block_size", [16, 64])
-def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
+def test_combined_q_and_kv(
+    num_tokens: int, n_heads: int, padded_heads: int, block_size: int
+):
     torch.manual_seed(2)
     device = "cuda"
     dtype = torch.bfloat16
@@ -345,10 +395,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
     )
 
     # Fused single call.
-    q_fused = q.clone()
     k_cache_fused = torch.zeros_like(k_cache_ref)
-    _call_fused(
-        q_fused,
+    q_out = _call_fused(
+        q,
+        padded_heads,
         kv,
         k_cache_fused,
         slot_mapping,
@@ -358,5 +408,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
         block_size,
     )
 
-    torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
+    if n_heads < padded_heads:
+        pad_region = q_out[:, n_heads:padded_heads]
+        assert pad_region.abs().max().item() == 0.0, (
+            "padded head slots must be exact zero"
+        )
     torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)

tests/kernels/test_fused_inv_rope_fp8_quant.py

@@ -8,7 +8,7 @@ the existing separate operations (inverse RoPE via rotate_neox + FP8 quant
 via per_token_group_quant_fp8).
 
 The reference faithfully reproduces the exact flow in
-deepseek_v4/nvidia/ops/attention.py:295-310:
+deepseek_v4/attention.py:295-310:
   1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head)
   2. Reshape [T, H, head_dim] -> [T, G, D]
   3. Transpose+flatten to [G*T, D], quantize, reshape back
@@ -668,7 +668,7 @@ def _unfused_inv_rope_fp8_quant(
     nope_dim: int = NOPE_DIM,
     rope_dim: int = ROPE_DIM,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Unfused path matching deepseek_v4/nvidia/ops/attention.py:295-310.
+    """Unfused path matching deepseek_v4/attention.py:295-310.
 
     Uses the production CUDA RoPE kernel + per_token_group_quant_fp8.
     """

tests/model_executor/test_nemotron_h_quantization.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import Mock, patch
+
+
+def test_nemotron_h_lm_head_receives_quant_config():
+    from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
+
+    mock_quant_config = Mock()
+
+    mock_hf_config = Mock()
+    mock_hf_config.vocab_size = 128
+    mock_hf_config.hidden_size = 64
+
+    mock_vllm_config = Mock()
+    mock_vllm_config.model_config.hf_config = mock_hf_config
+    mock_vllm_config.model_config.dtype = None
+    mock_vllm_config.scheduler_config = Mock()
+    mock_vllm_config.quant_config = mock_quant_config
+
+    with (
+        patch("vllm.model_executor.models.nemotron_h.NemotronHModel") as MockModel,
+        patch("vllm.model_executor.models.nemotron_h.ParallelLMHead") as MockLMHead,
+        patch("vllm.model_executor.models.nemotron_h.LogitsProcessor"),
+    ):
+        MockModel.return_value.make_empty_intermediate_tensors = Mock()
+        MockModel.return_value.has_moe = False
+
+        NemotronHForCausalLM(vllm_config=mock_vllm_config)
+
+        MockLMHead.assert_called_once()
+        call_kwargs = MockLMHead.call_args.kwargs
+        assert call_kwargs["quant_config"] is mock_quant_config

tests/model_executor/test_qwen3_5_quantization.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import Mock, patch
+
+
+def test_qwen3_5_lm_head_receives_quant_config():
+    from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLMBase
+
+    mock_quant_config = Mock()
+
+    mock_hf_config = Mock()
+    mock_hf_config.tie_word_embeddings = False
+    mock_hf_config.vocab_size = 128
+    mock_hf_config.hidden_size = 64
+
+    mock_vllm_config = Mock()
+    mock_vllm_config.model_config.hf_text_config = mock_hf_config
+    mock_vllm_config.cache_config.mamba_cache_mode = "align"
+    mock_vllm_config.scheduler_config = Mock()
+    mock_vllm_config.quant_config = mock_quant_config
+    mock_vllm_config.lora_config = None
+
+    mock_pp_group = Mock()
+    mock_pp_group.is_last_rank = True
+
+    with (
+        patch("vllm.model_executor.models.qwen3_5.Qwen3_5Model") as MockModel,
+        patch("vllm.model_executor.models.qwen3_5.ParallelLMHead") as MockLMHead,
+        patch("vllm.model_executor.models.qwen3_5.LogitsProcessor"),
+        patch(
+            "vllm.model_executor.models.qwen3_5.get_pp_group",
+            return_value=mock_pp_group,
+        ),
+    ):
+        MockModel.return_value.make_empty_intermediate_tensors = Mock()
+
+        Qwen3_5ForCausalLMBase(vllm_config=mock_vllm_config)
+
+        MockLMHead.assert_called_once()
+        call_kwargs = MockLMHead.call_args.kwargs
+        assert call_kwargs["quant_config"] is mock_quant_config
+
+
+def test_qwen3_5_mtp_lm_head_receives_quant_config():
+    from vllm.config import CompilationMode
+    from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP
+
+    mock_quant_config = Mock()
+
+    mock_hf_config = Mock()
+    mock_hf_config.tie_word_embeddings = False
+    mock_hf_config.vocab_size = 128
+    mock_hf_config.hidden_size = 64
+
+    mock_vllm_config = Mock()
+    mock_vllm_config.model_config.hf_text_config = mock_hf_config
+    mock_vllm_config.cache_config.mamba_cache_mode = "align"
+    mock_vllm_config.compilation_config.mode = CompilationMode.NONE
+    mock_vllm_config.quant_config = mock_quant_config
+
+    mock_pp_group = Mock()
+    mock_pp_group.is_last_rank = True
+
+    with (
+        patch("vllm.model_executor.models.qwen3_5_mtp.Qwen3_5MultiTokenPredictor"),
+        patch("vllm.model_executor.models.qwen3_5_mtp.ParallelLMHead") as MockLMHead,
+        patch("vllm.model_executor.models.qwen3_5_mtp.LogitsProcessor"),
+        patch(
+            "vllm.model_executor.models.qwen3_5_mtp.get_pp_group",
+            return_value=mock_pp_group,
+        ),
+    ):
+        Qwen3_5MTP(vllm_config=mock_vllm_config)
+
+        MockLMHead.assert_called_once()
+        call_kwargs = MockLMHead.call_args.kwargs
+        assert call_kwargs["quant_config"] is mock_quant_config

tests/quantization/test_modelopt.py

@@ -7,13 +7,24 @@ Run `pytest tests/quantization/test_modelopt.py`.
 
 import os
 from typing import Any, NoReturn
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, Mock, patch
 
 import pytest
 import torch
 
 from tests.quantization.utils import is_quant_method_supported
 from vllm.config.model import ModelConfig
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.modelopt import (
+    ModelOptFp8Config,
+    ModelOptMixedPrecisionConfig,
+    ModelOptNvFp4Config,
+    ModelOptNvFp4LinearMethod,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 
 
 @pytest.fixture(scope="function", autouse=True)
@@ -44,6 +55,87 @@ def _snapshot_download_or_skip(model_id: str) -> str:
         _skip(f"Failed to download {model_id} from the HF Hub: {e}")
 
 
+def _mock_lm_head() -> Mock:
+    lm_head = Mock(spec=ParallelLMHead)
+    lm_head.__class__ = ParallelLMHead
+    return lm_head
+
+
+def _mixed_precision_config(quantized_layers: dict) -> ModelOptMixedPrecisionConfig:
+    return ModelOptMixedPrecisionConfig(
+        kv_cache_quant_method=None,
+        exclude_modules=[],
+        quantized_layers=quantized_layers,
+        fp8_config=ModelOptFp8Config(
+            quant_method="FP8",
+            is_checkpoint_fp8_serialized=True,
+            kv_cache_quant_method=None,
+            exclude_modules=[],
+        ),
+        nvfp4_config=ModelOptNvFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            kv_cache_quant_algo=None,
+            exclude_modules=[],
+        ),
+        w4a16_nvfp4_config=ModelOptNvFp4Config(
+            quant_method="W4A16_NVFP4",
+            is_checkpoint_nvfp4_serialized=True,
+            kv_cache_quant_algo=None,
+            exclude_modules=[],
+        ),
+    )
+
+
+def test_modelopt_nvfp4_quantizes_parallel_lm_head():
+    config = ModelOptNvFp4Config(
+        is_checkpoint_nvfp4_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+    )
+
+    with patch(
+        "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
+    ):
+        method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
+
+    assert isinstance(method, ModelOptNvFp4LinearMethod)
+
+
+def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized():
+    config = ModelOptNvFp4Config(
+        is_checkpoint_nvfp4_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=["lm_head"],
+    )
+
+    method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
+
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+def test_modelopt_mixed_precision_quantizes_parallel_lm_head():
+    config = _mixed_precision_config(
+        {"lm_head": {"quant_algo": "NVFP4", "group_size": 16}}
+    )
+
+    with patch(
+        "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
+    ):
+        method = config.get_quant_method(_mock_lm_head(), prefix="lm_head")
+
+    assert isinstance(method, ModelOptNvFp4LinearMethod)
+
+
+def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale():
+    holder = Mock()
+    scale = torch.nn.Parameter(torch.empty(1))
+    loaded_scale = torch.tensor(2.0)
+
+    VocabParallelEmbedding.weight_loader(holder, scale, loaded_scale)
+
+    assert torch.equal(scale, loaded_scale.reshape(1))
+
+
 @pytest.mark.skipif(
     not is_quant_method_supported("modelopt"),
     reason="ModelOpt FP8 is not supported on this GPU type.",

tests/tool_parsers/test_minicpm5xml_tool_parser.py

@@ -0,0 +1,735 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
+
+import json
+import random
+from unittest.mock import MagicMock
+
+import pytest
+
+from tests.tool_parsers.utils import run_tool_extraction_streaming
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+    ChatCompletionToolsParam,
+)
+from vllm.entrypoints.openai.engine.protocol import FunctionCall, ToolCall
+from vllm.tool_parsers import ToolParser, ToolParserManager
+from vllm.tool_parsers.minicpm5xml_tool_parser import MiniCPM5XMLToolParser
+
+
+def _tool(name: str, parameters: dict) -> ChatCompletionToolsParam:
+    return ChatCompletionToolsParam(
+        type="function",
+        function={
+            "name": name,
+            "parameters": parameters,
+        },
+    )
+
+
+def make_tools_weather() -> list[ChatCompletionToolsParam]:
+    return [
+        _tool(
+            "get_weather",
+            {
+                "type": "object",
+                "properties": {
+                    "city": {"type": "string"},
+                    "date": {"type": "string"},
+                },
+                "required": ["city"],
+            },
+        )
+    ]
+
+
+def make_tools_sum() -> list[ChatCompletionToolsParam]:
+    return [
+        _tool(
+            "sum_values",
+            {
+                "type": "object",
+                "properties": {
+                    "nums": {"type": "array"},
+                    "exact": {"type": "boolean"},
+                },
+                "required": ["nums"],
+            },
+        )
+    ]
+
+
+def make_tools_no_required() -> list[ChatCompletionToolsParam]:
+    return [
+        _tool(
+            "noop",
+            {
+                "type": "object",
+                "properties": {"note": {"type": "string"}},
+                "required": [],
+            },
+        )
+    ]
+
+
+def make_request(
+    tools: list[ChatCompletionToolsParam],
+    tool_choice: str = "auto",
+) -> ChatCompletionRequest:
+    return ChatCompletionRequest(
+        model="test-model",
+        messages=[{"role": "user", "content": "hi"}],
+        tools=tools,
+        tool_choice=tool_choice,
+    )
+
+
+def make_tool_call(name: str, arguments: dict) -> ToolCall:
+    return ToolCall(
+        type="function",
+        function=FunctionCall(
+            name=name,
+            arguments=json.dumps(arguments, ensure_ascii=False),
+        ),
+    )
+
+
+def assert_tool_calls(
+    actual: list[ToolCall],
+    expected: list[ToolCall],
+) -> None:
+    assert len(actual) == len(expected)
+    for act, exp in zip(actual, expected):
+        assert act.type == "function"
+        assert act.function.name == exp.function.name
+        assert json.loads(act.function.arguments) == json.loads(exp.function.arguments)
+
+
+@pytest.fixture
+def parser() -> ToolParser:
+    mock_tokenizer = MagicMock()
+    return MiniCPM5XMLToolParser(mock_tokenizer)
+
+
+def test_registered_in_tool_parser_manager() -> None:
+    cls = ToolParserManager.get_tool_parser("minicpm5")
+    assert cls is MiniCPM5XMLToolParser
+
+
+def test_adjust_request_skip_special_tokens(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    assert request.skip_special_tokens is True
+    adjusted = parser.adjust_request(request)
+    assert adjusted.skip_special_tokens is False
+
+
+def test_adjust_request_tool_choice_none(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather(), tool_choice="none")
+    adjusted = parser.adjust_request(request)
+    assert adjusted.skip_special_tokens is True
+
+
+def test_no_tool_call(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    out = parser.extract_tool_calls("How can I help you?", request)
+    assert not out.tools_called
+    assert out.tool_calls == []
+    assert out.content == "How can I help you?"
+
+
+def test_single_call_with_surrounding_text(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        "Intro before.\n"
+        '<function name="get_weather">'
+        '<param name="city">上海</param>'
+        '<param name="date">2024-06-27</param>'
+        "</function>\n"
+        "Outro after.\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert out.tools_called
+    assert_tool_calls(
+        out.tool_calls,
+        [
+            make_tool_call(
+                "get_weather",
+                {
+                    "city": "上海",
+                    "date": "2024-06-27",
+                },
+            )
+        ],
+    )
+    assert out.content is None
+
+
+def test_cdata_multiline(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function name="get_weather">'
+        '<param name="city"><![CDATA[北\n京]]></param>'
+        '<param name="date">2024-06-27</param>'
+        "</function>\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args["city"] == "北\n京"
+    assert args["date"] == "2024-06-27"
+
+
+def test_tokenizer_space_marker(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function\u0120name="get_weather">'
+        '<param\u0120name="city">上海</param>'
+        '<param\u0120name="date">2024-06-27</param>'
+        "</function>\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert out.tools_called
+    assert_tool_calls(
+        out.tool_calls,
+        [
+            make_tool_call(
+                "get_weather",
+                {
+                    "city": "上海",
+                    "date": "2024-06-27",
+                },
+            )
+        ],
+    )
+
+
+def test_collapsed_function_and_param_tags(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<functionname="get_weather">'
+        '<paramname="city">上海</param>'
+        '<paramname="date">2024-06-27</param>'
+        "</function>\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert out.tools_called
+    assert_tool_calls(
+        out.tool_calls,
+        [
+            make_tool_call(
+                "get_weather",
+                {
+                    "city": "上海",
+                    "date": "2024-06-27",
+                },
+            )
+        ],
+    )
+
+
+def test_collapsed_param_tags_with_tokenizer_space_function(
+    parser: ToolParser,
+) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function\u0120name="get_weather">'
+        '<paramname="city">上海</param>'
+        '<paramname="date">2024-06-27</param>'
+        "</function>\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert out.tools_called
+    assert_tool_calls(
+        out.tool_calls,
+        [
+            make_tool_call(
+                "get_weather",
+                {
+                    "city": "上海",
+                    "date": "2024-06-27",
+                },
+            )
+        ],
+    )
+
+
+def test_extract_tool_calls_streaming_partial_chunks(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    chunks = [
+        '<function name="get_weather">',
+        '<param name="city">',
+        '上海</param><param name="date">2024-06-27</param></function>\n',
+    ]
+    reconstructor = run_tool_extraction_streaming(
+        parser,
+        chunks,
+        request,
+        assert_one_tool_per_delta=False,
+    )
+    assert len(reconstructor.tool_calls) == 1
+    assert reconstructor.tool_calls[0].function.name == "get_weather"
+    assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
+        "city": "上海",
+        "date": "2024-06-27",
+    }
+
+
+def test_extract_tool_calls_streaming_tokenizer_space_marker(
+    parser: ToolParser,
+) -> None:
+    request = make_request(make_tools_weather())
+    chunks = [
+        '<function\u0120name="get_weather">',
+        '<param\u0120name="city">',
+        '上海</param><param\u0120name="date">2024-06-27</param></function>\n',
+    ]
+    reconstructor = run_tool_extraction_streaming(
+        parser,
+        chunks,
+        request,
+        assert_one_tool_per_delta=False,
+    )
+    assert len(reconstructor.tool_calls) == 1
+    assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
+        "city": "上海",
+        "date": "2024-06-27",
+    }
+
+
+def test_extract_tool_calls_streaming_collapsed_tags_weather(
+    parser: ToolParser,
+) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<functionname="get_weather">'
+        '<paramname="city">上海</param>'
+        '<paramname="date">2024-06-27</param>'
+        "</function>\n"
+    )
+    random.seed(2)
+    reconstructor = run_tool_extraction_streaming(
+        parser,
+        _random_chunks(text, 1, 4),
+        request,
+    )
+    assert len(reconstructor.tool_calls) == 1
+    assert reconstructor.tool_calls[0].function.name == "get_weather"
+    assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
+        "city": "上海",
+        "date": "2024-06-27",
+    }
+
+
+def test_collapsed_tags_current_weather(parser: ToolParser) -> None:
+    request = make_request(
+        [
+            _tool(
+                "get_current_weather",
+                {
+                    "type": "object",
+                    "properties": {
+                        "city": {"type": "string"},
+                        "state": {"type": "string"},
+                        "unit": {"type": "string"},
+                    },
+                    "required": ["city", "state", "unit"],
+                },
+            )
+        ]
+    )
+    text = (
+        '<functionname="get_current_weather">'
+        '<paramname="city">Dallas</param>'
+        '<paramname="state">TX</param>'
+        '<paramname="unit">fahrenheit</param>'
+        "</function>"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert out.tools_called
+    assert json.loads(out.tool_calls[0].function.arguments) == {
+        "city": "Dallas",
+        "state": "TX",
+        "unit": "fahrenheit",
+    }
+
+
+def test_extract_tool_calls_streaming_incremental_arguments(
+    parser: ToolParser,
+) -> None:
+    request = make_request(
+        [
+            _tool(
+                "get_current_weather",
+                {
+                    "type": "object",
+                    "properties": {
+                        "city": {"type": "string"},
+                        "state": {"type": "string"},
+                        "unit": {"type": "string"},
+                    },
+                    "required": ["city", "state", "unit"],
+                },
+            )
+        ]
+    )
+    text = (
+        '<function name="get_current_weather">'
+        '<param name="city">Dallas</param>'
+        '<param name="state">TX</param>'
+        '<param name="unit">fahrenheit</param>'
+        "</function>"
+    )
+    prev = ""
+    arguments = ""
+    for chunk in [text[:37], text[37:70], text[70:]]:
+        current = prev + chunk
+        delta = parser.extract_tool_calls_streaming(
+            prev,
+            current,
+            chunk,
+            [],
+            [],
+            [],
+            request,
+        )
+        if delta and delta.tool_calls:
+            arg_delta = delta.tool_calls[0].function.arguments
+            if arg_delta:
+                arguments += arg_delta
+        prev = current
+
+    assert json.loads(arguments) == {
+        "city": "Dallas",
+        "state": "TX",
+        "unit": "fahrenheit",
+    }
+
+
+def test_unknown_tool_block_preserved(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = '<function name="unknown"><param name="x">1</param></function>\n'
+    out = parser.extract_tool_calls(text, request)
+    assert not out.tools_called
+    assert "unknown" in (out.content or "")
+
+
+def test_non_string_types(parser: ToolParser) -> None:
+    request = make_request(make_tools_sum())
+    text = (
+        '<function name="sum_values">'
+        '<param name="nums">[1, 2, 3]</param>'
+        '<param name="exact">true</param>'
+        "</function>\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args["nums"] == [1, 2, 3]
+    assert args["exact"] is True
+
+
+def test_multiple_calls_interleaved_text(parser: ToolParser) -> None:
+    tools = make_tools_weather() + make_tools_sum()
+    request = make_request(tools)
+    text = (
+        "Head\n"
+        '<function name="get_weather"><param name="city">北京</param></function>\n'
+        "TXT\n"
+        '<function name="sum_values"><param name="nums">[7,8,9]</param>'
+        '<param name="exact">false</param></function>\n'
+        "Tail\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 2
+    args0 = json.loads(out.tool_calls[0].function.arguments)
+    assert args0["city"] == "北京"
+    args1 = json.loads(out.tool_calls[1].function.arguments)
+    assert args1["nums"] == [7, 8, 9]
+    assert args1["exact"] is False
+    assert out.content is None
+
+
+def test_incomplete_missing_function_end(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = '<function name="get_weather"><param name="city">北京</param>'
+    out = parser.extract_tool_calls(text, request)
+    assert not out.tools_called
+    assert "get_weather" in (out.content or "")
+
+
+def test_param_missing_name_invalid(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function name="get_weather">'
+        "<param>北京</param>"
+        '<param name="date">2024-06-27</param>'
+        "</function>\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert not out.tools_called
+    assert "<param>北京</param>" in (out.content or "")
+
+
+def test_duplicate_param_names_invalid(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function name="get_weather">'
+        '<param name="city">北京</param>'
+        '<param name="city">上海</param>'
+        "</function>\n"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert not out.tools_called
+
+
+def test_case_sensitive_param_name_invalid(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = '<function name="get_weather"><param name="City">北京</param></function>\n'
+    out = parser.extract_tool_calls(text, request)
+    assert not out.tools_called
+
+
+def test_no_required_and_zero_param_valid(parser: ToolParser) -> None:
+    request = make_request(make_tools_no_required())
+    text = '<function name="noop"></function>\n'
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args == {}
+
+
+def test_thinking_only_sentencepiece_normalized_in_content(
+    parser: ToolParser,
+) -> None:
+    request = make_request(make_tools_weather())
+    text = "\u010aFirst,\u0120I\u0120need\u0120to\u0120check\u0120the\u0120weather."
+    out = parser.extract_tool_calls(text, request)
+    assert not out.tools_called
+    assert out.content is not None
+    assert "\u0120" not in out.content
+    assert "\u010a" not in out.content
+    assert "First, I need to check the weather." in out.content
+
+
+def test_properties_wrapped_arguments(parser: ToolParser) -> None:
+    tools = [
+        _tool(
+            "get_customer_by_phone",
+            {
+                "type": "object",
+                "properties": {"phone_number": {"type": "string"}},
+                "required": ["phone_number"],
+            },
+        )
+    ]
+    request = make_request(tools)
+    text = (
+        '<function name="get_customer_by_phone">'
+        "<param name=\"properties\">{'phone_number': '555-123-2002'}</param>"
+        "</function>"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    assert out.tool_calls[0].function.name == "get_customer_by_phone"
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args == {"phone_number": "555-123-2002"}
+
+
+def test_arguments_wrapped_arguments(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function name="get_weather">'
+        '<param name="arguments">{"city": "上海", "date": "2024-06-27"}</param>'
+        "</function>"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args == {"city": "上海", "date": "2024-06-27"}
+
+
+def test_wrapped_arguments_still_validate_schema(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function name="get_weather">'
+        '<param name="properties">{"unknown": "x"}</param>'
+        "</function>"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert not out.tools_called
+
+
+def test_extra_arguments_ignored_when_required_present(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function name="get_weather">'
+        '<param name="city">上海</param>'
+        '<param name="unknown">ignored</param>'
+        "</function>"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args == {"city": "上海"}
+
+
+def test_extra_arguments_do_not_satisfy_required(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        '<function name="get_weather"><param name="unknown">ignored</param></function>'
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert not out.tools_called
+
+
+def test_zero_arg_tool_ignores_extra_arguments(parser: ToolParser) -> None:
+    request = make_request(make_tools_no_required())
+    text = (
+        '<function name="noop">'
+        '<param name="note">ignored</param>'
+        '<param name="extra">ignored</param>'
+        "</function>"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args == {"note": "ignored"}
+
+
+def test_alias_get_details_by_phone(parser: ToolParser) -> None:
+    tools = [
+        _tool(
+            "get_customer_by_phone",
+            {
+                "type": "object",
+                "properties": {"phone_number": {"type": "string"}},
+                "required": ["phone_number"],
+            },
+        )
+    ]
+    request = make_request(tools)
+    text = (
+        '<function name="get_details_by_phone">'
+        '<param name="phone_number">555-123-2002</param>'
+        "</function>"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    assert out.tool_calls[0].function.name == "get_customer_by_phone"
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args == {"phone_number": "555-123-2002"}
+
+
+def test_alias_get_line_details(parser: ToolParser) -> None:
+    tools = [
+        _tool(
+            "get_details_by_id",
+            {
+                "type": "object",
+                "properties": {"id": {"type": "string"}},
+                "required": ["id"],
+            },
+        )
+    ]
+    request = make_request(tools)
+    text = (
+        '<function name="get_line_details">'
+        '<param name="customer_id">C1001</param>'
+        '<param name="line_id">L1001</param>'
+        "</function>"
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    assert out.tool_calls[0].function.name == "get_details_by_id"
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args == {"id": "L1001"}
+
+
+def test_alias_enable_roaming(parser: ToolParser) -> None:
+    tools = [
+        _tool(
+            "toggle_roaming",
+            {
+                "type": "object",
+                "properties": {
+                    "line_id": {"type": "string"},
+                    "enabled": {"type": "boolean"},
+                },
+                "required": ["line_id", "enabled"],
+            },
+        )
+    ]
+    request = make_request(tools)
+    text = (
+        '<function name="enable_roaming"><param name="line_id">L1001</param></function>'
+    )
+    out = parser.extract_tool_calls(text, request)
+    assert len(out.tool_calls) == 1
+    assert out.tool_calls[0].function.name == "toggle_roaming"
+    args = json.loads(out.tool_calls[0].function.arguments)
+    assert args == {"line_id": "L1001", "enabled": True}
+
+
+def _random_chunks(text: str, min_len: int, max_len: int) -> list[str]:
+    chunks: list[str] = []
+    index = 0
+    while index < len(text):
+        size = random.randint(min_len, max_len)
+        chunks.append(text[index : index + size])
+        index += size
+    return chunks
+
+
+def test_extract_tool_calls_streaming_single(parser: ToolParser) -> None:
+    request = make_request(make_tools_weather())
+    text = (
+        "Intro before.\n"
+        '<function name="get_weather">'
+        '<param name="city">上海</param>'
+        '<param name="date">2024-06-27</param>'
+        "</function>\n"
+        "Outro after.\n"
+    )
+    random.seed(0)
+    reconstructor = run_tool_extraction_streaming(
+        parser,
+        _random_chunks(text, 1, 4),
+        request,
+    )
+    assert "Intro before." in reconstructor.other_content
+    assert "Outro after." in reconstructor.other_content
+    assert len(reconstructor.tool_calls) == 1
+    assert reconstructor.tool_calls[0].function.name == "get_weather"
+    assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
+        "city": "上海",
+        "date": "2024-06-27",
+    }
+
+
+def test_extract_tool_calls_streaming_multiple(parser: ToolParser) -> None:
+    tools = make_tools_weather() + make_tools_sum()
+    request = make_request(tools)
+    text = (
+        "Head\n"
+        '<function name="get_weather"><param name="city">北京</param></function>\n'
+        "TXT\n"
+        '<function name="sum_values"><param name="nums">[7,8,9]</param>'
+        '<param name="exact">false</param></function>\n'
+        "Tail\n"
+    )
+    random.seed(1)
+    reconstructor = run_tool_extraction_streaming(
+        parser,
+        _random_chunks(text, 1, 5),
+        request,
+    )
+    assert "Head" in reconstructor.other_content
+    assert "TXT" in reconstructor.other_content
+    assert "Tail" in reconstructor.other_content
+    assert len(reconstructor.tool_calls) == 2
+    assert json.loads(reconstructor.tool_calls[0].function.arguments)["city"] == "北京"
+    assert json.loads(reconstructor.tool_calls[1].function.arguments) == {
+        "nums": [7, 8, 9],
+        "exact": False,
+    }

tests/utils.py

@@ -1363,43 +1363,38 @@ def multi_process_parallel(
 ) -> None:
     import ray
 
-    # Using ray helps debugging the error when it failed
-    # as compared to multiprocessing.
-    # NOTE: We need to set working_dir for distributed tests,
-    # otherwise we may get import errors on ray workers
-    # NOTE: Force ray not to use gitignore file as excluding, otherwise
-    # it will not move .so files to working dir.
-    # So we have to manually add some of large directories
-    os.environ["RAY_RUNTIME_ENV_IGNORE_GITIGNORE"] = "1"
+    # Using ray helps debugging the error when it failed as compared to
+    # multiprocessing. For local Ray workers, putting the repo root on
+    # PYTHONPATH is enough and avoids uploading the full source tree, which
+    # exceeds Ray's working_dir package size limit on CI.
+    env_vars = {
+        "PYTHONPATH": os.pathsep.join(
+            filter(None, [str(VLLM_PATH), os.environ.get("PYTHONPATH")])
+        ),
+        **{env_var: "1" for env_var in current_platform.ray_noset_device_env_vars},
+    }
     ray.init(
         runtime_env={
-            "working_dir": VLLM_PATH,
-            "excludes": [
-                "build",
-                ".git",
-                "cmake-build-*",
-                "shellcheck",
-                "dist",
-                "ep_kernels_workspace",
-            ],
+            "env_vars": env_vars,
         }
     )
 
     distributed_init_port = get_open_port()
-    refs = []
-    for rank in range(tp_size * pp_size):
-        refs.append(
-            test_target.remote(
-                monkeypatch,
-                tp_size,
-                pp_size,
-                rank,
-                distributed_init_port,
-            ),
-        )
-    ray.get(refs)
-
-    ray.shutdown()
+    try:
+        refs = []
+        for rank in range(tp_size * pp_size):
+            refs.append(
+                test_target.remote(
+                    monkeypatch,
+                    tp_size,
+                    pp_size,
+                    rank,
+                    distributed_init_port,
+                ),
+            )
+        ray.get(refs)
+    finally:
+        ray.shutdown()
 
 
 @contextmanager

tests/v1/attention/test_attention_backends.py

@@ -136,7 +136,8 @@ def create_and_prepopulate_kv_cache(
     block_table = common_attn_metadata.block_table_tensor
     slot_mapping = common_attn_metadata.slot_mapping
 
-    # Create KV cache
+    # Create KV cache and populate in (2, num_blocks, ...) layout for easy
+    # flat indexing, then transpose to (num_blocks, 2, ...) layout.
     kv_cache = torch.zeros(
         2, num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device
     )
@@ -155,6 +156,9 @@ def create_and_prepopulate_kv_cache(
         # Stay block aligned and allocate enough blocks for the new tokens
         start_block_idx += cdiv(int(seq_lens[i]), block_size)
 
+    # Transpose to (num_blocks, 2, ...) layout
+    kv_cache = kv_cache.transpose(0, 1).contiguous()
+
     blocks_end = start_block_idx
 
     # Permute the context blocks (excluding block 0 which is null)
@@ -168,7 +172,7 @@ def create_and_prepopulate_kv_cache(
     inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
     # Add 1 to account for starting from block 1
     inv_perm[1:] = torch.argsort(perm) + 1
-    kv_cache[:, 1:blocks_end, ...] = kv_cache[:, perm, ...]
+    kv_cache[1:blocks_end, ...] = kv_cache[perm, ...]
 
     # Construct the right block table
     # Start from block_id=1 since block_id=0 is considered the null block
@@ -473,28 +477,35 @@ def _test_backend_correctness(
     # Note: flex_attention has known Triton kernel compatibility issues
     # with test infrastructures
     for backend_name in backend_to_test:
-        # FlashAttentionm + FlexAttention:
-        #   [2, num_blocks, block_size, num_kv_heads, head_size]
-        # FlashInfer + Triton:
-        #   [num_blocks, 2, block_size, num_kv_heads, head_size]
-        # Select the appropriate KV cache format for each backend
-        kv_cache_for_backend = kv_cache
         reset_kv_cache_layout = False
-        if backend_name in (
-            AttentionBackendEnum.FLASHINFER,
-            AttentionBackendEnum.TRITON_ATTN,
-        ):
-            kv_cache_for_backend = kv_cache.transpose(0, 1)
+
+        # Resolve backend class for both enum and string names.
+        actual_backend = backend_name
+        if backend_name == "FLEX_ATTENTION_SLOW":
+            actual_backend = AttentionBackendEnum.FLEX_ATTENTION
+        if hasattr(actual_backend, "get_class"):
+            backend_cls = actual_backend.get_class()
+        else:
+            backend_cls = None
 
         if backend_name == AttentionBackendEnum.FLASHINFER:
-            # For FlashInfer default to HND layout and
-            kv_cache_for_backend = (
-                kv_cache_for_backend.transpose(2, 3).contiguous().transpose(2, 3)
-            )
             set_kv_cache_layout("HND")
             reset_kv_cache_layout = True
-        elif backend_name == AttentionBackendEnum.TRITON_ATTN:
-            kv_cache_for_backend = kv_cache_for_backend.contiguous()
+
+        # Apply stride order like runtime does in
+        # _reshape_kv_cache (attn_utils.py:182-210): permute to physical
+        # layout, make contiguous, then permute to logical layout.
+        kv_cache_for_backend = kv_cache
+        if backend_cls is not None:
+            try:
+                stride_order = backend_cls.get_kv_cache_stride_order()
+            except (AttributeError, NotImplementedError):
+                stride_order = tuple(range(kv_cache.ndim))
+            if stride_order != tuple(range(kv_cache.ndim)):
+                inv_order = [stride_order.index(i) for i in range(len(stride_order))]
+                kv_cache_for_backend = (
+                    kv_cache.permute(*stride_order).contiguous().permute(*inv_order)
+                )
 
         try:
             backend_output = run_attention_backend(

tests/v1/attention/test_kv_head_stride_canonicalization.py

@@ -19,7 +19,7 @@ size-1 dimensions via torch.as_strided — zero-copy.
 
 The degenerate stride manifests at different positions in different backends:
 - FlashInfer: stride(-3) after kv_cache.permute() → shape [..., 1, B, D]
-- FlashAttention: stride(-2) after kv_cache.unbind(0) → shape [N, B, 1, D]
+- FlashAttention: stride(-2) after kv_cache.unbind(1) → shape [N, B, 1, D]
 """
 
 import torch

tests/v1/attention/test_mla_prefill_registry.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for MLA prefill backend registry."""
+
+import pytest
+import torch
+
+from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
+from vllm.v1.attention.backends.mla.prefill.registry import (
+    MLAPrefillBackendEnum,
+    register_mla_prefill_backend,
+)
+
+
+class CustomMLAPrefillBackend(MLAPrefillBackend):
+    """Mock custom MLA prefill backend for testing."""
+
+    supported_dtypes = [torch.bfloat16, torch.float16]
+    requires_r1_mla_dimensions = False
+
+    @staticmethod
+    def get_name() -> str:
+        return "CUSTOM"
+
+    def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
+        raise NotImplementedError
+
+    def run_prefill_context_chunk(self, chunk_idx, q, k, v):
+        raise NotImplementedError
+
+
+@pytest.fixture(autouse=True)
+def cleanup_overrides():
+    """Clear any overrides after each test."""
+    yield
+    for member in MLAPrefillBackendEnum:
+        member.clear_override()
+
+
+def test_custom_is_not_alias_of_any_backend():
+    all_backends = list(MLAPrefillBackendEnum)
+
+    aliases = []
+    for backend in all_backends:
+        if backend.name != "CUSTOM" and backend is MLAPrefillBackendEnum.CUSTOM:
+            aliases.append(backend.name)
+
+    assert len(aliases) == 0, (
+        f"BUG! CUSTOM is an alias of: {', '.join(aliases)}!\n"
+        f"CUSTOM.value = {repr(MLAPrefillBackendEnum.CUSTOM.value)}\n"
+        f"All MLA prefill backend values:\n"
+        + "\n".join(f"  {b.name}: {repr(b.value)}" for b in all_backends)
+    )
+
+    assert MLAPrefillBackendEnum.CUSTOM.name == "CUSTOM"
+
+
+def test_custom_unregistered_raises():
+    with pytest.raises(ValueError, match="must be registered before use"):
+        MLAPrefillBackendEnum.CUSTOM.get_path()
+
+
+def test_register_custom_backend_with_class_path():
+    register_mla_prefill_backend(
+        backend=MLAPrefillBackendEnum.CUSTOM,
+        class_path=(
+            "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
+        ),
+    )
+
+    assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
+
+    class_path = MLAPrefillBackendEnum.CUSTOM.get_path()
+    assert class_path == (
+        "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
+    )
+
+    backend_cls = MLAPrefillBackendEnum.CUSTOM.get_class()
+    assert backend_cls.get_name() == "CUSTOM"
+
+
+def test_register_custom_backend_as_decorator():
+    @register_mla_prefill_backend(MLAPrefillBackendEnum.CUSTOM)
+    class DecoratedPrefillBackend(MLAPrefillBackend):
+        supported_dtypes = [torch.bfloat16]
+        requires_r1_mla_dimensions = False
+
+        @staticmethod
+        def get_name() -> str:
+            return "DECORATED"
+
+        def run_prefill_new_tokens(self, q, k, v, return_softmax_lse):
+            raise NotImplementedError
+
+        def run_prefill_context_chunk(self, chunk_idx, q, k, v):
+            raise NotImplementedError
+
+    assert MLAPrefillBackendEnum.CUSTOM.is_overridden()
+    assert "DecoratedPrefillBackend" in MLAPrefillBackendEnum.CUSTOM.get_path()
+
+
+def test_override_existing_backend():
+    original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path()
+
+    register_mla_prefill_backend(
+        backend=MLAPrefillBackendEnum.FLASH_ATTN,
+        class_path=(
+            "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
+        ),
+    )
+
+    assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
+    assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() != original_path
+
+    backend_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
+    assert backend_cls.get_name() == "CUSTOM"
+
+
+def test_clear_override():
+    original_path = MLAPrefillBackendEnum.FLASH_ATTN.get_path()
+
+    register_mla_prefill_backend(
+        backend=MLAPrefillBackendEnum.FLASH_ATTN,
+        class_path=(
+            "tests.v1.attention.test_mla_prefill_registry.CustomMLAPrefillBackend"
+        ),
+    )
+    assert MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
+
+    MLAPrefillBackendEnum.FLASH_ATTN.clear_override()
+    assert not MLAPrefillBackendEnum.FLASH_ATTN.is_overridden()
+    assert MLAPrefillBackendEnum.FLASH_ATTN.get_path() == original_path
+
+
+def test_unknown_backend_name_raises():
+    with pytest.raises(ValueError, match="Unknown MLA prefill backend"):
+        MLAPrefillBackendEnum["NONEXISTENT"]

tests/v1/core/test_scheduler.py

@@ -774,6 +774,27 @@ def test_scheduler_reset_prefix_cache():
         assert scheduler.waiting[i] == request
 
 
+def test_reset_connector_cache_no_connector_is_no_op_success():
+    """``reset_connector_cache`` must return True when no connector is
+    configured.
+
+    Without this, ``reset_prefix_cache(reset_connector=True)`` returns
+    ``False`` on every engine that doesn't have a KV connector configured —
+    even when the local prefix cache reset succeeded — and any caller that
+    interprets the return value as "did the reset I asked for succeed?"
+    sees a spurious failure.
+    """
+    scheduler = create_scheduler(enable_prefix_caching=True)
+    assert scheduler.connector is None
+
+    # No-connector reset is treated as success.
+    assert scheduler.reset_connector_cache() is True
+
+    # End-to-end: reset_prefix_cache(reset_connector=True) on an idle
+    # scheduler succeeds with or without a connector.
+    assert scheduler.reset_prefix_cache(reset_connector=True) is True
+
+
 # Note - these test cases mirror some of those in test_rejection_sampler.py
 @pytest.mark.parametrize(
     "spec_tokens,output_tokens,expected",

tests/v1/kv_connector/unit/offloading_connector/test_worker.py

@@ -1,14 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections import defaultdict
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 
 import pytest
 import torch
 
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import get_dtype_size
-from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
@@ -85,15 +84,6 @@ def _allocate_and_reshape_kv_caches(
         set_kv_cache_layout(None)
 
 
-def _make_mock_layer(backend_cls: type[AttentionBackend]):
-    """
-    Create a mock AttentionLayerBase whose get_attn_backend returns backend_cls.
-    """
-    layer = MagicMock()
-    layer.get_attn_backend.return_value = backend_cls
-    return layer
-
-
 def _make_worker(kv_cache_config: KVCacheConfig):
     """
     Create an OffloadingConnectorWorker with mocked dependencies.
@@ -119,11 +109,7 @@ def _make_worker(kv_cache_config: KVCacheConfig):
 
 
 @pytest.mark.parametrize("backend", ATTN_BACKENDS)
-@patch(
-    "vllm.distributed.kv_transfer.kv_connector.v1.offloading"
-    ".worker.get_layers_from_vllm_config"
-)
-def test_register_kv_caches(mock_get_layers, backend):
+def test_register_kv_caches(backend):
     """Test register_kv_caches with multiple groups covering all layer types.
 
     Creates one FullAttention group, one MLA group, one Mamba group, and
@@ -287,13 +273,6 @@ def test_register_kv_caches(mock_get_layers, backend):
         device=torch.device("cuda:0"),
     )
 
-    mock_layers: dict[str, MagicMock] = {}
-    for layer_name in attn_layer_names:
-        mock_layers[layer_name] = _make_mock_layer(backend_cls)
-    for layer_name in mla_layer_names:
-        mock_layers[layer_name] = _make_mock_layer(DeepseekV32IndexerBackend)
-    mock_get_layers.return_value = mock_layers
-
     worker, spec = _make_worker(kv_cache_config)
     worker.register_kv_caches(kv_caches)
 
@@ -360,11 +339,7 @@ def test_register_kv_caches(mock_get_layers, backend):
 
 
 @pytest.mark.parametrize("backend", ATTN_BACKENDS)
-@patch(
-    "vllm.distributed.kv_transfer.kv_connector.v1.offloading"
-    ".worker.get_layers_from_vllm_config"
-)
-def test_register_kv_caches_uniform_type(mock_get_layers, backend):
+def test_register_kv_caches_uniform_type(backend):
     """Test register_kv_caches with UniformTypeKVCacheSpecs.
 
     Two attention layers use the same backend but different num_kv_heads,
@@ -441,64 +416,29 @@ def test_register_kv_caches_uniform_type(mock_get_layers, backend):
         device=torch.device("cuda:0"),
     )
 
-    mock_get_layers.return_value = {
-        layer_a: _make_mock_layer(backend_cls),
-        layer_b: _make_mock_layer(backend_cls),
-    }
-
     worker, spec = _make_worker(kv_cache_config)
     worker.register_kv_caches(kv_caches)
 
     canonical = spec.get_handlers.call_args[0][0]
     assert isinstance(canonical, CanonicalKVCaches)
 
-    unbinds = backend_cls.get_name() in ("FLASH_ATTN", "FLEX_ATTENTION")
-    tensors_per_layer = 2 if unbinds else 1
-
     for block_tensor in canonical.tensors:
         assert block_tensor.tensor.dtype == torch.int8
 
     # Single group with refs from both layers
     assert len(canonical.group_data_refs) == 1
     group_refs = canonical.group_data_refs[0]
-    assert len(group_refs) == 2 * tensors_per_layer
+    assert len(group_refs) == 2
 
-    if unbinds:
-        half_a = spec_a.page_size_bytes // 2
-        half_b = spec_b.page_size_bytes // 2
+    assert len(canonical.tensors) == 2
+    assert canonical.tensors[0].page_size_bytes == spec_a.page_size_bytes
+    assert canonical.tensors[1].page_size_bytes == spec_b.page_size_bytes
+    assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
+    assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
 
-        assert len(canonical.tensors) == 4
-        assert canonical.tensors[0].page_size_bytes == half_a
-        assert canonical.tensors[1].page_size_bytes == half_a
-        assert canonical.tensors[2].page_size_bytes == half_b
-        assert canonical.tensors[3].page_size_bytes == half_b
-        assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, half_a)
-        assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, half_a)
-        assert canonical.tensors[2].tensor.shape == (NUM_BLOCKS, half_b)
-        assert canonical.tensors[3].tensor.shape == (NUM_BLOCKS, half_b)
-
-        assert group_refs[0] == CanonicalKVCacheRef(
-            tensor_idx=0, page_size_bytes=half_a
-        )
-        assert group_refs[1] == CanonicalKVCacheRef(
-            tensor_idx=1, page_size_bytes=half_a
-        )
-        assert group_refs[2] == CanonicalKVCacheRef(
-            tensor_idx=2, page_size_bytes=half_b
-        )
-        assert group_refs[3] == CanonicalKVCacheRef(
-            tensor_idx=3, page_size_bytes=half_b
-        )
-    else:
-        assert len(canonical.tensors) == 2
-        assert canonical.tensors[0].page_size_bytes == spec_a.page_size_bytes
-        assert canonical.tensors[1].page_size_bytes == spec_b.page_size_bytes
-        assert canonical.tensors[0].tensor.shape == (NUM_BLOCKS, spec_a.page_size_bytes)
-        assert canonical.tensors[1].tensor.shape == (NUM_BLOCKS, spec_b.page_size_bytes)
-
-        assert group_refs[0] == CanonicalKVCacheRef(
-            tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
-        )
-        assert group_refs[1] == CanonicalKVCacheRef(
-            tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
-        )
+    assert group_refs[0] == CanonicalKVCacheRef(
+        tensor_idx=0, page_size_bytes=spec_a.page_size_bytes
+    )
+    assert group_refs[1] == CanonicalKVCacheRef(
+        tensor_idx=1, page_size_bytes=spec_b.page_size_bytes
+    )

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -3,7 +3,7 @@
 from collections.abc import Iterable, Iterator
 from dataclasses import dataclass
 from typing import Any
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -25,7 +25,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
 )
 from vllm.forward_context import ForwardContext
 from vllm.utils.hashing import sha256
-from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
 from vllm.v1.core.kv_cache_utils import (
     get_request_block_hasher,
     init_none_hash,
@@ -239,37 +238,22 @@ class RequestRunner:
 
         # register worker kv_caches to enable OffloadingWorker creations
         # set_current_vllm_config is needed for get_kv_cache_layout() to work
-        # Mock get_layers_from_vllm_config so that mock layer names
-        # resolve to layers whose get_attn_backend() returns
-        # FlashAttentionBackend.
-        def _mock_get_layers(_vllm_config, _layer_type, layer_names):
-            mock_layer = MagicMock()
-            mock_layer.get_attn_backend.return_value = FlashAttentionBackend
-            return {name: mock_layer for name in layer_names}
-
         kv_caches: dict[str, torch.Tensor] = {}
         for group in kv_cache_groups:
             spec = group.kv_cache_spec
             for layer_name in group.layer_names:
                 # Shape follows FlashAttention layout:
-                # (2, num_blocks, block_size, num_kv_heads, head_size)
+                # Shape: (num_blocks, 2, block_size, num_kv_heads, head_size)
                 kv_caches[layer_name] = torch.empty(
-                    2,
                     num_gpu_blocks,
+                    2,
                     spec.block_size,
                     spec.num_kv_heads,
                     spec.head_size,
                     dtype=spec.dtype,
                 )
 
-        with (
-            set_current_vllm_config(vllm_config),
-            patch(
-                "vllm.distributed.kv_transfer.kv_connector.v1"
-                ".offloading.worker.get_layers_from_vllm_config",
-                side_effect=_mock_get_layers,
-            ),
-        ):
+        with set_current_vllm_config(vllm_config):
             self.worker_connector.register_kv_caches(kv_caches)
 
         # extract connector of scheduler

tests/v1/kv_connector/unit/test_hma_auto_config.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Regression tests for HMA auto-disable with KV transfer connectors."""
+
+import pytest
+
+from vllm.config import DeviceConfig, KVTransferConfig, SchedulerConfig, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
+from vllm.platforms import current_platform
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+pytestmark = pytest.mark.cpu_test
+
+
+@pytest.fixture(autouse=True)
+def mock_hybrid_kv_cache_supported(monkeypatch):
+    monkeypatch.setattr(current_platform, "support_hybrid_kv_cache", lambda: True)
+
+
+@pytest.mark.parametrize(
+    "kv_transfer_config,expect_disabled",
+    [
+        (  # HMA-supporting connector → HMA stays enabled
+            KVTransferConfig(
+                kv_connector="SimpleCPUOffloadConnector",
+                kv_role="kv_both",
+                kv_connector_extra_config={"cpu_bytes_to_use": 1 << 30},
+            ),
+            False,
+        ),
+        (  # Non-HMA connector → HMA is auto-disabled
+            KVTransferConfig(kv_connector="ExampleConnector", kv_role="kv_both"),
+            True,
+        ),
+        (  # MultiConnector: all HMA children → HMA stays enabled
+            KVTransferConfig(
+                kv_connector="MultiConnector",
+                kv_role="kv_both",
+                kv_connector_extra_config={
+                    "connectors": [
+                        {
+                            "kv_connector": "SimpleCPUOffloadConnector",
+                            "kv_role": "kv_both",
+                            "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
+                        },
+                        {
+                            "kv_connector": "OffloadingConnector",
+                            "kv_role": "kv_both",
+                            "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
+                        },
+                    ]
+                },
+            ),
+            False,
+        ),
+        (  # MultiConnector: mixed children → HMA is auto-disabled
+            KVTransferConfig(
+                kv_connector="MultiConnector",
+                kv_role="kv_both",
+                kv_connector_extra_config={
+                    "connectors": [
+                        {
+                            "kv_connector": "SimpleCPUOffloadConnector",
+                            "kv_role": "kv_both",
+                            "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
+                        },
+                        {"kv_connector": "ExampleConnector", "kv_role": "kv_both"},
+                    ]
+                },
+            ),
+            True,
+        ),
+    ],
+    ids=["hma_connector", "non_hma_connector", "multi_all_hma", "multi_mixed"],
+)
+def test_hma_auto_config(kv_transfer_config, expect_disabled):
+    vllm_config = VllmConfig(
+        device_config=DeviceConfig("cpu"),
+        kv_transfer_config=kv_transfer_config,
+    )
+    assert (
+        vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is expect_disabled
+    )
+
+
+def test_explicit_hma_with_non_hma_connector_errors_at_factory():
+    vllm_config = VllmConfig(
+        device_config=DeviceConfig("cpu"),
+        scheduler_config=SchedulerConfig(
+            max_model_len=16,
+            is_encoder_decoder=False,
+            disable_hybrid_kv_cache_manager=False,
+        ),
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="ExampleConnector",
+            kv_role="kv_both",
+        ),
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
+    )
+    with pytest.raises(ValueError, match="does not support HMA but HMA is enabled"):
+        KVConnectorFactory.create_connector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -369,15 +369,31 @@ async def test_kv_producer(monkeypatch):
         with patch.object(
             prefill_worker, "_send_blocks", return_value=0
         ) as mock_send_blocks:
+            # With blocks-first layout, each block is virtually split
+            # into K and V halves, producing non-coalesced transfers.
+            kv_half = block_len // 2
+
+            def expected_split_transfers(src_base, dst_base, src_blocks, dst_blocks):
+                """Build expected (src_ptrs, dst_ptrs, lengths) for
+                virtual-split K/V transfers."""
+                src_ptrs, dst_ptrs, lengths = [], [], []
+                for kv_offset in (0, kv_half):
+                    for sb, db in zip(src_blocks, dst_blocks):
+                        src_ptrs.append(src_base + sb * block_len + kv_offset)
+                        dst_ptrs.append(dst_base + db * block_len + kv_offset)
+                        lengths.append(kv_half)
+                return src_ptrs, dst_ptrs, lengths
+
             # Normal case: 2 blocks to 2 blocks
-            # Worker processes the consumer's request
             await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
-            # Verify transfer parameters are correct
-            src_ptr = 0x1000 + 10 * block_len
-            dst_ptr = 0x2000 + 20 * block_len
-            length = 2 * block_len
+            src, dst, lens = expected_split_transfers(
+                0x1000, 0x2000, [10, 11], [20, 21]
+            )
             mock_send_blocks.assert_called_once_with(
-                "consumer-host:54321", [src_ptr], [dst_ptr], [length]
+                "consumer-host:54321",
+                src,
+                dst,
+                lens,
             )
             mock_socket.send_multipart.assert_called_once()
 
@@ -404,11 +420,12 @@ async def test_kv_producer(monkeypatch):
             # Worker processes the consumer's request
             await prefill_worker.send_kv_to_decode(identity, mock_socket, xfer_meta)
             # Verify transfer parameters are correct: 11 to 20
-            src_ptr = 0x1000 + 11 * block_len
-            dst_ptr = 0x2000 + 20 * block_len
-            length = 1 * block_len
+            src, dst, lens = expected_split_transfers(0x1000, 0x2000, [11], [20])
             mock_send_blocks.assert_called_once_with(
-                "consumer-host:54321", [src_ptr], [dst_ptr], [length]
+                "consumer-host:54321",
+                src,
+                dst,
+                lens,
             )
             mock_socket.send_multipart.assert_called_once()
 
@@ -618,18 +635,14 @@ def test_register_kv_caches():
 
             mock_batch_register.assert_called_once()
             registered_ptrs, registered_lens = mock_batch_register.call_args[0]
-            expected_ptrs = {
-                tensor.data_ptr()
-                for kv_pair in kv_caches.values()
-                for tensor in kv_pair
-            }
+            expected_ptrs = {tensor.data_ptr() for tensor in kv_caches.values()}
             assert set(registered_ptrs) == expected_ptrs
-            assert set(registered_lens) == {tensor1[0].nbytes}
+            assert set(registered_lens) == {tensor1.nbytes}
 
             # Verify block_len_per_layer is set correctly.
             assert len(worker.block_len_per_layer) == len(registered_ptrs)
             for bl in worker.block_len_per_layer:
-                assert bl == tensor1[0].nbytes // tensor1.shape[1]
+                assert bl == tensor1.nbytes // tensor1.shape[0]
 
 
 def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
@@ -791,33 +804,49 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):
                 # Flatten nested per-group block IDs for assertions
                 flat_local = [b for g in local_block_ids for b in g]
                 flat_remote = [b for g in remote_block_ids for b in g]
+                num_blocks = len(flat_local)
 
-                # Heterogeneous TP: blocks cannot be coalesced because
-                # local and remote block_lens differ
-                assert len(src_ptrs) == len(flat_local)
-                assert len(dst_ptrs) == len(flat_local)
-                assert len(lengths) == len(flat_local)
+                # With blocks-first layout, virtual split halves block
+                # lengths and doubles transfer regions (K + V).
+                local_kv_block_len = local_block_len // 2
+                remote_kv_block_len = remote_block_len // 2
 
-                # Compute expected offsets based on TP ratio
+                assert len(src_ptrs) == 2 * num_blocks
+                assert len(dst_ptrs) == 2 * num_blocks
+                assert len(lengths) == 2 * num_blocks
+
+                # Compute expected offsets using kv_block_len
                 if d_tp_size <= P_TP_SIZE:
                     tp_ratio = P_TP_SIZE // d_tp_size
                     expected_src_off = 0
-                    expected_dst_off = (P_TP_RANK % tp_ratio) * local_block_len
-                    expected_xfer_len = local_block_len
+                    expected_dst_off = (P_TP_RANK % tp_ratio) * local_kv_block_len
+                    expected_xfer_len = local_kv_block_len
                 else:
                     ratio_abs = d_tp_size // P_TP_SIZE
-                    expected_src_off = (d_rank % ratio_abs) * remote_block_len
+                    expected_src_off = (d_rank % ratio_abs) * remote_kv_block_len
                     expected_dst_off = 0
-                    expected_xfer_len = remote_block_len
+                    expected_xfer_len = remote_kv_block_len
 
-                for idx, (lblk, rblk) in enumerate(zip(flat_local, flat_remote)):
-                    assert src_ptrs[idx] == (
-                        0x1000 + lblk * local_block_len + expected_src_off
-                    )
-                    assert dst_ptrs[idx] == (
-                        0x2000 + rblk * remote_block_len + expected_dst_off
-                    )
-                    assert lengths[idx] == expected_xfer_len
+                # First num_blocks entries are K region,
+                # next num_blocks are V region.
+                for region_idx in range(2):
+                    local_region_base = 0x1000 + region_idx * local_kv_block_len
+                    remote_region_base = 0x2000 + region_idx * remote_kv_block_len
+                    for blk_idx, (lblk, rblk) in enumerate(
+                        zip(flat_local, flat_remote)
+                    ):
+                        idx = region_idx * num_blocks + blk_idx
+                        assert src_ptrs[idx] == (
+                            local_region_base
+                            + lblk * local_block_len
+                            + expected_src_off
+                        )
+                        assert dst_ptrs[idx] == (
+                            remote_region_base
+                            + rblk * remote_block_len
+                            + expected_dst_off
+                        )
+                        assert lengths[idx] == expected_xfer_len
 
                 # Verify successful response sent back to consumer
                 mock_socket.send_multipart.assert_called_once()

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

@@ -11,6 +11,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
     connector as mooncake_store_connector,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
+    protocol,
+    scheduler,
+    worker,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
     MooncakeStoreConnectorMetadata,
 )
@@ -284,3 +289,328 @@ def test_update_connector_output_and_take_events():
     assert connector._kv_cache_events is kv_events
     assert list(connector.take_events()) == [event]
     assert connector._kv_cache_events is None
+
+
+# ============================================================
+# reset_cache() — RL hard-reset path via typed LookupKey protocol
+# ============================================================
+
+
+def test_reset_cache_scheduler_role_delegates_to_reset_store():
+    """SCHEDULER role reset_cache() routes to scheduler.reset_store()."""
+    vllm_config = _make_vllm_config()
+    kv_cache_config = _make_kv_cache_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ) as mock_scheduler_cls,
+    ):
+        conn = mooncake_store_connector.MooncakeStoreConnector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )
+
+    mock_scheduler_cls.return_value.reset_store.return_value = True
+    assert conn.reset_cache() is True
+    mock_scheduler_cls.return_value.reset_store.assert_called_once_with()
+
+
+def test_reset_cache_scheduler_role_propagates_failure():
+    """SCHEDULER role surfaces False when scheduler.reset_store() fails."""
+    vllm_config = _make_vllm_config()
+    kv_cache_config = _make_kv_cache_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ) as mock_scheduler_cls,
+    ):
+        conn = mooncake_store_connector.MooncakeStoreConnector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )
+
+    mock_scheduler_cls.return_value.reset_store.return_value = False
+    assert conn.reset_cache() is False
+
+
+def test_reset_cache_worker_role_returns_none():
+    """WORKER role reset_cache() is a no-op; reset is driven via ZMQ admin."""
+    vllm_config = _make_vllm_config()
+    kv_cache_config = _make_kv_cache_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreWorker"
+        ),
+    ):
+        conn = mooncake_store_connector.MooncakeStoreConnector(
+            vllm_config, KVConnectorRole.WORKER, kv_cache_config
+        )
+
+    assert conn.reset_cache() is None
+
+
+def test_scheduler_reset_store_returns_client_reset_result():
+    """MooncakeStoreScheduler.reset_store() returns LookupKeyClient.reset()."""
+    vllm_config = _make_vllm_config()
+    kv_cache_config = _make_kv_cache_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "scheduler.LookupKeyClient"
+        ) as mock_client_cls,
+    ):
+        sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)
+
+    mock_client_cls.return_value.reset.return_value = True
+    assert sched.reset_store() is True
+    mock_client_cls.return_value.reset.assert_called_once_with()
+
+
+def test_scheduler_reset_store_handles_rpc_exception():
+    """Exceptions from the ZMQ reset RPC convert to False, not raise."""
+    vllm_config = _make_vllm_config()
+    kv_cache_config = _make_kv_cache_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "scheduler.LookupKeyClient"
+        ) as mock_client_cls,
+    ):
+        sched = scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)
+
+    mock_client_cls.return_value.reset.side_effect = RuntimeError("rpc timed out")
+    assert sched.reset_store() is False
+
+
+def test_lookup_key_client_lookup_prepends_typed_tag():
+    """LookupKeyClient.lookup() puts LOOKUP_MSG tag at frame 0."""
+    vllm_config = _make_vllm_config()
+
+    with patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+        "worker.make_zmq_socket"
+    ) as mock_make_socket:
+        client = worker.LookupKeyClient(vllm_config)
+
+    fake_socket = mock_make_socket.return_value
+    fake_socket.recv.return_value = (5).to_bytes(4, "big")
+
+    assert client.lookup(token_len=128, block_hashes=[]) == 5
+
+    sent_frames = fake_socket.send_multipart.call_args[0][0]
+    assert sent_frames[0] == protocol.LOOKUP_MSG
+    assert int.from_bytes(sent_frames[1], "big") == 128
+
+
+def test_lookup_key_client_reset_uses_typed_protocol():
+    """LookupKeyClient.reset() sends RESET_MSG and parses RESP_OK / RESP_ERR."""
+    vllm_config = _make_vllm_config()
+
+    with patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+        "worker.make_zmq_socket"
+    ) as mock_make_socket:
+        client = worker.LookupKeyClient(vllm_config)
+
+    fake_socket = mock_make_socket.return_value
+
+    # ACK path: server returns RESP_OK -> client returns True.
+    fake_socket.recv.return_value = protocol.RESP_OK
+    assert client.reset() is True
+    assert fake_socket.send.call_args[0][0] == protocol.RESET_MSG
+
+    # NACK path: server returns RESP_ERR -> client returns False.
+    fake_socket.recv.return_value = protocol.RESP_ERR
+    assert client.reset() is False
+
+
+def test_protocol_tags_are_distinct_and_non_empty():
+    """Protocol tags must be unique and non-empty to avoid collision."""
+    tags = {protocol.LOOKUP_MSG, protocol.RESET_MSG}
+    assert len(tags) == 2
+    for tag in tags:
+        assert isinstance(tag, bytes)
+        assert len(tag) > 0
+    assert protocol.RESP_OK != protocol.RESP_ERR
+
+
+def test_scheduler_reset_connector_cache_invokes_connector_reset():
+    """Cascade test: Scheduler.reset_prefix_cache(reset_connector=True)
+    cascades into MooncakeStoreConnector.reset_cache without dragging in
+    the heavy KVCacheManager fixtures.
+    """
+    vllm_config = _make_vllm_config()
+    kv_cache_config = _make_kv_cache_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ) as mock_scheduler_cls,
+    ):
+        conn = mooncake_store_connector.MooncakeStoreConnector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )
+
+    mock_scheduler_cls.return_value.reset_store.return_value = True
+
+    class _StubScheduler:
+        def __init__(self, c):
+            self.connector = c
+
+        def reset_connector_cache(self):
+            return self.connector.reset_cache() is not False
+
+    sched = _StubScheduler(conn)
+    assert sched.reset_connector_cache() is True
+    mock_scheduler_cls.return_value.reset_store.assert_called_once_with()
+
+    mock_scheduler_cls.return_value.reset_store.reset_mock()
+    mock_scheduler_cls.return_value.reset_store.return_value = False
+    assert sched.reset_connector_cache() is False
+
+
+def test_reset_cache_scheduler_role_clears_local_state():
+    """SCHEDULER reset_cache() must clear scheduler-side state that points
+    at master keys we're about to wipe -- pending load_specs and
+    accumulated _kv_cache_events both reference keys whose blobs are
+    about to be remove_all'd, so reading them after reset would surface
+    stale references to wiped keys.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+        LoadSpec,
+    )
+
+    vllm_config = _make_vllm_config()
+    kv_cache_config = _make_kv_cache_config()
+
+    with (
+        set_current_vllm_config(vllm_config),
+        patch(
+            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+            "connector.MooncakeStoreScheduler"
+        ) as mock_scheduler_cls,
+    ):
+        conn = mooncake_store_connector.MooncakeStoreConnector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )
+
+    # Seed both sentinel pieces of stale-reference state.
+    sched_inst = mock_scheduler_cls.return_value
+    sched_inst.load_specs = {
+        "req-A": LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=128, can_load=True)
+    }
+    conn._kv_cache_events = mooncake_store_connector.MooncakeStoreKVEvents(
+        num_workers=1
+    )
+    sched_inst.reset_store.return_value = True
+
+    assert conn.reset_cache() is True
+
+    # Both stale references must be cleared by the time reset_store is
+    # invoked downstream (load_specs flushed dict, events nulled).
+    assert sched_inst.load_specs == {}
+    assert conn._kv_cache_events is None
+
+
+def test_lookup_key_server_reset_drains_send_queue_before_remove_all():
+    """LookupKeyServer RESET handler must drain the send thread's
+    request_queue BEFORE calling store.remove_all -- otherwise stale
+    puts that were already in flight when the caller paused generation
+    can land on the master AFTER remove_all and silently repopulate it
+    with KV hashed against the previous-policy weights.
+    """
+    # Exercise the handler logic directly with mocks for the send thread
+    # and store. We assert (a) join() is called, (b) remove_all is called,
+    # and (c) join() comes BEFORE remove_all in the call order. The full
+    # LookupKeyServer is heavy (binds a real ZMQ REP socket), so we drive
+    # just the dispatch branch here via a stub equivalent.
+    from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
+        protocol,
+    )
+
+    call_order: list[str] = []
+
+    fake_send_queue = MagicMock()
+    fake_send_queue.join.side_effect = lambda: call_order.append("join")
+
+    fake_store = MagicMock()
+    fake_store.remove_all.side_effect = lambda force: call_order.append(
+        f"remove_all(force={force})"
+    )
+
+    fake_send_thread = MagicMock()
+    fake_send_thread.request_queue = fake_send_queue
+
+    fake_store_worker = MagicMock()
+    fake_store_worker.kv_send_thread = fake_send_thread
+    fake_store_worker.store = fake_store
+
+    fake_socket = MagicMock()
+    sent: list[bytes] = []
+    fake_socket.send.side_effect = lambda frame: sent.append(frame)
+
+    # Mirror the body of LookupKeyServer.process_request RESET_MSG branch.
+    # Keeping this inline (instead of importing the closure) keeps the
+    # test independent of the live thread lifecycle.
+    msg_type = protocol.RESET_MSG
+    if msg_type == protocol.RESET_MSG:
+        try:
+            if fake_store_worker.kv_send_thread is not None:
+                fake_store_worker.kv_send_thread.request_queue.join()
+            fake_store_worker.store.remove_all(force=True)
+            fake_socket.send(protocol.RESP_OK)
+        except Exception:
+            fake_socket.send(protocol.RESP_ERR)
+
+    # Drain must happen before remove_all.
+    assert call_order == ["join", "remove_all(force=True)"]
+    # Worker reported success.
+    assert sent == [protocol.RESP_OK]
+
+
+def test_lookup_key_server_reset_skips_drain_when_no_send_thread():
+    """When the worker has no send thread (e.g. consumer-only role
+    configurations), the RESET handler must still call remove_all
+    instead of dereferencing a None send thread.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
+        protocol,
+    )
+
+    call_order: list[str] = []
+    fake_store = MagicMock()
+    fake_store.remove_all.side_effect = lambda force: call_order.append("remove_all")
+
+    fake_store_worker = MagicMock()
+    fake_store_worker.kv_send_thread = None
+    fake_store_worker.store = fake_store
+
+    fake_socket = MagicMock()
+    sent: list[bytes] = []
+    fake_socket.send.side_effect = lambda frame: sent.append(frame)
+
+    msg_type = protocol.RESET_MSG
+    if msg_type == protocol.RESET_MSG:
+        try:
+            if fake_store_worker.kv_send_thread is not None:
+                fake_store_worker.kv_send_thread.request_queue.join()
+            fake_store_worker.store.remove_all(force=True)
+            fake_socket.send(protocol.RESP_OK)
+        except Exception:
+            fake_socket.send(protocol.RESP_ERR)
+
+    assert call_order == ["remove_all"]
+    assert sent == [protocol.RESP_OK]

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import importlib.util
-import os
 import subprocess
 import uuid
 from unittest.mock import MagicMock, patch
@@ -202,6 +201,7 @@ def create_vllm_config(
     enable_chunked_prefill: bool = True,
     enable_permute_local_kv: bool = False,
     role="kv_consumer",
+    read_mode: bool = False,
 ) -> VllmConfig:
     """Initialize VllmConfig for testing."""
     scheduler_config = SchedulerConfig(
@@ -228,6 +228,7 @@ def create_vllm_config(
         kv_connector="MoRIIOConnector",
         kv_role=role,
         enable_permute_local_kv=enable_permute_local_kv,
+        kv_connector_extra_config={"read_mode": read_mode},
     )
     return VllmConfig(
         scheduler_config=scheduler_config,
@@ -238,15 +239,6 @@ def create_vllm_config(
     )
 
 
-@pytest.fixture
-def moriio_read_mode():
-    """Force the connector into read mode via env for tests."""
-    os.environ["VLLM_MORIIO_CONNECTOR_READ_MODE"] = "True"
-    yield
-    # Cleanup after test
-    os.environ.pop("VLLM_MORIIO_CONNECTOR_READ_MODE", None)
-
-
 def test_write_mode_saves_local_block_ids():
     """Write mode records local block ids in MoRIIOConnectorMetadata.reqs_to_save."""
 
@@ -358,11 +350,11 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids():
         assert block_id == block.block_id, f"{block_id} != {block.block_id}"
 
 
-def test_read_mode_loads_remote_block_ids(moriio_read_mode):
+def test_read_mode_loads_remote_block_ids():
     """Read mode loads remote block ids into local cache mapping."""
 
     # Setup Scheduler and Request
-    vllm_config = create_vllm_config(role="kv_consumer")
+    vllm_config = create_vllm_config(role="kv_consumer", read_mode=True)
     scheduler = create_scheduler(vllm_config)
 
     # 2 Full Blocks and 1 Half Block.

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -1000,11 +1000,8 @@ def _make_multi_connector(connector_names: list[str]) -> MultiConnector:
     )
 
 
-def test_multi_connector_hma_opt_in():
+def test_multi_connector_hma_support_detection():
     """
-    MultiConnector currently assumes HMA is opt-in: it needs
-    --no-disable-hybrid-kv-cache-manager to be enabled.
-
     At runtime, _all_support_hma is True only when every sub-connector
     implements SupportsHMA. Test all combinations of HMA / non-HMA
     sub-connectors.

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -1635,6 +1635,7 @@ def test_register_kv_caches(
             num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
         )
         is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
+        virtually_split = is_blocks_first and not connector.prefer_cross_layer_blocks
 
         if connector.prefer_cross_layer_blocks:
             with set_current_vllm_config(vllm_config):
@@ -1665,7 +1666,7 @@ def test_register_kv_caches(
             ]
             expected_num_entries = 1
 
-            expected_blocks_count = num_blocks * (2 if is_blocks_first else 1)
+            expected_blocks_count = num_blocks * (2 if virtually_split else 1)
 
             kv_caches = {"all-layers": cross_layers_kv_cache}
         else:
@@ -1739,7 +1740,7 @@ def test_register_kv_caches(
         else:
             num_blocks = kv_cache_config.num_blocks
 
-        if is_blocks_first:
+        if virtually_split:
             expected_block_len = expected_tensor_size // num_blocks // 2
         else:
             expected_block_len = expected_tensor_size // num_blocks

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -723,8 +723,7 @@ def test_has_mamba_init(
 
     block_size = 16
     vllm_config = create_vllm_config(block_size=block_size)
-    # VllmConfig.__post_init__ auto-disables HMA when kv_transfer_config
-    # is set; override so we can test the scheduler's own derivation.
+    # Explicitly enable HMA so we can test the scheduler's own derivation.
     vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
     kv_cache_config = make_kv_cache_config(
         block_size=block_size,

tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -280,7 +280,7 @@ def test_cpu_offloading(
         kv_events_config=kv_events_config,
         kv_transfer_config=kv_transfer_config,
         **({"attention_config": {"backend": attn_backend}} if attn_backend else {}),
-        # HMA models need explicit opt-in when kv_transfer_config is set
+        # Keep HMA explicitly enabled for HMA model coverage.
         **({"disable_hybrid_kv_cache_manager": False} if uses_hma else {}),
         **({"enable_prefix_caching": True} if force_prefix_caching else {}),
         # ROCm: batch size 1 to reduce variability

tests/v1/logits_processors/test_correctness.py

@@ -20,7 +20,7 @@ from tests.v1.sample.utils import (
 )
 from vllm.config import VllmConfig
 from vllm.platforms import current_platform
-from vllm.sampling_params import SamplingParams
+from vllm.sampling_params import SamplingParams, validate_thinking_token_budget
 from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.v1.sample.logits_processor import (
     BatchUpdate,
@@ -1194,3 +1194,37 @@ def test_thinking_budget_enforced_without_penalties():
         "Budget exceeded: in_end should be True so that apply_to_logits "
         "forces the end token"
     )
+
+
+@pytest.mark.parametrize(
+    ("raw_value", "expected"),
+    [
+        (None, None),
+        (-1, None),
+        (10, 10),
+        (0, 0),
+    ],
+)
+def test_validate_thinking_token_budget(raw_value, expected):
+    assert validate_thinking_token_budget(raw_value) == expected
+
+
+def test_sampling_params_minus_one_normalizes_to_none():
+    params = SamplingParams(thinking_token_budget=-1)
+    assert params.thinking_token_budget is None
+
+
+@pytest.mark.parametrize("invalid_budget", [-2, 0.6, 10.5, True])
+def test_validate_thinking_token_budget_rejects_invalid(invalid_budget):
+    from vllm.exceptions import VLLMValidationError
+
+    with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
+        validate_thinking_token_budget(invalid_budget)
+
+
+@pytest.mark.parametrize("invalid_budget", [-2, 0.6, 10.5])
+def test_thinking_budget_invalid_budget_rejected(invalid_budget):
+    from vllm.exceptions import VLLMValidationError
+
+    with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
+        SamplingParams(thinking_token_budget=invalid_budget)

tests/v1/worker/test_gpu_model_runner.py

@@ -34,7 +34,6 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
 from vllm.v1.kv_cache_interface import (
-    AttentionSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -44,7 +43,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
 from vllm.v1.worker.gpu_input_batch import InputBatch
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-from vllm.v1.worker.utils import AttentionGroup, select_common_block_size
+from vllm.v1.worker.utils import select_common_block_size
 
 BLOCK_SIZE = 16
 NUM_BLOCKS = 10
@@ -1195,33 +1194,6 @@ def test_hybrid_attention_mamba_tensor_shapes():
             assert torch.equal(actual_ssm, expected_ssm)
 
 
-def test_update_hybrid_attention_mamba_layout_with_num_block_2_rewrites_stride():
-    from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
-
-    ambiguous_cache = torch.empty((2, 2, BLOCK_SIZE, 1, 8), dtype=torch.float16)
-    """Ambiguous, because both dims[0=kv_dim] and dims[1=num_blocks] == 2"""
-    hidden_size = ambiguous_cache.shape[2:].numel()
-    assert ambiguous_cache.stride()[:2] == (2 * hidden_size, hidden_size)
-
-    attention_spec = AttentionSpec(
-        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=8, dtype=torch.float16
-    )
-    runner_stub = SimpleNamespace(
-        cache_config=SimpleNamespace(cache_dtype="auto"),
-        _kv_cache_spec_attn_group_iterator=lambda: iter(
-            [AttentionGroup(FlashAttentionBackend, ["attn"], attention_spec, 0)]
-        ),
-    )
-    GPUModelRunner._update_hybrid_attention_mamba_layout(
-        runner_stub, {"attn": ambiguous_cache}, [BLOCK_SIZE]
-    )
-
-    assert ambiguous_cache.stride()[:2] == (hidden_size, 2 * hidden_size), """\
-        We expect _update_hybrid_attention_mamba_layout to re-stride the cache from:
-        (2, num_blocks) -> (num_blocks, 2), even when num_blocks==2, 
-        which was ambiguous before get_kv_cache_block_dim was used"""
-
-
 def test_hybrid_block_table_initialization():
     """Test hybrid block table with different kernel and kvcache_manager block
     sizes."""

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -168,7 +168,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
         slot_mappings_by_layer=None,
         hidden_states=None,
         aux_hidden_states=None,
-        kv_connector_output=None,
+        finished_req_ids=set(),
         num_tokens_across_dp=None,
     )
     runner.postprocess = lambda *args, **kwargs: events.append("postprocess")

tools/pre_commit/generate_attention_backend_docs.py

@@ -508,7 +508,7 @@ def parse_mla_prefill_backends() -> list[dict[str, Any]]:
         metadata = backend_metadata.get(backend_name, {})
         display_name = backend_info.get("name", backend_name)
 
-        # Add marker for default Blackwell backend
+        # Add marker for the highest-priority automatic backend.
         marker = ""
         if backend_name == priority_order[0] and priorities.get("blackwell"):
             marker = "‡"
@@ -1595,8 +1595,9 @@ def generate_mla_section(
     lines.extend(
         [
             "",
-            "> **‡** 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",
             "",

vllm/_aiter_ops.py

@@ -144,6 +144,7 @@ def _rocm_aiter_fused_moe_impl(
     intermediate_pad: int = 0,
     bias1: torch.Tensor | None = None,
     bias2: torch.Tensor | None = None,
+    moe_sorting_dispatch_policy: int = 0,
 ) -> torch.Tensor:
     from aiter import ActivationType, QuantType
     from aiter.fused_moe import fused_moe
@@ -171,6 +172,7 @@ def _rocm_aiter_fused_moe_impl(
         intermediate_pad=intermediate_pad,
         bias1=bias1,
         bias2=bias2,
+        moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
     )
 
 
@@ -194,6 +196,7 @@ def _rocm_aiter_fused_moe_fake(
     intermediate_pad: int = 0,
     bias1: torch.Tensor | None = None,
     bias2: torch.Tensor | None = None,
+    moe_sorting_dispatch_policy: int = 0,
 ) -> torch.Tensor:
     if output_dtype is not None:
         return torch.empty_like(hidden_states, dtype=output_dtype)
@@ -1282,6 +1285,18 @@ class rocm_aiter_ops:
         - Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale
     """
 
+    _MOE_DISPATCH_POLICY: int | None = None
+
+    @classmethod
+    @if_aiter_supported
+    def get_moe_dispatch_policy(cls) -> int:
+        """Cached MoE sorting dispatch policy."""
+        if cls._MOE_DISPATCH_POLICY is None:
+            import vllm.envs as envs
+
+            cls._MOE_DISPATCH_POLICY = envs.VLLM_ROCM_AITER_MOE_DISPATCH_POLICY
+        return cls._MOE_DISPATCH_POLICY
+
     # Check if the env variable is set
     _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
     _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
@@ -1890,6 +1905,7 @@ class rocm_aiter_ops:
         intermediate_pad: int = 0,
         bias1: torch.Tensor | None = None,
         bias2: torch.Tensor | None = None,
+        moe_sorting_dispatch_policy: int = 0,
     ) -> torch.Tensor:
         return torch.ops.vllm.rocm_aiter_fused_moe(
             hidden_states,
@@ -1911,6 +1927,7 @@ class rocm_aiter_ops:
             intermediate_pad,
             bias1,
             bias2,
+            moe_sorting_dispatch_policy,
         )
 
     @staticmethod

vllm/benchmarks/datasets/datasets.py

@@ -2323,15 +2323,152 @@ class CustomImageDataset(CustomDataset):
         "prompt": "Which country has the most pokemons based on the given graphs?",
         "image_files": ["path/to/image.png"],
     }
+    {
+        "content": [
+            {"type": "text", "text": "Compare these images: "},
+            {"type": "image", "image": "path/to/image1.png"},
+            {"type": "text", "text": " and "},
+            {"type": "image_url", "image_url": {"url": "path/to/image2.png"}},
+        ],
+    }
     ```
 
-    NOTE: Only the first image file in "image_files" is used for each sample request.
-
     This is used to benchmark multimodal LLMs on arbitrary datasets.
     """
 
     IS_MULTIMODAL = True
 
+    def load_data(self) -> None:
+        if self.dataset_path is None:
+            raise ValueError("dataset_path must be provided for loading data.")
+
+        self.data: list[dict] = []
+
+        if not self.dataset_path.endswith(".jsonl"):
+            raise NotImplementedError(
+                "Only JSONL format is supported for CustomImageDataset."
+            )
+
+        with open(self.dataset_path, encoding="utf-8") as f:
+            for line_number, line in enumerate(f, start=1):
+                line = line.strip()
+                if not line:
+                    continue
+
+                try:
+                    item = json.loads(line)
+                except json.JSONDecodeError as e:
+                    raise ValueError(
+                        f"Invalid JSON in custom image dataset line {line_number}: {e}"
+                    ) from e
+
+                if not isinstance(item, dict):
+                    raise ValueError(
+                        "Each custom image dataset line must contain a JSON object. "
+                        f"Found {type(item)} on line {line_number}."
+                    )
+
+                has_legacy_fields = "prompt" in item and "image_files" in item
+                has_interleaved_content = "content" in item
+                if not has_legacy_fields and not has_interleaved_content:
+                    raise ValueError(
+                        "Each custom image dataset line must contain either "
+                        "'prompt' and 'image_files' fields, or a 'content' field. "
+                        f"Invalid line: {line_number}."
+                    )
+
+                self.data.append(item)
+
+        random.seed(self.random_seed)
+        if not getattr(self, "disable_shuffle", False):
+            random.shuffle(self.data)
+
+    @staticmethod
+    def _validate_content_parts(content: Any) -> list[dict[str, Any]]:
+        if not isinstance(content, list):
+            raise ValueError(
+                "'content' must be a list of text and image content dictionaries."
+            )
+
+        if not content:
+            raise ValueError("'content' must contain at least one item.")
+
+        parts: list[dict[str, Any]] = []
+        for part in content:
+            if not isinstance(part, dict):
+                raise ValueError(
+                    f"Each item in 'content' must be a dictionary. Found {type(part)}."
+                )
+            parts.append(part)
+
+        return parts
+
+    @classmethod
+    def _process_content_part(cls, part: dict[str, Any]) -> dict[str, Any]:
+        content_type = part.get("type")
+        if content_type == "text":
+            text = part.get("text")
+            if not isinstance(text, str):
+                raise ValueError("Text content parts must contain a string 'text'.")
+            return {"type": "text", "text": text}
+
+        if content_type == "image":
+            if "image" not in part:
+                raise ValueError("Image content parts must contain an 'image' field.")
+            return dict(process_image(part["image"]))
+
+        if content_type == "image_url":
+            image_url = part.get("image_url")
+            if isinstance(image_url, str):
+                return dict(process_image(image_url))
+
+            if isinstance(image_url, dict):
+                url = image_url.get("url")
+                if not isinstance(url, str):
+                    raise ValueError(
+                        "Image URL content parts must contain a string 'image_url.url'."
+                    )
+
+                processed_part = dict(process_image(url))
+                processed_image_url = dict(processed_part["image_url"])
+                processed_image_url.update(
+                    {key: value for key, value in image_url.items() if key != "url"}
+                )
+                processed_part["image_url"] = processed_image_url
+                return processed_part
+
+            raise ValueError(
+                "Image URL content parts must contain an 'image_url' string "
+                "or dictionary."
+            )
+
+        raise ValueError(
+            "Content parts must have type 'text', 'image', or 'image_url'. "
+            f"Found: {content_type!r}."
+        )
+
+    @classmethod
+    def _process_interleaved_content(cls, content: Any) -> list[dict[str, Any]]:
+        return [
+            cls._process_content_part(part)
+            for part in cls._validate_content_parts(content)
+        ]
+
+    @staticmethod
+    def _get_text_from_content(content: list[dict[str, Any]]) -> str:
+        return "".join(part["text"] for part in content if part.get("type") == "text")
+
+    @staticmethod
+    def _process_image_files(images: Any) -> dict[str, Any] | list[dict[str, Any]]:
+        if not isinstance(images, list) or not images:
+            raise ValueError("'image_files' must be a non-empty list.")
+
+        mm_content = [dict(process_image(image)) for image in images]
+        if len(mm_content) == 1:
+            return mm_content[0]
+
+        return mm_content
+
     def sample(
         self,
         tokenizer: TokenizerLike,
@@ -2356,17 +2493,33 @@ class CustomImageDataset(CustomDataset):
         for i, item in enumerate(self.data):
             if len(sampled_requests) >= num_requests:
                 break
+
+            if "content" in item:
+                content = self._process_interleaved_content(item["content"])
+                text_prompt = self._get_text_from_content(content)
+                prompt_len = len(tokenizer(text_prompt).input_ids)
+                prompt = (
+                    [{"role": "user", "content": content}]
+                    if enable_multimodal_chat
+                    else content
+                )
+                sampled_requests.append(
+                    SampleRequest(
+                        prompt=prompt,
+                        prompt_len=prompt_len,
+                        expected_output_len=output_len,
+                        multi_modal_data=None,
+                        request_id=request_id_prefix + str(i),
+                    )
+                )
+                continue
+
             prompt = item["prompt"]
+            if not isinstance(prompt, str):
+                raise ValueError("'prompt' must be a string.")
 
             prompt_len = len(tokenizer(prompt).input_ids)
-            images = item["image_files"]
-            if len(images) > 1:
-                logger.warning(
-                    "Multiple image files found for sample %d. "
-                    "Only the first image will be used.",
-                    i,
-                )
-            mm_content = process_image(images[0])
+            mm_content = self._process_image_files(item["image_files"])
             if enable_multimodal_chat:
                 # Note: when chat is enabled the request prompt_len is no longer
                 # accurate and we will be using request output to count the

vllm/benchmarks/lib/endpoint_request_func.py

@@ -66,7 +66,7 @@ class StreamedResponseHandler:
 class RequestFuncInput:
     """The input for the request function."""
 
-    prompt: str | list[str]
+    prompt: str | list[str] | list[dict[str, Any]]
     api_url: str
     prompt_len: int
     output_len: int
@@ -268,8 +268,6 @@ def _get_chat_content(
     request_func_input: RequestFuncInput,
     mm_position: Literal["first", "last"] = "last",
 ) -> list[dict[str, Any]]:
-    text_contents = [{"type": "text", "text": request_func_input.prompt}]
-
     mm_contents = []
     if request_func_input.multi_modal_content:
         mm_content = request_func_input.multi_modal_content
@@ -282,12 +280,60 @@ def _get_chat_content(
                 "multi_modal_content must be a dict or list[dict] for openai-chat"
             )
 
+    prompt = request_func_input.prompt
+    if (
+        isinstance(prompt, list)
+        and prompt
+        and all(
+            isinstance(item, dict) and isinstance(item.get("type"), str)
+            for item in prompt
+        )
+    ):
+        if mm_position == "first":
+            return mm_contents + prompt
+
+        return prompt + mm_contents
+
+    text_contents = [{"type": "text", "text": prompt}]
+
     if mm_position == "first":
         return mm_contents + text_contents
 
     return text_contents + mm_contents
 
 
+def _is_chat_messages(prompt: Any) -> bool:
+    return (
+        isinstance(prompt, list)
+        and prompt
+        and all(
+            isinstance(item, dict)
+            and isinstance(item.get("role"), str)
+            and isinstance(item.get("content"), (str, list))
+            for item in prompt
+        )
+    )
+
+
+def _get_chat_messages(
+    request_func_input: RequestFuncInput,
+    mm_position: Literal["first", "last"] = "last",
+) -> list[dict[str, Any]]:
+    prompt = request_func_input.prompt
+    if _is_chat_messages(prompt):
+        return prompt
+
+    return [
+        {
+            "role": "user",
+            "content": _get_chat_content(
+                request_func_input,
+                mm_position=mm_position,
+            ),
+        }
+    ]
+
+
 async def async_request_openai_chat_completions(
     request_func_input: RequestFuncInput,
     session: aiohttp.ClientSession,
@@ -297,15 +343,13 @@ async def async_request_openai_chat_completions(
     api_url = request_func_input.api_url
     _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions")
 
-    content = _get_chat_content(request_func_input, mm_position=mm_position)
+    messages = _get_chat_messages(request_func_input, mm_position=mm_position)
 
     payload = {
         "model": request_func_input.model_name
         if request_func_input.model_name
         else request_func_input.model,
-        "messages": [
-            {"role": "user", "content": content},
-        ],
+        "messages": messages,
         "max_completion_tokens": request_func_input.output_len,
         "stream": True,
         "stream_options": {
@@ -608,15 +652,13 @@ async def async_request_openai_embeddings_chat(
     api_url = request_func_input.api_url
     _validate_api_url(api_url, "OpenAI Embeddings API", "embeddings")
 
-    content = _get_chat_content(request_func_input, mm_position=mm_position)
+    messages = _get_chat_messages(request_func_input, mm_position=mm_position)
 
     payload = {
         "model": request_func_input.model_name
         if request_func_input.model_name
         else request_func_input.model,
-        "messages": [
-            {"role": "user", "content": content},
-        ],
+        "messages": messages,
         # Many embedding models have short context length,
         # this is to avoid dropping some of the requests.
         "truncate_prompt_tokens": -1,

vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py

@@ -1,340 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Fusion pass: replace MiniMax QK allreduce + RMS norm with the Lamport
-fused kernel (minimax_allreduce_rms_qk) for decode-size batches.
-
-Pattern (inlined forward_qk in compiled graph):
-    q, k, v = qkv.split([q_size, kv_size, kv_size], -1)
-    q_fp32 = q.to(float32); k_fp32 = k.to(float32)
-    q_var = q_fp32.pow(2).mean(-1, keepdim=True)
-    k_var = k_fp32.pow(2).mean(-1, keepdim=True)
-    qk_var = cat([q_var, k_var], -1)
-    qk_var = allreduce(qk_var) / tp_world
-    q_var, k_var = qk_var.chunk(2, -1)
-    q_out = (q_fp32 * rsqrt(q_var + eps) * q_weight).to(orig_dtype)
-    k_out = (k_fp32 * rsqrt(k_var + eps) * k_weight).to(orig_dtype)
-    return q_out, k_out, v
-
-Replacement (pure, no in-place on qkv/q/k):
-    q_out, k_out = minimax_qk_norm_fused(qkv, q_weight, k_weight, workspace, ...)
-    v = qkv.split([q_size, kv_size, kv_size], -1)[2]
-    return q_out, k_out, v
-
-is_applicable_for_range: only fires for compile_range.end <= max_decode_tokens
-so that large prefill batches fall through to the original forward_qk (= main).
-"""
-
-import torch
-import torch._inductor.pattern_matcher as pm
-import torch.fx as fx
-from torch._inductor.pattern_matcher import PatternMatcherPass
-
-from vllm.config import VllmConfig
-from vllm.config.utils import Range
-from vllm.distributed import tensor_model_parallel_all_reduce
-from vllm.distributed.parallel_state import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
-from vllm.logger import init_logger
-from vllm.utils.torch_utils import direct_register_custom_op
-
-from ..inductor_pass import enable_fake_mode
-from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
-
-logger = init_logger(__name__)
-
-MAX_TOKEN_NUM = 2048
-
-_MINIMAX_QK_NORM_FUSED_OP = None
-if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
-
-    def _minimax_qk_norm_fused(
-        qkv: torch.Tensor,
-        norm_weight_q: torch.Tensor,
-        norm_weight_k: torch.Tensor,
-        q_size: int,
-        kv_size: int,
-        rank: int,
-        nranks: int,
-        eps: float,
-        max_tokens: int,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        from vllm.distributed.parallel_state import get_tp_group
-        from vllm.model_executor.layers.mamba.lamport_workspace import (
-            get_allreduce_workspace,
-        )
-
-        workspace = get_allreduce_workspace(
-            rank=rank,
-            world_size=nranks,
-            max_tokens=max_tokens,
-            process_group=get_tp_group().cpu_group,
-        )
-        return torch.ops._C.minimax_allreduce_rms_qk(
-            qkv,
-            norm_weight_q,
-            norm_weight_k,
-            workspace,
-            q_size,
-            kv_size,
-            rank,
-            nranks,
-            eps,
-        )
-
-    def _minimax_qk_norm_fused_fake(
-        qkv: torch.Tensor,
-        norm_weight_q: torch.Tensor,
-        norm_weight_k: torch.Tensor,
-        q_size: int,
-        kv_size: int,
-        rank: int,
-        nranks: int,
-        eps: float,
-        max_tokens: int,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        T = qkv.shape[0]
-        return (
-            torch.empty([T, q_size], dtype=qkv.dtype, device=qkv.device),
-            torch.empty([T, kv_size], dtype=qkv.dtype, device=qkv.device),
-        )
-
-    direct_register_custom_op(
-        op_name="minimax_qk_norm_fused",
-        op_func=_minimax_qk_norm_fused,
-        fake_impl=_minimax_qk_norm_fused_fake,
-        mutates_args=[],
-    )
-    _MINIMAX_QK_NORM_FUSED_OP = torch.ops.vllm.minimax_qk_norm_fused.default
-
-
-class MiniMaxQKNormPattern:
-    """
-    Match the forward_qk allreduce+rms pattern and replace with Lamport kernel.
-    """
-
-    def __init__(
-        self,
-        q_size: int,
-        kv_size: int,
-        eps: float,
-        tp_world: int,
-        tp_rank: int,
-        max_tokens: int,
-        dtype: torch.dtype,
-        device: str | None,
-    ) -> None:
-        self.q_size = q_size
-        self.kv_size = kv_size
-        self.eps = eps
-        self.tp_world = tp_world
-        self.tp_rank = tp_rank
-        self.max_tokens = max_tokens
-        self.dtype = dtype
-        self.device = device
-
-    def get_inputs(self) -> list[torch.Tensor]:
-        T = 4
-        qkv = torch.empty(
-            [T, self.q_size + 2 * self.kv_size],
-            device=self.device,
-            dtype=self.dtype,
-        )
-        q_weight = torch.empty([self.q_size], device=self.device, dtype=self.dtype)
-        k_weight = torch.empty([self.kv_size], device=self.device, dtype=self.dtype)
-        return [qkv, q_weight, k_weight]
-
-    def register(self, pm_pass: PatternMatcherPass) -> None:
-        q_size = self.q_size
-        kv_size = self.kv_size
-        eps = self.eps
-        tp_world = self.tp_world
-        max_tokens = self.max_tokens
-        tp_rank = self.tp_rank
-        dtype = self.dtype
-
-        def pattern(
-            qkv: torch.Tensor,
-            q_weight: torch.Tensor,
-            k_weight: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
-            q_fp32 = q.to(torch.float32)
-            k_fp32 = k.to(torch.float32)
-            q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
-            k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
-            qk_var = torch.cat([q_var, k_var], dim=-1)
-            qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
-            q_var, k_var = qk_var.chunk(2, dim=-1)
-            q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
-            k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
-            return q_out, k_out, v
-
-        def replacement(
-            qkv: torch.Tensor,
-            q_weight: torch.Tensor,
-            k_weight: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            assert _MINIMAX_QK_NORM_FUSED_OP is not None
-            q_out, k_out = torch.ops.vllm.minimax_qk_norm_fused(
-                qkv,
-                q_weight,
-                k_weight,
-                q_size,
-                kv_size,
-                tp_rank,
-                tp_world,
-                eps,
-                max_tokens,
-            )
-            _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
-            return q_out, k_out, v
-
-        pm.register_replacement(
-            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
-        )
-
-        # Second pattern: three separate split_with_sizes nodes (one per output),
-        # each with _users=1. This occurs when the QKV projection uses a
-        # functional GEMM kernel (e.g. cutlass_scaled_mm via auto_functionalized),
-        # which causes inductor to generate one split per consumer.
-        def pattern_split3(
-            qkv: torch.Tensor,
-            q_weight: torch.Tensor,
-            k_weight: torch.Tensor,
-        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0]
-            k = qkv.split([q_size, kv_size, kv_size], dim=-1)[1]
-            v = qkv.split([q_size, kv_size, kv_size], dim=-1)[2]
-            q_fp32 = q.to(torch.float32)
-            k_fp32 = k.to(torch.float32)
-            q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
-            k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
-            qk_var = torch.cat([q_var, k_var], dim=-1)
-            qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
-            q_var, k_var = qk_var.chunk(2, dim=-1)
-            q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
-            k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
-            return q_out, k_out, v
-
-        pm.register_replacement(
-            pattern_split3, replacement, self.get_inputs(), pm.fwd_only, pm_pass
-        )
-
-
-class MiniMaxQKNormPass(VllmPatternMatcherPass):
-    """
-    Replace forward_qk allreduce+norm with the Lamport fused kernel.
-    Only applied for decode-size compile ranges (small token counts).
-    """
-
-    def __init__(self, config: VllmConfig) -> None:
-        super().__init__(config)
-        self.disabled = True
-
-        if _MINIMAX_QK_NORM_FUSED_OP is None:
-            logger.warning_once(
-                "minimax_allreduce_rms_qk op not found, MiniMaxQKNormPass disabled."
-            )
-            return
-
-        tp_world = get_tensor_model_parallel_world_size()
-        if tp_world <= 1:
-            logger.warning_once("MiniMaxQKNormPass disabled: tp_size <= 1.")
-            return
-
-        if config.model_config is None:
-            logger.warning_once("MiniMaxQKNormPass disabled: no model_config.")
-            return
-
-        hf_cfg = config.model_config.hf_config
-
-        model_name = getattr(hf_cfg, "architectures", "")[0]
-        if model_name != "MiniMaxM2ForCausalLM":
-            return
-
-        num_attention_heads = getattr(hf_cfg, "num_attention_heads", 0)
-        num_key_value_heads = getattr(hf_cfg, "num_key_value_heads", 0)
-        hidden_size = getattr(hf_cfg, "hidden_size", 0)
-        head_dim = getattr(hf_cfg, "head_dim", 0)
-        eps: float = getattr(hf_cfg, "rms_norm_eps", 1e-6)
-
-        if (
-            num_attention_heads != 48
-            or num_key_value_heads != 8
-            or hidden_size != 3072
-            or head_dim != 128
-        ):
-            logger.warning_once(
-                "MiniMaxQKNormPass disabled: cannot infer model info from hf_config."
-            )
-            return
-
-        num_heads_per_rank = num_attention_heads // tp_world
-        num_kv_heads_per_rank = max(1, num_key_value_heads // tp_world)
-        q_size = num_heads_per_rank * head_dim
-        kv_size = num_kv_heads_per_rank * head_dim
-
-        self.max_token_num = min(
-            MAX_TOKEN_NUM, config.scheduler_config.max_num_batched_tokens
-        )
-
-        tp_rank = get_tensor_model_parallel_rank()
-        # Allocate Lamport workspace first.
-        from vllm.distributed.parallel_state import get_tp_group
-        from vllm.model_executor.layers.mamba.lamport_workspace import (
-            get_allreduce_workspace,
-        )
-
-        get_allreduce_workspace(
-            rank=tp_rank,
-            world_size=tp_world,
-            max_tokens=self.max_token_num,
-            process_group=get_tp_group().cpu_group,
-        )
-
-        self.patterns: PatternMatcherPass = PatternMatcherPass(
-            pass_name="minimax_qk_norm_pass"
-        )
-        self._register_patterns(q_size, kv_size, eps, tp_world, tp_rank)
-        self.dump_patterns(config, self.patterns)
-        self.disabled = False
-
-    @enable_fake_mode
-    def _register_patterns(
-        self,
-        q_size: int,
-        kv_size: int,
-        eps: float,
-        tp_world: int,
-        tp_rank: int,
-    ) -> None:
-        MiniMaxQKNormPattern(
-            q_size=q_size,
-            kv_size=kv_size,
-            eps=eps,
-            tp_world=tp_world,
-            tp_rank=tp_rank,
-            max_tokens=self.max_token_num,
-            dtype=self.model_dtype,
-            device=self.device,
-        ).register(self.patterns)
-
-    def is_applicable_for_range(self, compile_range: Range) -> bool:
-        if self.disabled:
-            return False
-
-        return bool(compile_range.end <= self.max_token_num)
-
-    @VllmInductorPass.time_and_log
-    def __call__(self, graph: fx.Graph) -> None:
-        if self.disabled:
-            return
-        self.matched_count = self.patterns.apply(graph)
-        logger.debug("MiniMaxQKNormPass replaced %s patterns", self.matched_count)
-
-    def uuid(self) -> str:
-        return VllmInductorPass.hash_source(self, MiniMaxQKNormPattern)

vllm/compilation/passes/pass_manager.py

@@ -44,7 +44,6 @@ if current_platform.is_cuda_alike():
 if current_platform.is_cuda():
     from .fusion.allreduce_rms_fusion import AllReduceFusionPass
     from .fusion.collective_fusion import AsyncTPPass
-    from .fusion.minimax_qk_norm_fusion import MiniMaxQKNormPass
 
 from .inductor_pass import (
     CustomGraphPass,
@@ -154,9 +153,6 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
                 else:
                     self.passes += [AllReduceFusionPass(config)]
 
-            if self.pass_config.fuse_minimax_qk_norm:
-                self.passes += [MiniMaxQKNormPass(config)]
-
             if self.pass_config.fuse_norm_quant:
                 if rocm_aiter_ops.is_enabled():
                     self.passes += [

vllm/config/compilation.py

@@ -135,7 +135,9 @@ class PassConfig:
     fuse_allreduce_rms: bool = None  # type: ignore[assignment]
     """Enable flashinfer allreduce fusion."""
     fuse_minimax_qk_norm: bool = None  # type: ignore[assignment]
-    """Enable fused allreduce+RMSNorm for MiniMax QK norm."""
+    """Deprecated. The MiniMax QK norm fusion is now applied automatically at
+    runtime (see `MiniMaxText01RMSNormTP.forward_qkv`). This flag is kept for
+    backward compatibility and has no effect; it will be removed in v0.23."""
     enable_qk_norm_rope_fusion: bool = None  # type: ignore[assignment]
     """Enable fused Q/K RMSNorm + RoPE pass."""
     fuse_rope_kvcache_cat_mla: bool = None  # type: ignore[assignment]
@@ -294,6 +296,13 @@ class PassConfig:
                 "current platform is not CUDA or ROCm. The fusion will be disabled."
             )
             self.fuse_rope_kvcache_cat_mla = False
+        if self.fuse_minimax_qk_norm is not None:
+            logger.warning_once(
+                "`fuse_minimax_qk_norm` is deprecated and has no effect; "
+                "the MiniMax QK norm fusion is now applied automatically at "
+                "runtime when its conditions are met. This flag will be "
+                "removed in v0.23."
+            )
 
     def log_enabled_passes(self) -> None:
         """