m3 engine parser (reasoning-only) stage

Signed-off-by: Ting Sun <suntcrick@gmail.com>

.buildkite/test_areas/lm_eval.yaml

@@ -109,6 +109,7 @@ steps:
       - image-build-amd
       commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
       - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
 
 - label: MoE Refactor Integration Test (H100 - TEMPORARY)

cmake/external_projects/fmha_sm100.cmake

@@ -17,7 +17,7 @@ else()
   FetchContent_Declare(
     fmha_sm100
     GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
+    GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND ""
@@ -36,38 +36,13 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")
 
 install(FILES
   "${FMHA_SM100_PY_ROOT}/__init__.py"
-  "${FMHA_SM100_PY_ROOT}/api.py"
-  "${FMHA_SM100_PY_ROOT}/bench_utils.py"
-  "${FMHA_SM100_PY_ROOT}/jit.py"
   "${FMHA_SM100_PY_ROOT}/sparse.py"
-  "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
   DESTINATION vllm/third_party/fmha_sm100
   COMPONENT fmha_sm100)
 
-install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
-  DESTINATION vllm/third_party/fmha_sm100/csrc
-  COMPONENT fmha_sm100
-  PATTERN "__pycache__" EXCLUDE
-  PATTERN "*.pyc" EXCLUDE
-  PATTERN ".git*" EXCLUDE)
-
 install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
   DESTINATION vllm/third_party/fmha_sm100/cute
   COMPONENT fmha_sm100
   PATTERN "__pycache__" EXCLUDE
   PATTERN "*.pyc" EXCLUDE
   PATTERN ".git*" EXCLUDE)
-
-install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/include/"
-  DESTINATION vllm/third_party/fmha_sm100/cutlass/include
-  COMPONENT fmha_sm100
-  PATTERN "__pycache__" EXCLUDE
-  PATTERN "*.pyc" EXCLUDE
-  PATTERN ".git*" EXCLUDE)
-
-install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/tools/util/include/"
-  DESTINATION vllm/third_party/fmha_sm100/cutlass/tools/util/include
-  COMPONENT fmha_sm100
-  PATTERN "__pycache__" EXCLUDE
-  PATTERN "*.pyc" EXCLUDE
-  PATTERN ".git*" EXCLUDE)

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -67,13 +67,6 @@
   #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif
 
-// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs.
-#ifndef USE_ROCM
-  #include <cuda_fp8.h>
-#else
-  #include <hip/hip_fp8.h>
-#endif
-
 #ifndef FINAL_MASK
   #ifdef USE_ROCM
     #define FINAL_MASK 0xffffffffffffffffULL
@@ -82,19 +75,6 @@
   #endif
 #endif
 
-#ifdef USE_ROCM
-// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4
-// fused kernel).
-__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
-  #if defined(HIP_FP8_TYPE_OCP)
-  __hip_fp8_e4m3 fp8_val(val);
-  #else
-  __hip_fp8_e4m3_fnuz fp8_val(val);
-  #endif
-  return reinterpret_cast<uint8_t&>(fp8_val);
-}
-#endif
-
 namespace vllm {
 namespace minimax_m3_fused_ops {
 
@@ -213,8 +193,6 @@ __device__ __forceinline__ void storeElems(
   *reinterpret_cast<uint2*>(dst) = v;
 }
 
-// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache
-// dtypes use the scaled-convert path with identity scale.
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __device__ __forceinline__ void storeCacheElems(
     cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
@@ -230,32 +208,6 @@ __device__ __forceinline__ void storeCacheElems(
   }
 }
 
-// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast,
-// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale
-// (RMSNorm outputs are O(1) and the score path only needs relative block
-// ordering).
-__device__ __forceinline__ void storeElemsFp8(
-    uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
-  constexpr float kFp8Max = 448.0f;
-#ifndef USE_ROCM
-  __nv_fp8x2_storage_t out2[kElemsPerLane / 2];
-  #pragma unroll
-  for (int i = 0; i < kElemsPerLane / 2; i++) {
-    float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]);
-    vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max);
-    vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max);
-    out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3);
-  }
-  *reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
-#else
-  #pragma unroll
-  for (int i = 0; i < kElemsPerLane; i++) {
-    float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max);
-    dst[i] = rocm_cvt_float_to_fp8_e4m3(vv);
-  }
-#endif
-}
-
 // ────────────────────────────────────────────────────────────────────────────
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
@@ -272,14 +224,12 @@ __device__ __forceinline__ void storeElemsFp8(
 //     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
 //     IQ: niq  only if kIsSparse        (norm+RoPE)
 //     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
-// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
-// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
-          typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
+          bool kIsSparse, bool kInsertKV>
 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
-    scalar_t* __restrict__ q_out,         // [N, nq*128] contiguous, or nullptr
-    out_idx_t* __restrict__ index_q_out,  // [N, niq*128]; scalar_t or e4m3 byte
+    scalar_t* __restrict__ q_out,        // [N, nq*128] contiguous, or nullptr
+    scalar_t* __restrict__ index_q_out,  // [N, niq*128] contiguous, or nullptr
     scalar_t const* __restrict__ q_norm_w,
     scalar_t const* __restrict__ k_norm_w,
     scalar_t const* __restrict__ iq_norm_w,
@@ -288,8 +238,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     int64_t const* __restrict__ positions,       // [N] i64
     int64_t const* __restrict__ slot_mapping,    // main K/V slots or nullptr
     int64_t const* __restrict__ index_slot_mapping,  // index K slots/nullptr
-    cache_t* __restrict__ kv_cache,       // [nb,2,bs,nkv,128] or nullptr
-    out_idx_t* __restrict__ index_cache,  // [nb*bs, 128]; scalar_t or e4m3 byte
+    cache_t* __restrict__ kv_cache,      // [nb,2,bs,nkv,128] or nullptr
+    scalar_t* __restrict__ index_cache,  // [nb*bs, 128] or nullptr
     float const eps, int const rotary_dim, int const num_tokens, int const nq,
     int const nkv, int const niq, int const block_size,
     // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
@@ -384,12 +334,9 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
       store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
                   slot * kHeadDim;
     } else if (isIQ && index_q_out != nullptr) {
-      // bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
-      if constexpr (!kFp8Idx) {
-        store_ptr = index_q_out +
-                    static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
-                    (slot - iq_begin) * kHeadDim;
-      }
+      store_ptr = index_q_out +
+                  static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                  (slot - iq_begin) * kHeadDim;
     }
 
     // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
@@ -409,19 +356,7 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
       scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
       normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
                             cos_ptr, /*apply_norm=*/norm_w != nullptr);
-      if constexpr (kFp8Idx) {
-        // index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
-        if (isIQ && index_q_out != nullptr) {
-          storeElemsFp8(index_q_out +
-                            static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
-                            (slot - iq_begin) * kHeadDim + dim_base,
-                        elems);
-        } else {
-          storeElems<scalar_t>(store_ptr + dim_base, elems);
-        }
-      } else {
-        storeElems<scalar_t>(store_ptr + dim_base, elems);
-      }
+      storeElems<scalar_t>(store_ptr + dim_base, elems);
     }
 
     // ── Cache inserts (sparse serving only). ───────────────────────────────
@@ -432,11 +367,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
                              : (isIK ? index_slot_mapping[tokenIdx] : -1);
       if (sm >= 0) {  // skip padded / unscheduled tokens
         if (isIK) {
-          if constexpr (kFp8Idx) {
-            storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems);
-          } else {
-            storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
-          }
+          scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
+          storeElems<scalar_t>(dst, elems);
         } else if (isK || isV) {
           // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
           // Paging is logical (block = sm/block_size, token = sm%block_size);
@@ -466,19 +398,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
-void launchFusedMiniMaxM3(
-    scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w,
-    scalar_t const* k_norm_w, scalar_t const* iq_norm_w,
-    scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache,
-    int64_t const* positions, int64_t const* slot_mapping,
-    int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
-    float const eps, int const rotary_dim, int const num_tokens, int const nq,
-    int const nkv, int const niq, int const block_size,
-    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
-    int64_t const kv_s_head, bool const has_index, bool const insert_kv,
-    bool const fp8_idx, cudaStream_t stream) {
-  // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
-  // void* pointers per instantiation in the LAUNCH macro.
+void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
+                          scalar_t const* q_norm_w, scalar_t const* k_norm_w,
+                          scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
+                          scalar_t const* cos_sin_cache,
+                          int64_t const* positions, int64_t const* slot_mapping,
+                          int64_t const* index_slot_mapping, cache_t* kv_cache,
+                          scalar_t* index_cache, float const eps,
+                          int const rotary_dim, int const num_tokens,
+                          int const nq, int const nkv, int const niq,
+                          int const block_size, int64_t const kv_s_block,
+                          int64_t const kv_s_kv, int64_t const kv_s_token,
+                          int64_t const kv_s_head, bool const has_index,
+                          bool const insert_kv, cudaStream_t stream) {
   // Slot count must match the kernel's compile-time gating.
   int const v_slots = insert_kv ? nkv : 0;
   int const idx_slots = has_index ? niq + 1 : 0;
@@ -508,27 +440,25 @@ void launchFusedMiniMaxM3(
   config.attrs = attrs;
   config.numAttrs = (sm_version >= 90) ? 1 : 0;
 
-  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                                \
-    cudaLaunchKernelEx(                                                        \
-        &config,                                                               \
-        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
-                                              IS_SPARSE, INSERT, FP8>,         \
-        qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
-        iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping,          \
-        index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache),   \
-        eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block,     \
-        kv_s_kv, kv_s_token, kv_s_head)
+  #define LAUNCH(IS_SPARSE, INSERT)                                           \
+    cudaLaunchKernelEx(                                                       \
+        &config,                                                              \
+        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,       \
+                                              IS_SPARSE, INSERT>,             \
+        qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w,    \
+        cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
+        index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size,   \
+        kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
 #else
   // ROCm: standard kernel launch syntax (no PDL/stream serialization).
   // clang-format off
-  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                              \
-    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T,   \
-                                          IS_SPARSE, INSERT, FP8>            \
+  #define LAUNCH(IS_SPARSE, INSERT)                                          \
+    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,          \
+                                          IS_SPARSE, INSERT>                 \
         <<<grid, kBlockSize, 0, stream>>>(                                   \
-            qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w,     \
-            k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions,        \
-            slot_mapping, index_slot_mapping, kv_cache,                      \
-            reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim,          \
+            qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w,          \
+            ik_norm_w, cos_sin_cache, positions, slot_mapping,               \
+            index_slot_mapping, kv_cache, index_cache, eps, rotary_dim,      \
             num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv,       \
             kv_s_token, kv_s_head)
   // clang-format on
@@ -536,22 +466,14 @@ void launchFusedMiniMaxM3(
 
   if (has_index) {
     if (insert_kv) {
-      if (fp8_idx) {
-        LAUNCH(true, true, true, uint8_t);  // sparse serving, fp8 index outputs
-      } else {
-        LAUNCH(true, true, false, scalar_t);  // sparse serving, bf16
-      }
+      LAUNCH(true, true);  // sparse serving
     } else {
-      if (fp8_idx) {
-        LAUNCH(true, false, true, uint8_t);  // sparse profiling, fp8 index_q
-      } else {
-        LAUNCH(true, false, false, scalar_t);  // sparse profiling, bf16
-      }
+      LAUNCH(true, false);  // sparse profiling
     }
   } else {
     // Dense layer: never has an index branch and never inserts here (the
     // generic Attention layer owns the KV insert).
-    LAUNCH(false, false, false, scalar_t);
+    LAUNCH(false, false);
   }
 #undef LAUNCH
 }
@@ -563,9 +485,8 @@ void launchFusedMiniMaxM3(
   vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>(     \
       reinterpret_cast<st*>(qkv.data_ptr()),                                   \
       q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr,  \
-      index_q_out.has_value()                                                  \
-          ? reinterpret_cast<void*>(index_q_out->data_ptr())                   \
-          : nullptr,                                                           \
+      index_q_out.has_value() ? reinterpret_cast<st*>(index_q_out->data_ptr()) \
+                              : nullptr,                                       \
       reinterpret_cast<st const*>(q_norm_weight.data_ptr()),                   \
       reinterpret_cast<st const*>(k_norm_weight.data_ptr()),                   \
       has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
@@ -581,11 +502,11 @@ void launchFusedMiniMaxM3(
                 : nullptr,                                                     \
       insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr,  \
       (insert_kv && has_index)                                                 \
-          ? reinterpret_cast<void*>(index_cache->data_ptr())                   \
+          ? reinterpret_cast<st*>(index_cache->data_ptr())                     \
           : nullptr,                                                           \
       static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq,   \
       nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
-      kv_s_head, has_index, insert_kv, fp8_idx, stream)
+      kv_s_head, has_index, insert_kv, stream)
 
 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
@@ -691,7 +612,6 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
                  torch::headeronly::ScalarType::Long &&
              index_slot_mapping->numel() == slot_mapping->numel()),
         "index_slot_mapping must be int64 CUDA with slot_mapping length");
-    // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
     if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
       STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
                       "auto kv_cache dtype must match qkv");
@@ -700,13 +620,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
           kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
           "fp8 kv_cache must use uint8 storage");
     }
-    // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
-    STD_TORCH_CHECK(
-        index_cache.has_value() &&
-            (index_cache->scalar_type() == qkv.scalar_type() ||
-             index_cache->scalar_type() ==
-                 torch::headeronly::ScalarType::Float8_e4m3fn),
-        "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
+    STD_TORCH_CHECK(index_cache.has_value() &&
+                        index_cache->scalar_type() == qkv.scalar_type(),
+                    "insert mode requires matching index_cache");
     STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
                     "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
                     "head_dim (stride(4)==1)");
@@ -736,31 +652,14 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
         "index_q_out requires the index branch (num_index_heads > 0)");
     STD_TORCH_CHECK(
         index_q_out->is_cuda() && index_q_out->is_contiguous() &&
-            (index_q_out->scalar_type() == qkv.scalar_type() ||
-             index_q_out->scalar_type() ==
-                 torch::headeronly::ScalarType::Float8_e4m3fn),
-        "index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3");
+            index_q_out->scalar_type() == qkv.scalar_type(),
+        "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
     STD_TORCH_CHECK(index_q_out->numel() ==
                         static_cast<int64_t>(num_tokens) * niq * kHeadDim,
                     "index_q_out must have num_tokens * num_index_heads * 128 "
                     "elements");
   }
 
-  // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while
-  // q/k/v + q_out stay qkv dtype. Both index outputs must agree.
-  auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
-  bool const fp8_idx =
-      (index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
-      (index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
-  if (fp8_idx) {
-    STD_TORCH_CHECK(
-        !index_cache.has_value() || index_cache->scalar_type() == kFp8,
-        "fp8 index path: index_cache must be fp8 e4m3");
-    STD_TORCH_CHECK(
-        !index_q_out.has_value() || index_q_out->scalar_type() == kFp8,
-        "fp8 index path: index_q_out must be fp8 e4m3");
-  }
-
   const torch::stable::accelerator::DeviceGuard device_guard(
       qkv.get_device_index());
   auto stream = get_current_cuda_stream(qkv.get_device_index());

docker/Dockerfile

@@ -792,7 +792,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.13rc2
+ARG FLASHINFER_VERSION=0.6.12
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

docker/versions.json

@@ -68,7 +68,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.13rc2"
+      "default": "0.6.12"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

requirements/cuda.txt

@@ -9,8 +9,8 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.13rc2
-flashinfer-cubin==0.6.13rc2
+flashinfer-python==0.6.12
+flashinfer-cubin==0.6.12
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 nvidia-cudnn-frontend>=1.19.1

setup.py

@@ -1171,15 +1171,7 @@ package_data = {
         "third_party/deep_gemm/include/**/*.h",
         "third_party/deep_gemm/include/**/*.hpp",
         # fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake)
-        "third_party/fmha_sm100/csrc/**/*.cu",
-        "third_party/fmha_sm100/csrc/**/*.h",
-        "third_party/fmha_sm100/csrc/**/*.jinja",
-        "third_party/fmha_sm100/csrc/**/*.cu.jinja",
         "third_party/fmha_sm100/cute/**/*.cu",
-        "third_party/fmha_sm100/cutlass/include/**/*.h",
-        "third_party/fmha_sm100/cutlass/include/**/*.hpp",
-        "third_party/fmha_sm100/cutlass/tools/util/include/**/*.h",
-        "third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp",
     ]
 }
 

tests/entrypoints/anthropic/test_anthropic_messages_conversion.py

@@ -8,6 +8,8 @@ AnthropicServingMessages._convert_anthropic_to_openai_request().
 Also covers extended-thinking edge cases such as ``redacted_thinking``
 blocks echoed back by Anthropic clients, and streaming conversion in
 ``message_stream_converter``.
+
+Also covers cache usage computation in ``_build_anthropic_usage``.
 """
 
 import json
@@ -18,7 +20,11 @@ import pytest
 from vllm.entrypoints.anthropic.protocol import (
     AnthropicMessagesRequest,
 )
-from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
+from vllm.entrypoints.anthropic.serving import (
+    AnthropicServingMessages,
+    _build_anthropic_usage,
+    _get_cached_tokens,
+)
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionResponseStreamChoice,
     ChatCompletionStreamResponse,
@@ -27,6 +33,7 @@ from vllm.entrypoints.openai.engine.protocol import (
     DeltaFunctionCall,
     DeltaMessage,
     DeltaToolCall,
+    PromptTokenUsageInfo,
     UsageInfo,
 )
 
@@ -653,6 +660,108 @@ class TestThinkingBlockConversion:
         assert asst.get("content") == "Hi!"
 
 
+# ======================================================================
+# Cache usage computation
+# ======================================================================
+
+
+class TestGetCachedTokens:
+    """Tests for _get_cached_tokens helper."""
+
+    def test_none_usage(self):
+        assert _get_cached_tokens(None) is None
+
+    def test_no_prompt_tokens_details(self):
+        usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
+        assert _get_cached_tokens(usage) is None
+
+    def test_cached_tokens_present(self):
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+        )
+        assert _get_cached_tokens(usage) == 80
+
+    def test_cached_tokens_zero(self):
+        """Zero cached tokens should return 0, not None."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+        )
+        assert _get_cached_tokens(usage) == 0
+
+    def test_cached_tokens_none_in_details(self):
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
+        )
+        assert _get_cached_tokens(usage) is None
+
+
+class TestBuildAnthropicUsage:
+    """Tests for _build_anthropic_usage helper.
+
+    Anthropic defines: total_input = input_tokens + cache_read + cache_creation
+    vLLM's prompt_tokens is the total.
+    """
+
+    def test_no_cache_info(self):
+        """When cache info is unavailable, return raw prompt_tokens."""
+        result = _build_anthropic_usage(100, 10, None)
+        assert result.input_tokens == 100
+        assert result.output_tokens == 10
+        assert result.cache_read_input_tokens is None
+        assert result.cache_creation_input_tokens is None
+
+    def test_cache_hit(self):
+        """When cache is hit, input_tokens excludes cached tokens."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 20  # 100 - 80
+        assert result.output_tokens == 10
+        assert result.cache_read_input_tokens == 80
+        assert result.cache_creation_input_tokens == 0
+
+    def test_zero_cached_tokens(self):
+        """Zero cached tokens should still set cache_creation to 0."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 100  # 100 - 0
+        assert result.cache_read_input_tokens == 0
+        assert result.cache_creation_input_tokens == 0
+
+    def test_all_tokens_cached(self):
+        """When all tokens are cached, input_tokens should be 0."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 0
+        assert result.cache_read_input_tokens == 100
+        assert result.cache_creation_input_tokens == 0
+
+    def test_no_prompt_tokens_details(self):
+        """UsageInfo without prompt_tokens_details returns no cache info."""
+        usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 100
+        assert result.cache_read_input_tokens is None
+        assert result.cache_creation_input_tokens is None
+
+
 class TestInlineSystemMessageInMessagesArray:
     """Verify that ``role: system`` messages embedded inside the ``messages``
     array are preserved in their original position.
@@ -1098,6 +1207,135 @@ class TestMessageStartIncludesTypeAndRole:
         assert message["role"] == "assistant"
 
 
+class TestStreamingCacheUsageSemantics:
+    """Locks in the documented streaming behavior of cache usage fields.
+
+    vLLM's OpenAI chat completion streaming only attaches
+    ``prompt_tokens_details`` to the terminal usage chunk. The Anthropic layer
+    mirrors that contract: cache fields are omitted on ``message_start`` (key
+    absence signals "unknown") and populated on ``message_delta`` (the final
+    cumulative count). This is intentionally consistent with vLLM's OpenAI
+    behavior, even though Anthropic's upstream API populates cache fields on
+    ``message_start``; closing that gap requires plumbing cache info into the
+    first chunk at the OpenAI layer, which is out of scope here.
+    """
+
+    @pytest.mark.asyncio
+    async def test_streaming_cache_fields_absent_then_populated(self):
+        """First chunk lacks prompt_tokens_details (vLLM contract);
+        message_start omits cache fields. The final chunk carries
+        prompt_tokens_details, so message_delta carries resolved values."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant", content="hi"),
+                usage=UsageInfo(prompt_tokens=100, total_tokens=100),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(
+                    prompt_tokens=100,
+                    completion_tokens=5,
+                    total_tokens=105,
+                    prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+                ),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        # message_start: cache fields unknown → omitted from JSON entirely.
+        start_usage = events[0][1]["message"]["usage"]
+        assert events[0][0] == "message_start"
+        assert start_usage["input_tokens"] == 100
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+
+        # message_delta: authoritative usage with cache fields populated.
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert delta_usage["input_tokens"] == 20  # 100 - 80
+        assert delta_usage["cache_read_input_tokens"] == 80
+        assert delta_usage["cache_creation_input_tokens"] == 0
+
+    @pytest.mark.asyncio
+    async def test_streaming_no_cache_hit(self):
+        """When the final chunk reports cached_tokens=0, message_delta carries
+        cache fields = 0 (cache miss); message_start still omits them."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant"),
+                usage=UsageInfo(prompt_tokens=50, total_tokens=50),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(
+                    prompt_tokens=50,
+                    completion_tokens=5,
+                    total_tokens=55,
+                    prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+                ),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        start_usage = events[0][1]["message"]["usage"]
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert start_usage["input_tokens"] == 50
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+        assert delta_usage["input_tokens"] == 50  # 50 - 0
+        assert delta_usage["cache_read_input_tokens"] == 0
+        assert delta_usage["cache_creation_input_tokens"] == 0
+
+    @pytest.mark.asyncio
+    async def test_streaming_no_prompt_tokens_details_at_all(self):
+        """If --enable-prompt-tokens-details is off, no chunk carries cache
+        info; both message_start and message_delta omit cache fields."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant"),
+                usage=UsageInfo(prompt_tokens=30, total_tokens=30),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        start_usage = events[0][1]["message"]["usage"]
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+        assert "cache_read_input_tokens" not in delta_usage
+        assert "cache_creation_input_tokens" not in delta_usage
+
+
 # ======================================================================
 # Auto-detection of system-first template requirement
 # ======================================================================

tests/entrypoints/serve/render/test_derender.py

@@ -8,6 +8,7 @@ import pytest
 import pytest_asyncio
 
 from tests.utils import RemoteLaunchRenderServer
+from vllm.tokenizers import get_tokenizer
 
 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"
 
@@ -486,3 +487,438 @@ async def test_derender_completion_kv_transfer_params_passthrough(client):
     )
     assert response.status_code == 200
     assert response.json()["kv_transfer_params"] == kv
+
+
+# ---------------------------------------------------------------------------
+# E2E: render -> derender roundtrip with parser (reasoning + tool calls)
+# ---------------------------------------------------------------------------
+
+PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+
+_E2E_TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get weather for a city",
+            "parameters": {
+                "type": "object",
+                "properties": {"city": {"type": "string"}},
+            },
+        },
+    }
+]
+
+
+@pytest.fixture(scope="module")
+def parser_server():
+    args = [
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "hermes",
+        "--reasoning-parser",
+        "deepseek_r1",
+    ]
+    with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def parser_client(parser_server):
+    async with httpx.AsyncClient(
+        base_url=parser_server.url_for(""), timeout=60.0
+    ) as http_client:
+        yield http_client
+
+
+@pytest.fixture(scope="module")
+def parser_tokenizer():
+    return get_tokenizer(PARSER_MODEL)
+
+
+def _encode(tokenizer, text: str) -> list[int]:
+    return tokenizer.encode(text, add_special_tokens=False)
+
+
+def _decoded(tokenizer, token_ids: list[int]) -> str:
+    return tokenizer.decode(token_ids, skip_special_tokens=True)
+
+
+def _require_markers_survive(tokenizer, text: str, *markers: str) -> list[int]:
+    """Encode text and skip the test if any marker is lost in roundtrip."""
+    ids = _encode(tokenizer, text)
+    decoded = tokenizer.decode(ids, skip_special_tokens=False)
+    for m in markers:
+        if m not in decoded:
+            pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip")
+    return ids
+
+
+async def _e2e_render_chat(
+    client: httpx.AsyncClient,
+    model: str,
+    messages: list[dict],
+) -> dict:
+    resp = await client.post(
+        "/v1/chat/completions/render",
+        json={"model": model, "messages": messages},
+    )
+    assert resp.status_code == 200, resp.text
+    return resp.json()
+
+
+def _e2e_generate_response(
+    token_ids: list[int],
+    request_id: str = "chatcmpl-e2e-test",
+) -> dict:
+    return {
+        "request_id": request_id,
+        "choices": [
+            {
+                "index": 0,
+                "token_ids": token_ids,
+                "finish_reason": "stop",
+            }
+        ],
+    }
+
+
+@pytest.mark.asyncio
+async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer):
+    """Plain text without reasoning markers roundtrips correctly."""
+    messages = [{"role": "user", "content": "What is 2+2?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "The answer is four."
+    output_ids = _encode(parser_tokenizer, answer)
+    expected = _decoded(parser_tokenizer, output_ids)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert content == expected
+
+
+@pytest.mark.asyncio
+async def test_e2e_token_identity(parser_client, parser_tokenizer):
+    """encode(derender(token_ids)) == token_ids (RL invariant)."""
+    messages = [{"role": "user", "content": "Hi"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "Hello! How can I help?"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    re_encoded = _encode(parser_tokenizer, content)
+    assert output_ids == re_encoded
+
+
+@pytest.mark.asyncio
+async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer):
+    """CJK + emoji roundtrip without U+FFFD."""
+    messages = [{"role": "user", "content": "Reply in Chinese"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "你好世界 😀"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert "�" not in content
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer):
+    """<think>...</think> splits into reasoning + content."""
+    messages = [{"role": "user", "content": "What is 2+3?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    reasoning_text = "The user wants 2 plus 3. That is 5."
+    answer_text = "The answer is 5."
+    output_text = f"<think>{reasoning_text}</think>{answer_text}"
+    output_ids = _require_markers_survive(parser_tokenizer, output_text, "</think>")
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    msg = resp.json()["choices"][0]["message"]
+    assert msg["reasoning"] is not None
+    assert reasoning_text in msg["reasoning"]
+    assert answer_text in msg["content"]
+    assert "<think>" not in msg["content"]
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer):
+    """<tool_call> extracted into tool_calls field."""
+    messages = [{"role": "user", "content": "Weather in Paris?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    output_text = (
+        "<think>Let me check the weather.</think>"
+        '<tool_call>\n{"name": "get_weather", '
+        '"arguments": {"city": "Paris"}}\n</tool_call>'
+    )
+    output_ids = _require_markers_survive(
+        parser_tokenizer,
+        output_text,
+        "</think>",
+        "<tool_call>",
+        "</tool_call>",
+    )
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "tools": _E2E_TOOLS,
+                "tool_choice": "auto",
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    choice = resp.json()["choices"][0]
+    assert choice["message"]["tool_calls"]
+    assert choice["message"]["tool_calls"][0]["function"]["name"] == "get_weather"
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenizer):
+    """Reasoning + tool call in the same output."""
+    messages = [{"role": "user", "content": "Weather in Paris?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    reasoning_text = "I should look up the weather."
+    tool_text = (
+        '<tool_call>\n{"name": "get_weather", '
+        '"arguments": {"city": "Paris"}}\n</tool_call>'
+    )
+    output_text = f"<think>{reasoning_text}</think>{tool_text}"
+    output_ids = _require_markers_survive(
+        parser_tokenizer, output_text, "</think>", "<tool_call>"
+    )
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "tools": _E2E_TOOLS,
+                "tool_choice": "auto",
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    choice = resp.json()["choices"][0]
+    assert choice["message"]["reasoning"] is not None
+    assert reasoning_text in choice["message"]["reasoning"]
+    assert choice["message"]["tool_calls"]
+
+
+@pytest.mark.asyncio
+async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer):
+    """Without chat_request, derender falls back to plain detokenization."""
+    messages = [{"role": "user", "content": "Hello"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "Hi there!"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert "Hi" in content
+
+
+# ---------------------------------------------------------------------------
+# E2E: HarmonyParser + GPT-OSS
+# ---------------------------------------------------------------------------
+
+HARMONY_MODEL = "openai/gpt-oss-20b"
+
+
+def _ensure_harmony_vocab():
+    """Pre-cache the o200k_base BPE file needed by openai-harmony.
+
+    The Rust tiktoken-rs backend downloads from Azure Blob Storage, which
+    may be unreachable in some environments.  When the cache is cold we
+    fetch the file ourselves and place it in ``/tmp/tiktoken-rs-cache/``
+    using the SHA-1(URL) filename that tiktoken-rs expects.
+    """
+    import hashlib
+    import urllib.request
+    from pathlib import Path
+
+    url = "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
+    cache_dir = Path("/tmp/tiktoken-rs-cache")
+    cache_key = hashlib.sha1(url.encode()).hexdigest()
+    cache_file = cache_dir / cache_key
+    if not cache_file.exists():
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, cache_file)
+
+
+@pytest.fixture(scope="module")
+def harmony_server():
+    _ensure_harmony_vocab()
+    args = [
+        "--trust-remote-code",
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "openai",
+        "--reasoning-parser",
+        "openai_gptoss",
+    ]
+    with RemoteLaunchRenderServer(HARMONY_MODEL, args) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def harmony_client(harmony_server):
+    async with httpx.AsyncClient(
+        base_url=harmony_server.url_for(""), timeout=60.0
+    ) as http_client:
+        yield http_client
+
+
+@pytest.fixture(scope="module")
+def harmony_tokenizer():
+    return get_tokenizer(HARMONY_MODEL, trust_remote_code=True)
+
+
+def _harmony_extract_assistant_ids(
+    tokenizer, assistant_msg: dict, user_content: str = "test"
+) -> list[int]:
+    """Extract assistant token IDs via apply_chat_template diff."""
+    prompt = [{"role": "user", "content": user_content}]
+    full = prompt + [assistant_msg]
+    text_prompt = tokenizer.apply_chat_template(
+        prompt, add_generation_prompt=True, tokenize=False
+    )
+    text_full = tokenizer.apply_chat_template(
+        full, add_generation_prompt=False, tokenize=False
+    )
+    prompt_ids = tokenizer.encode(text_prompt)
+    full_ids = tokenizer.encode(text_full)
+    assistant_ids = list(full_ids[len(prompt_ids) :])
+    if not assistant_ids:
+        pytest.skip("Could not extract assistant tokens for Harmony")
+    return assistant_ids
+
+
+@pytest.mark.asyncio
+async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer):
+    """GPT-OSS content-only roundtrip."""
+    messages = [{"role": "user", "content": "What is 2+2?"}]
+    gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
+
+    assistant_msg = {"role": "assistant", "content": "Four."}
+    output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
+
+    resp = await harmony_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": HARMONY_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": HARMONY_MODEL,
+                "messages": messages,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert content is not None and len(content) > 0
+    assert "Four" in content
+
+
+@pytest.mark.asyncio
+async def test_e2e_harmony_reasoning(harmony_client, harmony_tokenizer):
+    """GPT-OSS reasoning: analysis channel extracted."""
+    messages = [{"role": "user", "content": "Add 2 and 3."}]
+    gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
+
+    reasoning_text = "The user wants 2 plus 3."
+    answer_text = "The answer is 5."
+    assistant_msg = {
+        "role": "assistant",
+        "thinking": reasoning_text,
+        "content": answer_text,
+    }
+    output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
+
+    decoded = harmony_tokenizer.decode(output_ids)
+    if reasoning_text not in decoded:
+        pytest.skip("Harmony template did not render thinking")
+
+    resp = await harmony_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": HARMONY_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": HARMONY_MODEL,
+                "messages": messages,
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    msg = resp.json()["choices"][0]["message"]
+    assert msg["reasoning"] is not None
+    assert reasoning_text in msg["reasoning"]
+    assert answer_text in (msg["content"] or "")

tests/kernels/attention/test_minimax_m3.py

@@ -134,7 +134,6 @@ def _reference_index_topk(
     topk: int,
     init_blocks: int,
     local_blocks: int,
-    sm_scale: float = 1.0,
 ) -> torch.Tensor:
     total_q, num_idx_heads, _ = idx_q.shape
     out = torch.full(
@@ -150,7 +149,7 @@ def _reference_index_topk(
         num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
         pages = block_table[req_id, :num_blocks]
         k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
-        score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale
+        score = torch.einsum("qhd,kd->hqk", q.float(), k.float())
 
         q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
         k_pos = torch.arange(k.shape[0], device=idx_q.device)
@@ -245,270 +244,6 @@ def test_prefill_index_topk_correctness():
     _assert_topk_indices_equal_unordered(actual, expected)
 
 
-# MSA indexer (SM100): fmha_sm100 OnlyScore for the per-block scores, then the
-# Triton minimax_m3_index_topk for selection (no sparse_topk_select). Uses a
-# deterministic construction (idx_q == 1, distinct e4m3-exact per-block values)
-# so scores are strictly monotonic in the block id -> exact top-k agreement.
-def _fmha_indexer_topk(
-    idx_q: torch.Tensor,  # [total_q, H, 128] bf16/e4m3
-    index_cache: torch.Tensor,  # [num_pages, 128, 128] bf16/e4m3
-    block_table: torch.Tensor,
-    q_lens: torch.Tensor,
-    seq_lens: torch.Tensor,
-    prefix_lens: torch.Tensor,
-    sm_scale: float,
-    topk: int,
-) -> torch.Tensor:
-    """Replicate MiniMaxM3IndexerMSAImpl's score path (single decode/prefill side)."""
-    from vllm.third_party.fmha_sm100.api import _fmha_sm100, _fmha_sm100_plan
-
-    num_idx_heads, head_dim = idx_q.shape[1], idx_q.shape[2]
-    nvp = [(s + 127) // 128 for s in seq_lens.tolist()]
-    kv_indices = torch.cat([block_table[r, : nvp[r]] for r in range(len(nvp))]).to(
-        torch.int32
-    )
-
-    qo = q_lens.cpu().to(torch.int32)
-    kv = seq_lens.cpu().to(torch.int32)
-    plan = _fmha_sm100_plan(
-        qo,
-        kv,
-        num_idx_heads,
-        num_kv_heads=1,
-        qo_offset=kv - qo,
-        page_size=128,
-        output_maxscore=True,
-        causal=True,
-        num_kv_splits=1,
-    )
-    k_pages = index_cache.view(index_cache.shape[0], 1, 128, head_dim)
-    _, max_score = _fmha_sm100(
-        idx_q,
-        k_pages,
-        k_pages,
-        plan,
-        kv_indices=kv_indices,
-        output_o=False,
-        output_maxscore=True,
-        sm_scale=sm_scale,
-    )
-
-    batch = q_lens.numel()
-    cu = torch.zeros(batch + 1, dtype=torch.int32, device=idx_q.device)
-    cu[1:] = q_lens.to(torch.int32).cumsum(0)
-    # max_score [H, k_tiles, total_q] -> transpose to [H, total_q, k_tiles].
-    return minimax_m3_index_topk(
-        max_score.transpose(1, 2),
-        cu,
-        prefix_lens.to(torch.int32),
-        int(q_lens.max()),
-        topk,
-        0,  # init_blocks
-        0,  # local_blocks
-    )
-
-
-# e4m3-exact, strictly-increasing per-block values: with idx_q == 1 (also exact)
-# the per-block scores are exact and distinct in BOTH bf16 and e4m3, so the fp8
-# score path selects the same top-k as the reference (no quantization ties).
-_E4M3_EXACT_VALUES = [
-    *range(1, 17),  # 1..16  (step 1)
-    *range(18, 33, 2),  # 18..32 (step 2)
-    *range(36, 65, 4),  # 36..64 (step 4)
-    *range(72, 129, 8),  # 72..128 (step 8)
-]
-
-
-@pytest.mark.skipif(
-    not current_platform.is_device_capability_family(100),
-    reason="fmha_sm100 indexer requires SM100 (Blackwell).",
-)
-@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
-@pytest.mark.parametrize(
-    ("q_lens", "prefix_lens"),
-    [
-        ((4, 3), (2048, 2560)),  # prefill: every token sees >= 16 causal blocks
-        ((1, 1, 1), (2048, 3000, 4096)),  # decode: one query token per request
-    ],
-)
-def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
-    torch.manual_seed(0)
-    num_idx_heads, head_dim = 4, HEAD_DIM
-    device = "cuda"
-
-    q_lens_t = torch.tensor(q_lens, device=device, dtype=torch.int32)
-    prefix_lens_t = torch.tensor(prefix_lens, device=device, dtype=torch.int32)
-    seq_lens = prefix_lens_t + q_lens_t
-    batch = len(q_lens)
-    max_blocks = (int(seq_lens.max()) + BLOCK_SIZE - 1) // BLOCK_SIZE
-    assert max_blocks <= len(_E4M3_EXACT_VALUES)
-    num_pages = batch * max_blocks
-    block_table = torch.randperm(num_pages, device=device, dtype=torch.int32).reshape(
-        batch, max_blocks
-    )
-
-    idx_q = torch.ones(
-        int(q_lens_t.sum()), num_idx_heads, head_dim, device=device, dtype=index_dtype
-    )
-    index_cache = torch.empty(
-        num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
-    )
-    for r in range(batch):
-        for b in range(max_blocks):
-            index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[b])
-
-    sm_scale = head_dim**-0.5
-    actual = _fmha_indexer_topk(
-        idx_q,
-        index_cache,
-        block_table,
-        q_lens_t,
-        seq_lens,
-        prefix_lens_t,
-        sm_scale,
-        TOPK,
-    )
-    expected = _reference_index_topk(
-        idx_q,
-        index_cache,
-        block_table,
-        q_lens_t,
-        seq_lens,
-        prefix_lens_t,
-        TOPK,
-        init_blocks=0,
-        local_blocks=0,
-        sm_scale=sm_scale,
-    )
-    _assert_topk_indices_equal_unordered(actual, expected)
-
-
-# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score +
-# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata
-# builders on the SAME CommonAttentionMetadata + index cache, and assert the
-# selected blocks agree. This exercises all the metadata the impl/kernels consume
-# (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,
-# decode_pages split) -- a metadata bug on either side shifts the causal window
-# or the block->page mapping and breaks the comparison.
-@pytest.mark.skipif(
-    not current_platform.is_device_capability_family(100),
-    reason="fmha_sm100 indexer requires SM100 (Blackwell).",
-)
-@pytest.mark.parametrize("topk", [8, 16])
-def test_msa_indexer_impl_matches_triton(topk, monkeypatch):
-    import vllm.models.minimax_m3.common.indexer as indexer_mod
-    from tests.v1.attention.utils import (
-        BatchSpec,
-        create_common_attn_metadata,
-        create_vllm_config,
-    )
-    from vllm.config import set_current_vllm_config
-    from vllm.forward_context import set_forward_context
-    from vllm.models.minimax_m3.common.indexer import (
-        MiniMaxM3IndexerTritonImpl,
-        MiniMaxM3IndexerTritonMetadataBuilder,
-    )
-    from vllm.models.minimax_m3.nvidia.indexer_msa import (
-        MiniMaxM3IndexerMSAImpl,
-        MiniMaxM3IndexerMSAMetadataBuilder,
-    )
-
-    torch.manual_seed(0)
-    device = torch.device("cuda")
-    num_idx_heads, head_dim = 4, HEAD_DIM
-    # TP=1: avoid requiring an initialized distributed group in a unit test.
-    monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
-
-    vllm_config = create_vllm_config(
-        block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
-    )
-    vllm_config.model_config.hf_config.sparse_attention_config = {
-        "sparse_num_index_heads": num_idx_heads
-    }
-
-    # Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long
-    # prefixes so every token sees > TOPK causal blocks (non-trivial selection).
-    batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96])
-    common = create_common_attn_metadata(
-        batch, BLOCK_SIZE, device, arange_block_indices=True
-    )
-    num_tokens = batch.compute_num_tokens()
-
-    # Deterministic index cache: distinct, monotonic per-logical-block values so
-    # the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
-    block_table = common.block_table_tensor
-    num_pages = int(block_table.max().item()) + 1
-    index_cache = torch.zeros(
-        num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE
-    )
-    for r, seq_len in enumerate(batch.seq_lens):
-        for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):
-            index_cache[block_table[r, b]] = float(b + 1)
-    index_q = torch.ones(
-        num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE
-    )
-
-    spec = MLAAttentionSpec(
-        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE
-    )
-    impl_kwargs = dict(
-        num_kv_heads=num_idx_heads,
-        scale=head_dim**-0.5,
-        topk_blocks=topk,
-        sparse_block_size=BLOCK_SIZE,
-        num_index_heads=num_idx_heads,
-        index_head_dim=head_dim,
-        init_blocks=0,
-        local_blocks=0,
-    )
-
-    with set_current_vllm_config(vllm_config):
-        msa_impl = MiniMaxM3IndexerMSAImpl(prefix="idx_msa", **impl_kwargs)
-        triton_impl = MiniMaxM3IndexerTritonImpl(prefix="idx_triton", **impl_kwargs)
-        msa_builder = MiniMaxM3IndexerMSAMetadataBuilder(
-            spec, [msa_impl.index_cache.prefix], vllm_config, device
-        )
-        triton_builder = MiniMaxM3IndexerTritonMetadataBuilder(
-            spec, [triton_impl.index_cache.prefix], vllm_config, device
-        )
-
-    # Both impls score against the same index keys.
-    msa_impl.index_cache.kv_cache = index_cache
-    triton_impl.index_cache.kv_cache = index_cache
-
-    # Exercise the shared persistent top-k buffer for BOTH impls: each must write
-    # decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views.
-    # Separate buffers so the two forwards don't clobber each other.
-    nd = sum(q for q in batch.query_lens if q <= 1)
-    msa_impl.topk_indices_buffer = torch.full(
-        (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
-    )
-    triton_impl.topk_indices_buffer = torch.full(
-        (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
-    )
-
-    attn_metadata = {
-        msa_impl.index_cache.prefix: msa_builder.build(0, common),
-        triton_impl.index_cache.prefix: triton_builder.build(0, common),
-    }
-    with set_forward_context(attn_metadata, vllm_config):
-        msa_decode, msa_prefill = msa_impl(index_q)
-        tri_decode, tri_prefill = triton_impl(index_q)
-
-    assert msa_decode is not None and tri_decode is not None
-    assert msa_prefill is not None and tri_prefill is not None
-    _assert_topk_indices_equal_unordered(msa_decode, tri_decode)
-    _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill)
-    # decode/prefill outputs are views into each impl's persistent buffer.
-    for impl, dec, pre in (
-        (msa_impl, msa_decode, msa_prefill),
-        (triton_impl, tri_decode, tri_prefill),
-    ):
-        buf = impl.topk_indices_buffer
-        assert dec.data_ptr() == buf[:, :nd, :].data_ptr()
-        assert pre.data_ptr() == buf[:, nd:, :].data_ptr()
-
-
 @pytest.mark.parametrize(
     ("decode_query_len", "max_decode_query_len"),
     [
@@ -582,65 +317,6 @@ def test_decode_index_topk_correctness(
     _assert_topk_indices_equal_unordered(actual, expected)
 
 
-@pytest.mark.skipif(
-    not current_platform.is_device_capability_family(100),
-    reason="fp8 e4m3 indexer cache is the SM100 (MSA) path.",
-)
-@pytest.mark.parametrize("num_idx_heads", [1, 4])
-def test_decode_index_topk_fp8(num_idx_heads: int):
-    """The fp8 (e4m3) indexer cache feeds the Triton decode kernel on the MSA
-    path. The kernel must score in fp32 (no scaling) so its top-k matches a
-    reference computed from the dequantized fp8 values."""
-    torch.manual_seed(0)
-    topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128
-    decode_query_len = 1
-    active_seq_lens = torch.tensor((129, 1025, 4097), device="cuda", dtype=torch.int32)
-    q_lens = torch.full_like(active_seq_lens, decode_query_len)
-    prefix_lens = active_seq_lens - decode_query_len
-    batch = active_seq_lens.numel()
-    max_seq_len = int(active_seq_lens.max())
-    max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
-    num_pages = batch * max_blocks
-    block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
-        batch, max_blocks
-    )
-    idx_q = torch.randn(
-        batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
-    ).to(torch.float8_e4m3fn)
-    index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to(
-        torch.float8_e4m3fn
-    )
-
-    actual = minimax_m3_index_decode(
-        idx_q,
-        index_kv_cache,
-        block_table,
-        active_seq_lens,
-        max_seq_len=max_seq_len,
-        topk=topk,
-        init_blocks=init_blocks,
-        local_blocks=local_blocks,
-        num_kv_heads=num_idx_heads,
-        sm_scale=head_dim**-0.5,
-        decode_query_len=decode_query_len,
-    )
-    # Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK
-    # in fp32, so it must match an fp32 matmul of the same e4m3 values).
-    expected = _reference_index_topk(
-        idx_q.float(),
-        index_kv_cache.float(),
-        block_table,
-        q_lens,
-        active_seq_lens,
-        prefix_lens,
-        topk,
-        init_blocks,
-        local_blocks,
-        head_dim**-0.5,
-    )
-    _assert_topk_indices_equal_unordered(actual, expected)
-
-
 # Sparse attention kernels.
 def _reference_sparse_attn(
     q: torch.Tensor,

tests/kernels/moe/test_trtllm_nvfp4_moe.py

@@ -5,9 +5,8 @@ Tests for the FlashInfer TRTLLM NvFP4 MoE backend
 (`TrtLlmNvFp4ExpertsModular`).
 
 Covers the activations the wrapper claims to support — SiLU, RELU^2 (non-gated),
-GELU, and clamped SwiGLU-OAI (MiniMax-M3) — including a Gemma4-shaped case
-(128 experts, top-k 8, intermediate_size 704) that exercises the non-256-aligned
-padding path.
+and GELU — including a Gemma4-shaped case (128 experts, top-k 8,
+intermediate_size 704) that exercises the non-256-aligned padding path.
 """
 
 import pytest
@@ -81,29 +80,6 @@ if _CLAMP_OP_NAME not in op_registry:
 
 SILU_WITH_CLAMP = op_registry[_CLAMP_OP_NAME]
 
-# Clamped SwiGLU-OAI (MiniMax-M3): non-default alpha/beta so the kernel must
-# honor gemm1_alpha (raw) and gemm1_beta (folded by g1_alphas), not just clamp.
-_SWIGLU_ALPHA = 1.702
-_SWIGLU_BETA = 1.0
-_OAI_OP_NAME = "test_swigluoai_with_clamp"
-
-if _OAI_OP_NAME not in op_registry:
-
-    @CustomOp.register(_OAI_OP_NAME)
-    class _SwigluOAIWithClampTest(SiluAndMulWithClamp):
-        custom_op_name = _OAI_OP_NAME
-
-        def __init__(self, *, compile_native: bool = True) -> None:
-            super().__init__(
-                _SWIGLU_LIMIT,
-                alpha=_SWIGLU_ALPHA,
-                beta=_SWIGLU_BETA,
-                compile_native=compile_native,
-            )
-
-
-SWIGLUOAI_REF = op_registry[_OAI_OP_NAME]
-
 
 ACTIVATION_CASES = [
     pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"),
@@ -115,12 +91,6 @@ ACTIVATION_CASES = [
         id="relu2_no_mul",
     ),
     pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, id="gelu"),
-    pytest.param(
-        MoEActivation.SWIGLUOAI_UNINTERLEAVE,
-        SWIGLUOAI_REF,
-        _SWIGLU_LIMIT,
-        id="swigluoai_uninterleave",
-    ),
 ]
 
 
@@ -178,10 +148,6 @@ def test_trtllm_fp4_moe_no_graph(
             is_scale_swizzled=False,
         )
         quant_config.gemm1_clamp_limit = swiglu_limit
-        is_oai = activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE
-        if is_oai:
-            quant_config.gemm1_alpha = _SWIGLU_ALPHA
-            quant_config.gemm1_beta = _SWIGLU_BETA
         if swiglu_limit is not None:
             assert quant_config.g1_alphas is not None
             assert quant_config.a2_gscale is not None
@@ -226,27 +192,6 @@ def test_trtllm_fp4_moe_no_graph(
         fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas)
         trtllm_inner.process_weights_after_loading(fake_layer)
 
-        if is_oai:
-            # alpha stays raw; beta and clamp are folded by g1_alphas
-            # (== _LARGE_OUTPUT1_SCALE here), so the fold is load-bearing.
-            assert torch.allclose(
-                trtllm_inner.gemm1_alpha,
-                torch.full_like(trtllm_inner.gemm1_alpha, _SWIGLU_ALPHA),
-            )
-            assert torch.allclose(
-                trtllm_inner.gemm1_beta,
-                torch.full_like(
-                    trtllm_inner.gemm1_beta, _SWIGLU_BETA / _LARGE_OUTPUT1_SCALE
-                ),
-            )
-            assert torch.allclose(
-                trtllm_inner.gemm1_clamp_limit,
-                torch.full_like(
-                    trtllm_inner.gemm1_clamp_limit,
-                    _SWIGLU_LIMIT / _LARGE_OUTPUT1_SCALE,
-                ),
-            )
-
         trtllm_experts = mk.FusedMoEKernel(
             maybe_make_prepare_finalize(
                 moe=moe_config,

tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py

@@ -278,99 +278,3 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
     torch.testing.assert_close(
         index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
     )
-
-
-# ── Test 3: fp8 (e4m3) index outputs ─────────────────────────────────────────
-# The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
-# q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
-# (the index dtype must not perturb the main branch), and (2) the e4m3 index
-# outputs dequantize close to the bf16 reference.
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9),
-    reason="e4m3 conversion requires CUDA SM89+.",
-)
-@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
-@pytest.mark.parametrize("block_size", [16, 64])
-def test_sparse_full_fp8_index(num_tokens, block_size):
-    torch.manual_seed(1)
-    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
-    base, max_pos = 5_000_000.0, 4096
-    num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
-
-    q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
-    k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
-    iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
-    ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
-    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
-    positions = torch.randint(
-        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
-    )
-
-    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
-    iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
-    qkv0 = torch.randn(
-        num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
-    )
-
-    num_blocks = (num_tokens + block_size - 1) // block_size + 1
-    slot_mapping = torch.randperm(
-        num_blocks * block_size, dtype=torch.int64, device=device
-    )[:num_tokens]
-    index_slot_mapping = torch.roll(slot_mapping, shifts=1)
-
-    def run(index_dtype):
-        qkv = qkv0.clone()
-        kv_cache = torch.zeros(
-            num_blocks,
-            2,
-            block_size,
-            num_kv_heads,
-            HEAD_DIM,
-            dtype=dtype,
-            device=device,
-        )
-        index_cache = torch.zeros(
-            num_blocks, block_size, HEAD_DIM, dtype=index_dtype, device=device
-        )
-        q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
-        index_q = torch.empty(num_tokens, iqsz, dtype=index_dtype, device=device)
-        ops.fused_minimax_m3_qknorm_rope_kv_insert(
-            qkv,
-            q_w,
-            k_w,
-            cos_sin,
-            positions,
-            num_heads,
-            num_kv_heads,
-            ROTARY_DIM,
-            eps,
-            iq_w,
-            ik_w,
-            num_idx_heads,
-            slot_mapping,
-            index_slot_mapping,
-            kv_cache,
-            index_cache,
-            block_size,
-            q_out,
-            index_q,
-        )
-        return qkv, kv_cache, index_cache, q_out, index_q
-
-    qkv_bf, kvc_bf, idxc_bf, qo_bf, iq_bf = run(torch.bfloat16)
-    qkv_fp, kvc_fp, idxc_fp, qo_fp, iq_fp = run(torch.float8_e4m3fn)
-
-    assert iq_fp.dtype == torch.float8_e4m3fn
-    assert idxc_fp.dtype == torch.float8_e4m3fn
-
-    # (1) The main branch (q/k/v in qkv, q_out, kv cache) must be bit-identical:
-    # the index output dtype must not perturb anything else.
-    torch.testing.assert_close(qo_fp, qo_bf, rtol=0, atol=0)
-    torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0)
-    torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0)
-
-    # (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp.
-    torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05)
-    torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05)

tests/reasoning/test_minimax_m3_reasoning_parser.py

@@ -83,6 +83,20 @@ class MiniMaxM3Tokenizer:
         return "".join(tokens)
 
 
+class SplitMiniMaxM3Tokenizer(MiniMaxM3Tokenizer):
+    """Tokenizer that exposes marker vocab entries but encodes them as text."""
+
+    def tokenize(self, text: str) -> list[str]:
+        return list(text)
+
+
+class RuntimeSplitMiniMaxM3Tokenizer(MiniMaxM3Tokenizer):
+    """Tokenizer whose runtime output splits markers despite atomic encodes."""
+
+    def encode_runtime(self, text: str) -> list[int]:
+        return [self._add_token(token) for token in list(text)]
+
+
 def make_parser(
     chat_template_kwargs: dict[str, str] | None = None,
 ) -> tuple[MiniMaxM3ReasoningParser, MiniMaxM3Tokenizer]:
@@ -105,7 +119,8 @@ def run_streaming(
     reasoning_end_states: list[bool] = []
 
     for chunk in chunks:
-        delta_token_ids = tokenizer.encode(chunk, add_special_tokens=False)
+        encode_runtime = getattr(tokenizer, "encode_runtime", tokenizer.encode)
+        delta_token_ids = encode_runtime(chunk)
         current_text = previous_text + chunk
         current_token_ids = previous_token_ids + delta_token_ids
         delta = parser.extract_reasoning_streaming(
@@ -174,14 +189,14 @@ def test_nonstreaming_drops_leading_end_tag():
     assert content == "answer"
 
 
-def test_nonstreaming_non_leading_end_tag_is_content():
+def test_nonstreaming_end_tag_in_content_state_is_dropped():
     parser, _ = make_parser()
     request = ChatCompletionRequest(messages=[], model="test-model")
 
     reasoning, content = parser.extract_reasoning("XXX</mm:think>YYY", request)
 
     assert reasoning is None
-    assert content == "XXX</mm:think>YYY"
+    assert content == "XXXYYY"
 
 
 def test_nonstreaming_enabled_mode_starts_in_reasoning():
@@ -246,7 +261,7 @@ def test_streaming_drops_leading_end_tag():
     assert end_states == [True, True]
 
 
-def test_streaming_non_leading_end_tag_is_content():
+def test_streaming_end_tag_in_content_state_is_dropped():
     parser, tokenizer = make_parser()
 
     reasoning, content, end_states = run_streaming(
@@ -256,7 +271,7 @@ def test_streaming_non_leading_end_tag_is_content():
     )
 
     assert reasoning is None
-    assert content == "XXX</mm:think>YYY"
+    assert content == "XXXYYY"
     assert end_states == [True]
 
 
@@ -288,6 +303,110 @@ def test_streaming_plain_content_ends_reasoning_phase():
     assert end_states == [True, True]
 
 
+def test_streaming_split_marker_tokens_are_not_returned():
+    tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
+    parser = MiniMaxM3ReasoningParser(tokenizer)
+
+    reasoning, content, end_states = run_streaming(
+        parser,
+        tokenizer,
+        ["<mm:think>", "Reasoning", " content", "</mm:think>", "content"],
+    )
+
+    assert reasoning == "Reasoning content"
+    assert content == "content"
+    assert end_states == [False, False, False, True, True]
+
+
+def test_streaming_split_marker_text_drives_end_state():
+    tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
+    parser = MiniMaxM3ReasoningParser(tokenizer)
+    previous_text = ""
+    previous_token_ids: list[int] = []
+
+    for chunk in ["<mm:think>", "Reasoning", " content", "</mm:think>"]:
+        delta_token_ids = tokenizer.encode_runtime(chunk)
+        current_text = previous_text + chunk
+        current_token_ids = previous_token_ids + delta_token_ids
+        parser.extract_reasoning_streaming(
+            previous_text=previous_text,
+            current_text=current_text,
+            delta_text=chunk,
+            previous_token_ids=previous_token_ids,
+            current_token_ids=current_token_ids,
+            delta_token_ids=delta_token_ids,
+        )
+        previous_text = current_text
+        previous_token_ids = current_token_ids
+
+    assert parser.is_reasoning_end_streaming(previous_token_ids, []) is True
+
+
+def test_streaming_split_marker_tokens_enabled_mode():
+    tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
+    parser = MiniMaxM3ReasoningParser(
+        tokenizer, chat_template_kwargs={"thinking_mode": "enabled"}
+    )
+
+    reasoning, content, end_states = run_streaming(
+        parser,
+        tokenizer,
+        ["Reasoning", " content", "</mm:think>", "content"],
+    )
+
+    assert reasoning == "Reasoning content"
+    assert content == "content"
+    assert end_states == [False, False, True, True]
+
+
+def test_streaming_split_marker_text_across_deltas():
+    tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
+    parser = MiniMaxM3ReasoningParser(tokenizer)
+
+    reasoning, content, end_states = run_streaming(
+        parser,
+        tokenizer,
+        ["<mm:", "think>", "Reasoning", " content", "</mm:", "think>", "content"],
+    )
+
+    assert reasoning == "Reasoning content"
+    assert content == "content"
+    assert end_states == [False, False, False, False, False, True, True]
+
+
+def test_streaming_split_leading_end_marker_text_across_deltas():
+    tokenizer = RuntimeSplitMiniMaxM3Tokenizer()
+    parser = MiniMaxM3ReasoningParser(tokenizer)
+
+    reasoning, content, end_states = run_streaming(
+        parser,
+        tokenizer,
+        ["</mm:", "think>", "content"],
+    )
+
+    assert reasoning is None
+    assert content == "content"
+    assert end_states == [False, True, True]
+
+
+def test_token_id_helpers_with_split_marker_tokens():
+    tokenizer = SplitMiniMaxM3Tokenizer()
+    parser = MiniMaxM3ReasoningParser(tokenizer)
+    output_ids = tokenizer.encode(
+        "<mm:think>abc</mm:think>def", add_special_tokens=False
+    )
+    open_reasoning_ids = tokenizer.encode("<mm:think>abc", add_special_tokens=False)
+    content_ids = tokenizer.encode("plain", add_special_tokens=False)
+
+    assert parser.is_reasoning_end(output_ids)
+    assert not parser.is_reasoning_end(open_reasoning_ids)
+    assert not parser.is_reasoning_end(content_ids)
+    assert tokenizer.decode(parser.extract_content_ids(output_ids)) == "def"
+    assert parser.extract_content_ids(open_reasoning_ids) == []
+    assert parser.extract_content_ids(content_ids) == content_ids
+    assert parser.count_reasoning_tokens(output_ids) == len(tokenizer.encode("abc"))
+
+
 def test_token_id_helpers():
     parser, tokenizer = make_parser()
     output_ids = tokenizer.encode(

tests/v1/core/test_contiguous_kv_packing.py

@@ -1,16 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for contiguous KV cache packing in _get_kv_cache_config_deepseek_v4."""
+"""Tests for contiguous KV cache packing."""
 
 from unittest.mock import MagicMock
 
 import pytest
 import torch
 
-from vllm.v1.core.kv_cache_utils import _get_kv_cache_config_deepseek_v4
+from vllm import envs
+from vllm.v1.core.kv_cache_utils import (
+    _get_kv_cache_config_deepseek_v4,
+    get_kv_cache_config_from_groups,
+)
 from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
     KVCacheGroupSpec,
+    KVCacheTensor,
     MLAAttentionSpec,
+    SlidingWindowSpec,
     UniformTypeKVCacheSpecs,
 )
 
@@ -28,6 +35,25 @@ def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
     )
 
 
+def _make_full_spec() -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=64,
+        dtype=torch.float16,
+    )
+
+
+def _make_sw_spec() -> SlidingWindowSpec:
+    return SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=128,
+    )
+
+
 def _make_groups(n_c4, n_c128, n_swa):
     PS_C4_MLA = 37440
     PS_C4_IDX = 8640
@@ -130,6 +156,73 @@ class TestInterleavedPacking:
         for i, v in enumerate(views):
             assert (v == i + 1).all(), f"View {i} was corrupted"
 
+    def test_hma_attention_groups_keep_default_backing(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False)
+        full = _make_full_spec()
+        sw = _make_sw_spec()
+        page_size = full.page_size_bytes
+        groups = [
+            KVCacheGroupSpec(["full.0", "full.1"], full),
+            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
+            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
+        ]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
+        )
+
+        assert config.num_blocks == 32
+        assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32
+        assert config.kv_cache_tensors == [
+            KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),
+            KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
+        ]
+
+    def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False)
+        full = _make_full_spec()
+        sw = _make_sw_spec()
+        page_size = full.page_size_bytes
+        groups = [
+            KVCacheGroupSpec(["full.0", "full.1"], full),
+            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
+            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
+        ]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
+        )
+
+        assert config.num_blocks == 32
+        assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}
+        assert config.kv_cache_tensors == [
+            KVCacheTensor(
+                size=page_size * 2 * 32,
+                shared_by=["full.0", "sw.0", "sw.1"],
+                offset=0,
+                block_stride=page_size * 2,
+            ),
+            KVCacheTensor(
+                size=page_size * 2 * 32,
+                shared_by=["full.1", "sw.2", "sw.3"],
+                offset=page_size,
+                block_stride=page_size * 2,
+            ),
+        ]
+
+    def test_single_group_attention_keeps_unpacked_layout(self):
+        spec = _make_full_spec()
+        groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
+        )
+
+        assert sum(t.size for t in config.kv_cache_tensors) == (
+            spec.page_size_bytes * 2 * 32
+        )
+        assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]
+
 
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

tests/v1/core/test_scheduler.py

@@ -144,6 +144,43 @@ def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders():
     assert req.request_id in output.num_scheduled_tokens
 
 
+def test_cached_request_data_resumed_all_token_ids_mrv1_only():
+    """all_token_ids carries a resumed request's token ids to the connector
+    for the V1 model runner, but is skipped entirely for the V2 model runner.
+    """
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+
+    scheduler = create_scheduler()
+    (req,) = create_requests(num_requests=1, num_tokens=8)
+    req.append_output_token_ids([101, 102, 103])
+
+    # A resumed request was not scheduled in the previous step.
+    assert req.request_id not in scheduler.prev_step_scheduled_req_ids
+
+    empty_blocks = KVCacheBlocks(blocks=((),))
+
+    def make_cached():
+        return scheduler._make_cached_request_data(
+            running_reqs=[],
+            resumed_reqs=[req],
+            num_scheduled_tokens={req.request_id: 1},
+            spec_decode_tokens={},
+            req_to_new_blocks={req.request_id: empty_blocks},
+        )
+
+    # V1 model runner: the full token id list is propagated.
+    assert not scheduler.use_v2_model_runner
+    cached = make_cached()
+    assert req.request_id in cached.resumed_req_ids
+    assert cached.all_token_ids[req.request_id] == list(req.all_token_ids)
+
+    # V2 model runner: all_token_ids is skipped entirely.
+    scheduler.use_v2_model_runner = True
+    cached = make_cached()
+    assert req.request_id in cached.resumed_req_ids
+    assert cached.all_token_ids == {}
+
+
 def test_schedule_partial_requests():
     """Test scheduling behavior with partial requests.
 

tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py

@@ -7,7 +7,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp
     ExternalCachedBlockPool,
     MooncakeStoreCoordinator,
 )
-from vllm.v1.core.kv_cache_utils import BlockHash, BlockHashListWithBlockSize
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
+    chunk_hashes_for_block_size,
+)
+from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheGroupSpec,
@@ -182,7 +185,7 @@ def test_coordinator_group_block_size_double_hash():
     ]
     coord = _make_coord(groups, hash_block_size=16)
     hs = _hashes(4)
-    big_hashes = list(BlockHashListWithBlockSize(hs, 16, 32))
+    big_hashes = list(chunk_hashes_for_block_size(hs, 16, 32))
     exists = {(0, bytes(h)) for h in hs}
     exists |= {(1, bytes(bh)) for bh in big_hashes}
     cmap = ExternalCachedBlockPool(exists)

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

@@ -323,8 +323,8 @@ def test_recv_skips_swa_blocks_before_window():
 
 def test_chunked_token_database_hash_block_size_smaller_than_block_size():
     """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens
-    must merge every 4 fine hashes into one chunk hash via
-    BlockHashListWithBlockSize."""
+    keys each 16-token chunk by its last fine hash, keeping the Mooncake key
+    at one digest instead of concatenating all 4 fine hashes."""
     md = KeyMetadata("m", 0, 0, 0, 0, group_id=3)
     db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4)
     db.set_kv_caches_base_addr([0])
@@ -335,8 +335,7 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size():
     assert len(out) == 2
     assert out[0][0] == 0 and out[0][1] == 16
     assert out[1][0] == 16 and out[1][1] == 32
-    # Each chunk's hash is the concatenation of 4 fine hashes.
-    expected0 = b"".join(fine_hashes[0:4]).hex()
-    expected1 = b"".join(fine_hashes[4:8]).hex()
-    assert out[0][2].chunk_hash == expected0
-    assert out[1][2].chunk_hash == expected1
+    # Each chunk's hash is its last (4th) fine hash, which already chains the
+    # prior three.
+    assert out[0][2].chunk_hash == fine_hashes[3].hex()
+    assert out[1][2].chunk_hash == fine_hashes[7].hex()

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
     worker as mooncake_store_worker,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
+    BlobBlockHashes,
     ChunkedTokenDatabase,
     KeyMetadata,
     LoadSpec,
@@ -32,6 +33,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
     MooncakeStoreConnectorStats,
 )
+from vllm.v1.core.kv_cache_utils import BlockHash
 
 
 def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:
@@ -1179,9 +1181,9 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():
     assert full_event.group_idx == 0
     assert full_event.block_size == 32
     assert full_event.token_ids == list(range(32))
-    assert full_event.block_hashes == [
-        maybe_convert_block_hash(BlockHash(b"".join(hs)))
-    ]
+    # block_size=32 over hash_block_size=8 (scale 4): the chunk is keyed by its
+    # last sub-hash, not the concatenation of all four.
+    assert full_event.block_hashes == [maybe_convert_block_hash(BlockHash(hs[3]))]
 
     assert swa_event.group_idx == 1
     assert swa_event.block_size == 8
@@ -1749,3 +1751,33 @@ def test_store_worker_close_swallows_store_errors():
     worker.close()
 
     assert worker.store is None
+
+
+def test_blob_block_hashes_wire_roundtrip():
+    """The lookup wire format sends a ``hash_len`` frame plus the raw hashes
+    concatenated back-to-back; the server rebuilds them through a zero-copy
+    ``BlobBlockHashes`` view over the frame buffer."""
+    hashes = [BlockHash(bytes([i]) * 16) for i in range(5)]
+    hash_len = len(hashes[0])
+
+    # Client side (LookupKeyClient._lookup): flat payload frame.
+    blob = b"".join(hashes)
+
+    # Server side (LookupKeyServer): view over the frame buffer (a memoryview),
+    # never materializing the full hash list upfront.
+    view = BlobBlockHashes(memoryview(blob), hash_len)
+
+    assert len(view) == 5
+    assert list(view) == hashes  # default Sequence iter terminates via IndexError
+    assert [bytes(h) for h in view] == hashes
+    assert bytes(view[-1]) == hashes[-1]
+    assert [bytes(h) for h in view[1:3]] == hashes[1:3]
+    with pytest.raises(IndexError):
+        _ = view[5]
+
+
+def test_blob_block_hashes_empty():
+    """Empty lookups send hash_len=0 and an empty payload."""
+    view = BlobBlockHashes(memoryview(b""), 0)
+    assert len(view) == 0
+    assert list(view) == []

tests/v1/kv_offload/cpu/test_manager.py

@@ -14,12 +14,13 @@ from vllm.v1.kv_offload.base import (
     ReqContext,
     make_offload_key,
 )
-from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
+from vllm.v1.kv_offload.cpu.common import (
+    CPULoadStoreSpec,
+    CPUOffloadingMetrics,
+)
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
 
-STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
-
 
 def make_req_context(
     req_id: str = "", kv_transfer_params: dict | None = None
@@ -181,10 +182,45 @@ def test_filter_reused_manager_reports_stores_skipped_counter():
     )
     stats = manager.get_stats()
     assert stats is not None
-    assert stats.reduce()[STORES_SKIPPED] == 3
+    assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 3
     stats = manager.get_stats()
     assert stats is not None
-    assert stats.reduce()[STORES_SKIPPED] == 0
+    assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 0
+
+
+def test_cpu_manager_reports_cache_usage_gauge():
+    def check_usage_stats(manager: CPUOffloadingManager, value: float):
+        stats = manager.get_stats()
+        assert stats is not None
+        assert stats.reduce()[
+            CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC
+        ] == pytest.approx(value)
+
+    # Zero-capacity manager always reports 0.0
+    manager = make_cpu_manager(num_blocks=0)
+    check_usage_stats(manager, 0.0)
+
+    # Empty manager (4 blocks, none allocated): usage = 0.0
+    manager = make_cpu_manager(num_blocks=4)
+    check_usage_stats(manager, 0.0)
+
+    # After allocating 2 of 4 blocks: usage = 0.5
+    manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    check_usage_stats(manager, 0.5)
+
+    # After filling all 4 blocks: usage = 1.0
+    manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
+    check_usage_stats(manager, 1.0)
+
+    # After completing store, the blocks becomes evictable as it is not actively used
+    # and usage drops.
+    manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    check_usage_stats(manager, 0.5)
+
+    # After completing store, the blocks becomes evictable as it is not actively used
+    # and usage drops.
+    manager.complete_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
+    check_usage_stats(manager, 0.0)
 
 
 def test_cpu_manager():

tests/v1/logits_processors/test_correctness.py

@@ -145,7 +145,6 @@ def _generate_fake_sampling_metadata(
         vllm_config.scheduler_config.max_num_seqs,
         num_spec,
         device,
-        PIN_MEMORY_AVAILABLE,
     )
     fake_sampling_metadata = SamplingMetadata(
         temperature=torch.full((batch_size,), 0.0),
@@ -880,7 +879,6 @@ def test_maybe_create_thinking_budget_holder_without_reasoning():
             cfg.scheduler_config.max_num_seqs,
             0,
             torch.device("cpu"),
-            False,
         )
         is None
     )

tests/v1/simple_kv_offload/test_scheduler.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
+import pytest
 import torch
 
 from vllm import SamplingParams
@@ -1528,3 +1529,232 @@ def test_reset_pending_loads() -> None:
     # All GPU blocks free
     num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()
     assert num_used == 1, f"Expected only null block in use, got {num_used}"
+
+
+def _make_cp_vllm_config(
+    dcp_world_size: int = 1,
+    pcp_world_size: int = 1,
+) -> VllmConfig:
+    """VllmConfig with context-parallel sizes set for scheduler-only tests."""
+    cfg = _make_vllm_config()
+
+    cfg.parallel_config.decode_context_parallel_size = dcp_world_size
+    cfg.parallel_config.prefill_context_parallel_size = pcp_world_size
+    return cfg
+
+
+def _make_cp_scheduler(
+    *,
+    dcp_world_size: int = 1,
+    pcp_world_size: int = 1,
+    num_cpu_blocks: int = 8,
+    num_gpu_blocks: int = 16,
+    lazy: bool = False,
+) -> SchedulerFixture:
+    """Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size."""
+    cp_world_size = dcp_world_size * pcp_world_size
+    virtual_block_size = BLOCK_SIZE * cp_world_size
+
+    kv_cache_config = _make_kv_cache_config(num_gpu_blocks)
+    vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size)
+    cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks
+
+    sched = SimpleCPUOffloadScheduler(
+        vllm_config=vllm_config,
+        kv_cache_config=kv_cache_config,
+        cpu_capacity_bytes=cpu_capacity_bytes,
+        scheduler_block_size=virtual_block_size,
+        hash_block_size=virtual_block_size,
+        lazy_offload=lazy,
+    )
+
+    gpu_block_pool = BlockPool(
+        num_gpu_blocks=num_gpu_blocks,
+        enable_caching=True,
+        hash_block_size=virtual_block_size,
+    )
+    sched.bind_gpu_block_pool(gpu_block_pool)
+
+    return SchedulerFixture(
+        scheduler=sched,
+        gpu_block_pool=gpu_block_pool,
+        vllm_config=vllm_config,
+        kv_cache_config=kv_cache_config,
+    )
+
+
+def _make_cp_request(
+    num_blocks: int,
+    virtual_block_size: int,
+    request_id: str | None = None,
+) -> Request:
+    """Create a request whose block hashes are computed at the virtual
+    (CP-scaled) block size, matching what the real scheduler does.
+    """
+    global _req_counter
+    _req_counter += 1
+    if request_id is None:
+        request_id = f"req-cp-{_req_counter}"
+
+    num_tokens = num_blocks * virtual_block_size + 1
+    start = _req_counter * 10000
+    prompt_token_ids = list(range(start, start + num_tokens))
+    sampling_params = SamplingParams(max_tokens=1)
+
+    return Request(
+        request_id=request_id,
+        prompt_token_ids=prompt_token_ids,
+        sampling_params=sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=get_request_block_hasher(virtual_block_size, sha256),
+    )
+
+
+def _allocate_cp_gpu_blocks(
+    gpu_block_pool: BlockPool,
+    request: Request,
+    num_blocks: int,
+    virtual_block_size: int,
+    group_id: int = 0,
+) -> list:
+    """Allocate GPU blocks and cache them using the CP-scaled block size."""
+    blocks = gpu_block_pool.get_new_blocks(num_blocks)
+    num_full = min(num_blocks, len(request.block_hashes))
+    if num_full > 0:
+        gpu_block_pool.cache_full_blocks(
+            request=request,
+            blocks=blocks,
+            num_cached_blocks=0,
+            num_full_blocks=num_full,
+            block_size=virtual_block_size,
+            kv_cache_group_id=group_id,
+        )
+    return blocks
+
+
+# ---------------------------------------------------------------------------
+# Test 15: CP block size scaling is correct
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+    "dcp_world_size, pcp_world_size",
+    [
+        (2, 1),  # DCP only
+        (1, 2),  # PCP only
+        (2, 2),  # DCP + PCP
+    ],
+)
+def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None:
+    """Verify that the scheduler's block_size and cp_world_size are correctly
+    scaled when context parallelism is enabled."""
+    fix = _make_cp_scheduler(
+        dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
+    )
+    sched = fix.scheduler
+
+    expected_cp = dcp_world_size * pcp_world_size
+    assert sched.cp_world_size == expected_cp
+    assert sched.block_size == BLOCK_SIZE * expected_cp
+
+
+# ---------------------------------------------------------------------------
+# Test 16: CP eager store-and-load roundtrip
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+    "dcp_world_size, pcp_world_size",
+    [
+        (2, 1),
+        (1, 2),
+    ],
+)
+def test_cp_eager_store_and_load_roundtrip(
+    dcp_world_size: int, pcp_world_size: int
+) -> None:
+    """With CP enabled, store blocks to CPU and reload them for a new request
+    with matching tokens.  Verifies that hash matching and transfer-pair
+    construction work with the virtual block size."""
+    fix = _make_cp_scheduler(
+        dcp_world_size=dcp_world_size,
+        pcp_world_size=pcp_world_size,
+        num_cpu_blocks=8,
+        num_gpu_blocks=16,
+        lazy=False,
+    )
+    sched = fix.scheduler
+    cp = dcp_world_size * pcp_world_size
+    vbs = BLOCK_SIZE * cp
+
+    num_blocks = 2
+    req = _make_cp_request(num_blocks, vbs)
+
+    # Allocate GPU blocks and register hashes
+    gpu_blocks = _allocate_cp_gpu_blocks(fix.gpu_block_pool, req, num_blocks, vbs)
+    kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,))
+    req.num_computed_tokens = num_blocks * vbs
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    block_ids = kv_blocks.get_block_ids()
+    sched_out = make_scheduler_output(
+        {req.request_id: num_blocks * vbs},
+        new_reqs={req.request_id: block_ids},
+    )
+
+    meta = sched.build_connector_meta(sched_out)
+    assert meta.store_event >= 0, "Expected a store event"
+    assert len(meta.store_gpu_blocks) == num_blocks
+    assert len(meta.store_cpu_blocks) == num_blocks
+    simulate_store_completion(sched, meta.store_event)
+
+    # New request with same tokens — should get a full CPU cache hit.
+    req2 = Request(
+        request_id="req-cp-load",
+        prompt_token_ids=req.prompt_token_ids,
+        sampling_params=req.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=req._block_hasher,
+    )
+
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens == num_blocks * vbs
+    assert is_async is True
+
+    # Allocate fresh GPU blocks for the load.
+    gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
+    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
+    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
+
+    sched_out2 = make_scheduler_output(
+        {req2.request_id: 1},
+        new_reqs={req2.request_id: kv_blocks2.get_block_ids()},
+    )
+    meta2 = sched.build_connector_meta(sched_out2)
+    assert meta2.load_event >= 0, "Expected a load event"
+    assert len(meta2.load_gpu_blocks) == num_blocks
+    assert len(meta2.load_cpu_blocks) == num_blocks
+
+
+# ---------------------------------------------------------------------------
+# Test 17: CP lazy target blocks are scaled correctly
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize("cp_world_size", [1, 2, 4])
+def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None:
+    """_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1
+    because each virtual block covers more tokens."""
+    kv_cache_config = _make_kv_cache_config(num_blocks=16)
+    max_batched = 64
+
+    target_base = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
+        kv_cache_config, max_batched, cp_world_size=1
+    )
+    target_cp = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
+        kv_cache_config, max_batched, cp_world_size=cp_world_size
+    )
+
+    if cp_world_size == 1:
+        assert target_cp == target_base
+    else:
+        assert target_cp < target_base, (
+            f"cp_world_size={cp_world_size}: target_cp={target_cp} should be "
+            f"less than target_base={target_base}"
+        )

tests/v1/streaming_input/test_gpu_model_runner_streaming.py

@@ -35,7 +35,6 @@ def mock_model_runner_with_input_batch():
         max_model_len=1024,
         max_num_batched_tokens=1024,
         device="cpu",
-        pin_memory=False,
         vocab_size=32000,
         block_sizes=[16],
         kernel_block_sizes=[16],

tests/v1/worker/test_gpu_input_batch.py

@@ -10,7 +10,6 @@ import torch
 
 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
-from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.utils.torch_utils import make_tensor_with_pad
 from vllm.v1.pool.metadata import PoolingMetadata
 from vllm.v1.sample.logits_processor import LogitsProcessors
@@ -236,7 +235,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
         max_model_len=1024,
         max_num_batched_tokens=1024,
         device=torch.device(device),
-        pin_memory=is_pin_memory_available(),
         vocab_size=1024,
         block_sizes=[1],
         kernel_block_sizes=[1],
@@ -331,7 +329,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
         max_model_len=1024,
         max_num_batched_tokens=1024,
         device=torch.device(device),
-        pin_memory=is_pin_memory_available(),
         vocab_size=1024,
         block_sizes=[1],
         kernel_block_sizes=[1],
@@ -341,7 +338,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
         max_model_len=1024,
         max_num_batched_tokens=1024,
         device=torch.device(device),
-        pin_memory=is_pin_memory_available(),
         vocab_size=1024,
         block_sizes=[1],
         kernel_block_sizes=[1],
@@ -410,7 +406,6 @@ def test_pooling_prompt_lens_not_aliased(device: str):
         max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
         max_num_batched_tokens=batch_size * (MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS),
         device=torch.device(device),
-        pin_memory=is_pin_memory_available(),
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
         kernel_block_sizes=[16],
@@ -459,7 +454,6 @@ def test_pooling_metadata_token_id_buffers(
         max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
         max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
         device=torch.device("cpu"),
-        pin_memory=False,
         vocab_size=VOCAB_SIZE,
         block_sizes=[16],
         kernel_block_sizes=[16],

tests/v1/worker/test_gpu_model_runner.py

@@ -85,7 +85,6 @@ def initialize_kv_cache(runner: GPUModelRunner):
         max_model_len=runner.max_model_len,
         max_num_batched_tokens=runner.max_num_tokens,
         device=runner.device,
-        pin_memory=runner.pin_memory,
         vocab_size=runner.model_config.get_vocab_size(),
         block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
         kernel_block_sizes=[
@@ -1405,7 +1404,6 @@ def test_input_batch_with_kernel_block_sizes():
     max_model_len = 512
     max_num_batched_tokens = 512
     device = torch.device(DEVICE_TYPE)
-    pin_memory = False
     vocab_size = 50272
 
     # Test with different kernel block sizes
@@ -1417,7 +1415,6 @@ def test_input_batch_with_kernel_block_sizes():
         max_model_len=max_model_len,
         max_num_batched_tokens=max_num_batched_tokens,
         device=device,
-        pin_memory=pin_memory,
         vocab_size=vocab_size,
         block_sizes=block_sizes,
         kernel_block_sizes=kernel_block_sizes,
@@ -1478,7 +1475,6 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
         max_model_len=runner.max_model_len,
         max_num_batched_tokens=runner.max_num_tokens,
         device=runner.device,
-        pin_memory=runner.pin_memory,
         vocab_size=runner.model_config.get_vocab_size(),
         block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
         kernel_block_sizes=[16],

vllm/device_allocator/cumem.py

@@ -18,8 +18,8 @@ import torch
 
 from vllm.device_allocator import AllocationData, HandleType
 from vllm.logger import init_logger
-from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.utils.system_utils import find_loaded_library
+from vllm.utils.torch_utils import PIN_MEMORY
 
 logger = init_logger(__name__)
 
@@ -196,7 +196,7 @@ class CuMemAllocator:
                     size_in_bytes,
                     dtype=torch.uint8,
                     device="cpu",
-                    pin_memory=is_pin_memory_available(),
+                    pin_memory=PIN_MEMORY,
                 )
                 cpu_ptr = cpu_backup_tensor.data_ptr()
                 libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes)

vllm/device_allocator/xpumem.py

@@ -11,7 +11,7 @@ import torch
 
 from vllm.device_allocator import AllocationData, HandleType
 from vllm.logger import init_logger
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 
 logger = init_logger(__name__)
 
@@ -188,7 +188,7 @@ class XpuMemAllocator:
                 size_in_bytes,
                 dtype=torch.uint8,
                 device="cpu",
-                pin_memory=is_pin_memory_available(),
+                pin_memory=PIN_MEMORY,
             )
             cpu_ptr = cpu_backup_tensor.data_ptr()
             _xpu_memcpy_sync(

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

@@ -2,13 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """External-store cache-hit coordinator for MooncakeStoreConnector."""
 
+from collections.abc import Sequence
 from typing import cast
 
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
+    chunk_hashes_for_block_size,
+)
 from vllm.v1.core.block_pool import BlockPool
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
-    BlockHashList,
-    BlockHashListWithBlockSize,
     KVCacheBlock,
 )
 from vllm.v1.core.single_type_kv_cache_manager import (
@@ -120,7 +122,7 @@ class MooncakeStoreCoordinator:
 
     def find_longest_cache_hit(
         self,
-        block_hashes: list[BlockHash],
+        block_hashes: Sequence[BlockHash],
         max_length: int,
         cached_block_pool: ExternalCachedBlockPool,
         *,
@@ -147,7 +149,7 @@ class MooncakeStoreCoordinator:
 
     def load_mask(
         self,
-        block_hashes: list[BlockHash],
+        block_hashes: Sequence[BlockHash],
         token_len: int,
     ) -> tuple[list[bool], ...]:
         """Per-group load masks: ``mask[g][i]`` is True iff group ``g``'s
@@ -236,17 +238,15 @@ class MooncakeStoreCoordinator:
         return tuple(masks)
 
     def block_hashes_for_spec(
-        self, block_hashes: list[BlockHash], spec: KVCacheSpec
-    ) -> BlockHashList:
-        if spec.block_size == self.hash_block_size:
-            return block_hashes
-        return BlockHashListWithBlockSize(
+        self, block_hashes: Sequence[BlockHash], spec: KVCacheSpec
+    ) -> Sequence[BlockHash]:
+        return chunk_hashes_for_block_size(
             block_hashes, self.hash_block_size, spec.block_size
         )
 
     def _find_hit_blocks(
         self,
-        block_hashes: list[BlockHash],
+        block_hashes: Sequence[BlockHash],
         max_length: int,
         cached_block_pool: ExternalCachedBlockPool,
         *,
@@ -264,7 +264,7 @@ class MooncakeStoreCoordinator:
             spec, group_ids, manager_cls = self.attention_groups[0]
             hashes = self.block_hashes_for_spec(block_hashes, spec)
             hit_blocks = manager_cls.find_longest_cache_hit(
-                block_hashes=hashes,
+                block_hashes=hashes,  # type: ignore[arg-type]
                 max_length=max_length,
                 kv_cache_group_ids=group_ids,
                 block_pool=cast(BlockPool, cached_block_pool),
@@ -304,7 +304,7 @@ class MooncakeStoreCoordinator:
                     _max_length = min(curr_hit_length + spec.block_size, max_length)
                 hashes = self.block_hashes_for_spec(block_hashes, spec)
                 hit_blocks = manager_cls.find_longest_cache_hit(
-                    block_hashes=hashes,
+                    block_hashes=hashes,  # type: ignore[arg-type]
                     max_length=_max_length,
                     kv_cache_group_ids=group_ids,
                     block_pool=cast(BlockPool, cached_block_pool),

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

@@ -5,8 +5,9 @@
 # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
 """Data classes for MooncakeStoreConnector."""
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
+from typing import cast
 
 import torch
 
@@ -23,6 +24,77 @@ from vllm.v1.core.kv_cache_utils import (
 logger = init_logger(__name__)
 
 
+class BlobBlockHashes(Sequence[BlockHash]):
+    """Lazy view over a flat buffer of fixed-size block hashes to avoid the overhead
+    of materializing all hashes upfront.
+    """
+
+    def __init__(self, blob: memoryview, hash_len: int):
+        self._blob = blob
+        self._hash_len = hash_len
+        self._n = len(blob) // hash_len if hash_len else 0
+
+    def __len__(self) -> int:
+        return self._n
+
+    def __getitem__(self, idx):
+        if isinstance(idx, slice):
+            return [self[i] for i in range(*idx.indices(self._n))]
+        if idx < 0:
+            idx += self._n
+        if not 0 <= idx < self._n:
+            raise IndexError(idx)
+        off = idx * self._hash_len
+        return BlockHash(self._blob[off : off + self._hash_len])
+
+
+class _CompactChunkHashList(BlockHashListWithBlockSize):
+    """View that keys each ``block_size`` chunk by the last constituent
+    ``hash_block_size`` hash instead of concatenating all of them.
+
+    The engine chains block hashes (each hash folds in the previous one), so the
+    final sub-block hash of a chunk already uniquely identifies the whole chunk
+    and its prefix. Using it keeps a Mooncake key at a single hash digest
+    regardless of the ``block_size`` / ``hash_block_size`` ratio, instead of
+    growing the key linearly with it (e.g. 64x for ``block_size=256``,
+    ``hash_block_size=4``).
+    """
+
+    def __init__(
+        self,
+        block_hashes: Sequence[BlockHash],
+        hash_block_size: int,
+        target_block_size: int,
+    ):
+        # Accept any indexable sequence (e.g. the lazy ``BlobBlockHashes``), not
+        # just ``list``; the base only indexes/sizes it.
+        assert target_block_size % hash_block_size == 0
+        self.block_hashes = block_hashes  # type: ignore[assignment]
+        self.scale_factor = target_block_size // hash_block_size
+
+    def _get_value_at(self, idx: int) -> BlockHash:
+        return self.block_hashes[idx * self.scale_factor + self.scale_factor - 1]
+
+
+def chunk_hashes_for_block_size(
+    block_hashes: Sequence[BlockHash],
+    hash_block_size: int,
+    block_size: int,
+) -> Sequence[BlockHash]:
+    """Map ``hash_block_size``-granular block hashes to one compact hash per
+    ``block_size`` chunk (the chunk's last sub-hash). Returns ``block_hashes``
+    unchanged when the two sizes are equal.
+    """
+    if block_size == hash_block_size:
+        return block_hashes
+    # Structurally a Sequence[BlockHash] (indexable + sized); the base class
+    # just isn't declared as one.
+    return cast(
+        "Sequence[BlockHash]",
+        _CompactChunkHashList(block_hashes, hash_block_size, block_size),
+    )
+
+
 @dataclass
 class KeyMetadata:
     """Metadata for constructing pool keys."""
@@ -138,18 +210,15 @@ class ChunkedTokenDatabase:
         Args:
             token_len: Total number of tokens.
             block_hashes: Block hashes computed at ``hash_block_size`` granularity.
-                When ``block_size > hash_block_size`` consecutive hashes are merged
-                up to the group's ``block_size`` via ``BlockHashListWithBlockSize``.
+                When ``block_size > hash_block_size`` each group's ``block_size`` chunk
+                is keyed by its last sub-hash via ``chunk_hashes_for_block_size``.
             mask_num: Number of tokens to skip from the beginning.
         """
         if not block_hashes:
             return
-        if self.block_size == self.hash_block_size:
-            chunk_hashes: Iterable[BlockHash] = block_hashes
-        else:
-            chunk_hashes = BlockHashListWithBlockSize(
-                block_hashes, self.hash_block_size, self.block_size
-            )
+        chunk_hashes: Iterable[BlockHash] = chunk_hashes_for_block_size(
+            block_hashes, self.hash_block_size, self.block_size
+        )
         for chunk_id, h in enumerate(chunk_hashes):
             start_idx = chunk_id * self.block_size
             if start_idx >= token_len:

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/protocol.py

@@ -11,7 +11,10 @@ Wire format (REQ/REP over IPC):
 
       msg_type == LOOKUP_MSG:
           frame 1: token_len (u32 big-endian, 4 bytes)
-          frame 2..n: msgpack-encoded list[str] of block-hash hex digests
+          frame 2: hash_len (u16 big-endian, 2 bytes) — byte length of each
+                   fixed-size block hash (0 when there are no hashes)
+          frame 3: raw block hashes concatenated back-to-back (each hash_len
+                   bytes); the server splits on hash_len
         Response: [hit_count: u32 big-endian, 4 bytes]
 
       msg_type == RESET_MSG:

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

@@ -18,7 +18,7 @@ import socket
 import threading
 import time
 from collections import defaultdict
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass
 from typing import Any, Literal, TypeVar
@@ -45,6 +45,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp
     MooncakeStoreCoordinator,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+    BlobBlockHashes,
     ChunkedTokenDatabase,
     KeyMetadata,
     MooncakeStoreConnectorMetadata,
@@ -65,7 +66,6 @@ from vllm.v1.core.kv_cache_utils import (
     resolve_kv_cache_block_sizes,
 )
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec
-from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
 
 from .metrics import MooncakeStoreConnectorStats
 
@@ -1372,7 +1372,7 @@ class MooncakeStoreWorker:
 
         return finished_sending
 
-    def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
+    def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int:
         """Check how many prefix tokens exist in the store.
 
         Checks across all TP ranks and PP ranks.
@@ -1392,6 +1392,11 @@ class MooncakeStoreWorker:
             group_hashes = self.coord.block_hashes_for_spec(
                 block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
             )
+            metadata_templates = [
+                dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)
+                for tp in range(tp_count)
+                for pp in range(self.pp_size)
+            ]
             for chunk_id, h in enumerate(group_hashes):
                 start_idx = chunk_id * spec_block_size
                 if start_idx >= token_len:
@@ -1400,11 +1405,11 @@ class MooncakeStoreWorker:
                     chunk_id >= len(lookup_mask) or not lookup_mask[chunk_id]
                 ):
                     continue
-                for tp in range(tp_count):
-                    for pp in range(self.pp_size):
-                        md = dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)
-                        candidate_keys.append(PoolKey(md, h.hex()).to_string())
-                        candidate_meta.append((g_idx, bytes(h)))
+                h_hex = h.hex()
+                h_bytes = bytes(h)
+                for md in metadata_templates:
+                    candidate_keys.append(PoolKey(md, h_hex).to_string())
+                    candidate_meta.append((g_idx, h_bytes))
 
         if not candidate_keys:
             return 0
@@ -1483,7 +1488,6 @@ class LookupKeyServer:
         store_worker: MooncakeStoreWorker,
         vllm_config: VllmConfig,
     ):
-        self.decoder = MsgpackDecoder()
         self.ctx = zmq.Context()  # type: ignore[attr-defined]
         socket_path = get_zmq_rpc_path_lookup(vllm_config)
         self._ipc_path = socket_path.removeprefix("ipc://")
@@ -1506,9 +1510,9 @@ class LookupKeyServer:
 
                 if msg_type == LOOKUP_MSG:
                     token_len = int.from_bytes(all_frames[1], byteorder="big")
-                    hash_frames = all_frames[2:]
-                    hashes_str = self.decoder.decode(hash_frames)
-                    block_hashes = [BlockHash(bytes.fromhex(s)) for s in hashes_str]
+                    hash_len = int.from_bytes(all_frames[2], byteorder="big")
+                    blob = all_frames[3].buffer
+                    block_hashes = BlobBlockHashes(blob, hash_len)
                     result = self.store_worker.lookup(token_len, block_hashes)
                     self.socket.send(result.to_bytes(4, "big"))
 
@@ -1557,7 +1561,6 @@ class LookupKeyClient:
     """
 
     def __init__(self, vllm_config: VllmConfig):
-        self.encoder = MsgpackEncoder()
         self.ctx = zmq.Context()  # type: ignore[attr-defined]
         socket_path = get_zmq_rpc_path_lookup(vllm_config)
         self.socket = make_zmq_socket(
@@ -1574,14 +1577,16 @@ class LookupKeyClient:
         self.futures: dict[str, Future[int]] = {}
 
     def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
-        hash_strs = [h.hex() for h in block_hashes]
-        hash_frames = self.encoder.encode(hash_strs)
-        token_len_bytes = token_len.to_bytes(4, byteorder="big")
-        all_frames = [LOOKUP_MSG, token_len_bytes] + list(hash_frames)
+        hash_len = len(block_hashes[0]) if block_hashes else 0
+        all_frames = (
+            LOOKUP_MSG,
+            token_len.to_bytes(4, byteorder="big"),
+            hash_len.to_bytes(2, byteorder="big"),
+            b"".join(block_hashes),
+        )
         self.socket.send_multipart(all_frames, copy=False)
         resp = self.socket.recv()
-        result = int.from_bytes(resp, "big")
-        return result
+        return int.from_bytes(resp, "big")
 
     def lookup(
         self,

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

@@ -50,7 +50,8 @@ class OffloadingConnectorWorker:
     def register_kv_caches(
         self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
     ):
-        num_blocks = self.spec.kv_cache_config.num_blocks
+        kv_cache_config = self.spec.kv_cache_config
+        num_blocks = kv_cache_config.num_blocks
 
         # layer_name -> (num_blocks, page_size_bytes) tensor
         tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
@@ -58,7 +59,7 @@ class OffloadingConnectorWorker:
         unpadded_page_size_bytes: dict[str, int] = {}
         # layer_name -> size of page in bytes
         page_size_bytes: dict[str, int] = {}
-        for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
             group_layer_names = kv_cache_group.layer_names
             group_kv_cache_spec = kv_cache_group.kv_cache_spec
             if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -122,9 +123,35 @@ class OffloadingConnectorWorker:
                 else:
                     raise NotImplementedError
 
+        packed_kv_cache_tensor = next(
+            (t for t in kv_cache_config.kv_cache_tensors if t.block_stride), None
+        )
+        is_dsv4 = all(
+            isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
+            for group in kv_cache_config.kv_cache_groups
+        )
+        if packed_kv_cache_tensor is not None and not is_dsv4:
+            (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]]
+            block_stride = tensor.stride(0)
+            packed_tensor = tensor.as_strided(
+                (num_blocks, block_stride),
+                (block_stride, 1),
+                storage_offset=0,
+            )
+            self._register_handlers(
+                CanonicalKVCaches(
+                    [CanonicalKVCacheTensor(packed_tensor, block_stride)],
+                    [
+                        [CanonicalKVCacheRef(0, block_stride)]
+                        for _ in kv_cache_config.kv_cache_groups
+                    ],
+                )
+            )
+            return
+
         block_tensors: list[CanonicalKVCacheTensor] = []
         block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
-        for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors:
+        for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
             # Filter to layers that were actually processed above.
             # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for
             # every (tuple_idx, page_size) slot; slots where no group has a
@@ -166,7 +193,7 @@ class OffloadingConnectorWorker:
                     )
 
         group_data_refs: list[list[CanonicalKVCacheRef]] = []
-        for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
             group_refs: list[CanonicalKVCacheRef] = []
             for layer_name in kv_cache_group.layer_names:
                 group_refs += block_data_refs[layer_name]

vllm/entrypoints/anthropic/serving.py

@@ -43,6 +43,7 @@ from vllm.entrypoints.openai.engine.protocol import (
     JsonSchemaResponseFormat,
     ResponseFormat,
     StreamOptions,
+    UsageInfo,
 )
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.serve.utils.api_utils import sanitize_message
@@ -54,6 +55,49 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _get_cached_tokens(usage: UsageInfo | None) -> int | None:
+    """Extract cached token count from OpenAI UsageInfo."""
+    if usage is None or usage.prompt_tokens_details is None:
+        return None
+    return usage.prompt_tokens_details.cached_tokens
+
+
+def _build_anthropic_usage(
+    prompt_tokens: int,
+    completion_tokens: int | None,
+    usage: UsageInfo | None,
+) -> AnthropicUsage:
+    """Build an AnthropicUsage from OpenAI-style token counts.
+
+    Anthropic defines ``total_input == input_tokens + cache_read +
+    cache_creation``.  vLLM's ``prompt_tokens`` is the total, so
+    ``input_tokens = prompt_tokens - cached_tokens``.
+
+    OpenAI usage only exposes ``cached_tokens`` (hits); there is no
+    cache-creation analog, so ``cache_creation_input_tokens`` is ``0``
+    when cache info is present.  When cache info is absent (e.g.
+    ``--enable-prompt-tokens-details`` off, or a streaming chunk that
+    hasn't carried it yet), cache fields are left **unset** so
+    ``exclude_unset=True`` serialization omits them entirely.
+
+    ``completion_tokens`` follows ``UsageInfo`` and may be ``None`` on
+    intermediate stream chunks; we coerce to ``0`` for the wire format.
+    """
+    output_tokens = completion_tokens or 0
+    cached = _get_cached_tokens(usage)
+    if cached is not None:
+        return AnthropicUsage(
+            input_tokens=prompt_tokens - cached,
+            output_tokens=output_tokens,
+            cache_read_input_tokens=cached,
+            cache_creation_input_tokens=0,
+        )
+    return AnthropicUsage(
+        input_tokens=prompt_tokens,
+        output_tokens=output_tokens,
+    )
+
+
 def wrap_data_with_event(data: str, event: str):
     return f"event: {event}\ndata: {data}\n\n"
 
@@ -582,9 +626,10 @@ class AnthropicServingMessages(OpenAIServingChat):
             id=generator.id,
             content=[],
             model=generator.model,
-            usage=AnthropicUsage(
-                input_tokens=generator.usage.prompt_tokens,
-                output_tokens=generator.usage.completion_tokens,
+            usage=_build_anthropic_usage(
+                generator.usage.prompt_tokens,
+                generator.usage.completion_tokens,
+                generator.usage,
             ),
             kv_transfer_params=generator.kv_transfer_params,
         )
@@ -765,11 +810,12 @@ class AnthropicServingMessages(OpenAIServingChat):
                                     model=origin_chunk.model,
                                     stop_reason=None,
                                     stop_sequence=None,
-                                    usage=AnthropicUsage(
-                                        input_tokens=origin_chunk.usage.prompt_tokens
+                                    usage=_build_anthropic_usage(
+                                        origin_chunk.usage.prompt_tokens
                                         if origin_chunk.usage
                                         else 0,
-                                        output_tokens=0,
+                                        0,
+                                        origin_chunk.usage,
                                     ),
                                 ),
                             )
@@ -788,13 +834,14 @@ class AnthropicServingMessages(OpenAIServingChat):
                             chunk = AnthropicStreamEvent(
                                 type="message_delta",
                                 delta=AnthropicDelta(stop_reason=stop_reason),
-                                usage=AnthropicUsage(
-                                    input_tokens=origin_chunk.usage.prompt_tokens
+                                usage=_build_anthropic_usage(
+                                    origin_chunk.usage.prompt_tokens
                                     if origin_chunk.usage
                                     else 0,
-                                    output_tokens=origin_chunk.usage.completion_tokens
+                                    origin_chunk.usage.completion_tokens
                                     if origin_chunk.usage
                                     else 0,
+                                    origin_chunk.usage,
                                 ),
                             )
                             data = chunk.model_dump_json(exclude_unset=True)

vllm/entrypoints/openai/api_server.py

@@ -455,7 +455,7 @@ async def init_render_app_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
-        reasoning_parser=args.structured_outputs_config.reasoning_parser,
+        reasoning_parser=args.reasoning_parser,
         default_chat_template_kwargs=args.default_chat_template_kwargs,
         log_error_stack=args.log_error_stack,
     )

vllm/entrypoints/serve/disagg/protocol.py

@@ -219,10 +219,14 @@ class GenerateResponse(BaseModel):
 
 
 class DerenderChatRequest(BaseModel):
-    """Request for the /v1/chat/completions/derender endpoint.
+    """Request for the /v1/chat/completions/derender endpoint (non-streaming).
 
-    Wraps a GenerateResponse and caller-supplied metadata needed to produce
-    a fully-formed ChatCompletionResponse without a GPU.
+    Wraps a complete GenerateResponse and caller-supplied metadata needed to
+    produce a fully-formed ChatCompletionResponse without a GPU.
+
+    Streaming derender would require a separate endpoint design with
+    incremental token delivery, ``OutputProcessor``-based detokenization,
+    and ``parser.parse_delta()`` instead of ``parser.parse()``.
     """
 
     model: str
@@ -244,7 +248,7 @@ class DerenderChatRequest(BaseModel):
 
 
 class DerenderCompletionRequest(BaseModel):
-    """Request for the /v1/completions/derender endpoint.
+    """Request for the /v1/completions/derender endpoint (non-streaming).
 
     Parallel to DerenderChatRequest but handles the multi-prompt completions
     case: one GenerateResponse per prompt, mirroring the list[GenerateRequest]

vllm/entrypoints/serve/render/serving.py

@@ -27,6 +27,7 @@ from vllm.entrypoints.openai.completion.protocol import (
 )
 from vllm.entrypoints.openai.engine.protocol import (
     ErrorResponse,
+    ToolCall,
     UsageInfo,
 )
 from vllm.entrypoints.openai.engine.serving import resolve_token_id_placeholder
@@ -43,7 +44,6 @@ from vllm.entrypoints.serve.disagg.protocol import (
     DerenderChatRequest,
     DerenderCompletionRequest,
     GenerateRequest,
-    GenerateResponseChoice,
     MultiModalFeatures,
     PlaceholderRangeInfo,
 )
@@ -76,21 +76,83 @@ from vllm.utils.mistral import mt as _mt
 logger = init_logger(__name__)
 
 
+def _parse_token_id_placeholder(token: str) -> int | None:
+    """Extract token ID from a 'token_id:N' placeholder string."""
+    if not token.startswith("token_id:"):
+        return None
+    try:
+        return int(token[len("token_id:") :])
+    except ValueError:
+        return None
+
+
+def _correct_decoded_token(
+    token_id: int, context_token_ids: list[int], tokenizer: TokenizerLike
+) -> str:
+    """Use preceding tokens as context to fix U+FFFD from byte-fallback.
+
+    Mirrors LogprobsProcessor._correct_decoded_token in v1/engine/logprobs.py.
+    """
+    max_ctx = min(len(context_token_ids), 4)
+
+    for num_ctx in range(1, max_ctx + 1):
+        context = context_token_ids[-num_ctx:]
+        full_decoded = tokenizer.decode(context + [token_id])
+
+        if full_decoded.endswith("�"):
+            continue
+
+        clean_end = len(context)
+        for j in range(len(context) - 1, -1, -1):
+            if tokenizer.decode([context[j]]).endswith("�"):
+                clean_end = j
+            else:
+                break
+
+        clean_prefix = tokenizer.decode(context[:clean_end]) if clean_end > 0 else ""
+
+        if full_decoded.startswith(clean_prefix):
+            return full_decoded[len(clean_prefix) :]
+
+        common_len = 0
+        for a, b in zip(clean_prefix, full_decoded):
+            if a != b:
+                break
+            common_len += 1
+        return full_decoded[common_len:]
+
+    return ""
+
+
 def _resolve_logprobs(
     logprobs: ChatCompletionLogProbs, tokenizer: TokenizerLike
 ) -> ChatCompletionLogProbs:
-    """Resolve all token_id:N placeholders in a ChatCompletionLogProbs object."""
+    """Resolve token_id:N placeholders in a ChatCompletionLogProbs object."""
     if logprobs.content is None:
         return logprobs
+
+    context_token_ids: list[int] = []
     resolved_content = []
+
     for entry in logprobs.content:
         token_str, token_bytes = resolve_token_id_placeholder(entry.token, tokenizer)
+        sampled_id = _parse_token_id_placeholder(entry.token)
+
+        if token_str.endswith("�") and sampled_id is not None:
+            token_str = _correct_decoded_token(sampled_id, context_token_ids, tokenizer)
+            token_bytes = list(token_str.encode("utf-8"))
+
         resolved_top = []
         for top in entry.top_logprobs:
             top_str, top_bytes = resolve_token_id_placeholder(top.token, tokenizer)
+            top_id = _parse_token_id_placeholder(top.token)
+            if top_str.endswith("�") and top_id is not None:
+                top_str = _correct_decoded_token(top_id, context_token_ids, tokenizer)
+                top_bytes = list(top_str.encode("utf-8"))
             resolved_top.append(
                 top.model_copy(update={"token": top_str, "bytes": top_bytes})
             )
+
         resolved_content.append(
             entry.model_copy(
                 update={
@@ -100,6 +162,10 @@ def _resolve_logprobs(
                 }
             )
         )
+
+        if sampled_id is not None:
+            context_token_ids.append(sampled_id)
+
     return ChatCompletionLogProbs(content=resolved_content)
 
 
@@ -136,30 +202,6 @@ def _convert_chat_logprobs_to_completion_logprobs(
     )
 
 
-def _build_chat_choice(
-    choice: GenerateResponseChoice, tokenizer: TokenizerLike
-) -> ChatCompletionResponseChoice:
-    """Detokenize and resolve logprobs for a single GenerateResponseChoice.
-
-    Raises:
-        ValueError: if choice.token_ids is empty or None.
-    """
-    if not choice.token_ids:
-        raise ValueError(f"choice {choice.index} has empty or null token_ids")
-    decoded_text = tokenizer.decode(choice.token_ids, skip_special_tokens=True)
-    resolved_logprobs = (
-        _resolve_logprobs(choice.logprobs, tokenizer)
-        if choice.logprobs is not None
-        else None
-    )
-    return ChatCompletionResponseChoice(
-        index=choice.index,
-        message=ChatMessage(role="assistant", content=decoded_text),
-        logprobs=resolved_logprobs,
-        finish_reason=choice.finish_reason,
-    )
-
-
 class OpenAIServingRender:
     def __init__(
         self,
@@ -536,9 +578,12 @@ class OpenAIServingRender:
     ) -> ChatCompletionResponse | ErrorResponse:
         """Postprocess a GenerateResponse into a ChatCompletionResponse.
 
-        This is the symmetric inverse of render_chat_request: it detokenizes
-        output token IDs, resolves token_id:N logprob placeholders, and
-        formats the result as an OpenAI-compatible chat completion response.
+        Non-streaming only: expects the complete GenerateResponse with all
+        token IDs present.  Uses ``parser.parse()`` for one-shot extraction.
+
+        When ``request.chat_request`` is provided, the parser splits the
+        output into (reasoning, content, tool_calls).  Otherwise falls
+        back to plain detokenization.
         """
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
@@ -546,11 +591,89 @@ class OpenAIServingRender:
 
         tokenizer = self.renderer.get_tokenizer()
         gen = request.generate_response
+        chat_request = request.chat_request
         choices: list[ChatCompletionResponseChoice] = []
 
         try:
             for choice in gen.choices:
-                choices.append(_build_chat_choice(choice, tokenizer))
+                if not choice.token_ids:
+                    raise ValueError(
+                        f"choice {choice.index} has empty or null token_ids"
+                    )
+
+                resolved_logprobs = (
+                    _resolve_logprobs(choice.logprobs, tokenizer)
+                    if choice.logprobs is not None
+                    else None
+                )
+
+                if self.parser is not None and chat_request is not None:
+                    # Parser path: decode with special tokens preserved
+                    # so the parser can see markers like </think>,
+                    # <tool_call>, or Harmony channel tokens.
+                    decoded_text = tokenizer.decode(
+                        choice.token_ids, skip_special_tokens=False
+                    )
+
+                    chat_template_kwargs: dict[str, Any] = {}
+                    if not self.use_harmony:
+                        chat_template_kwargs = (
+                            chat_request.build_chat_params(
+                                self.chat_template,
+                                self.chat_template_content_format,
+                            )
+                            .with_defaults(self.default_chat_template_kwargs)
+                            .chat_template_kwargs
+                        )
+
+                    parser = self.parser(
+                        tokenizer,
+                        chat_request.tools,
+                        chat_template_kwargs=chat_template_kwargs,
+                    )
+                    reasoning, content, tool_calls = parser.parse(
+                        decoded_text,
+                        chat_request,
+                        enable_auto_tools=self.enable_auto_tools,
+                        model_output_token_ids=choice.token_ids,
+                    )
+
+                    if not getattr(chat_request, "include_reasoning", True):
+                        reasoning = None
+
+                    tc_items = (
+                        [
+                            ToolCall(
+                                id=random_uuid(),
+                                function=tc,
+                            )
+                            for tc in tool_calls
+                        ]
+                        if tool_calls
+                        else []
+                    )
+
+                    message = ChatMessage(
+                        role="assistant",
+                        reasoning=reasoning,
+                        content=content,
+                        tool_calls=tc_items,
+                    )
+                else:
+                    # No parser: plain detokenization.
+                    decoded_text = tokenizer.decode(
+                        choice.token_ids, skip_special_tokens=True
+                    )
+                    message = ChatMessage(role="assistant", content=decoded_text)
+
+                choices.append(
+                    ChatCompletionResponseChoice(
+                        index=choice.index,
+                        message=message,
+                        logprobs=resolved_logprobs,
+                        finish_reason=choice.finish_reason,
+                    )
+                )
         except ValueError as exc:
             return self.create_error_response(str(exc))
 
@@ -587,8 +710,9 @@ class OpenAIServingRender:
     ) -> CompletionResponse | ErrorResponse:
         """Postprocess a list of GenerateResponses into a CompletionResponse.
 
-        Mirrors the multi-prompt completions case: one GenerateResponse per
-        prompt, parallel to the list[GenerateRequest] from /v1/completions/render.
+        Non-streaming only.  Mirrors the multi-prompt completions case: one
+        GenerateResponse per prompt, parallel to the list[GenerateRequest]
+        from /v1/completions/render.
         """
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:

vllm/envs.py

@@ -209,6 +209,7 @@ if TYPE_CHECKING:
     VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
     VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5
     VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
+    VLLM_USE_PACKED_HMA_KV_CACHE: bool = False
     VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
     VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
     VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
@@ -485,7 +486,7 @@ def get_vllm_port() -> int | None:
             raise ValueError(
                 f"VLLM_PORT '{port}' appears to be a URI. "
                 "This may be caused by a Kubernetes service discovery issue,"
-                "check the warning in: https://docs.vllm.ai/en/stable/serving/env_vars.html"
+                "check the warning in: https://docs.vllm.ai/en/latest/configuration/env_vars.html"
             ) from None
         raise ValueError(f"VLLM_PORT '{port}' must be a valid integer") from err
 
@@ -1579,9 +1580,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
         os.getenv("VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS", "1")
     ),
     # Enforce function parameter schemas in structural-tag based tool calling.
-    "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: (
-        os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "True").lower() in ("true", "1")
-    ),
+    "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: os.getenv(
+        "VLLM_ENFORCE_STRICT_TOOL_CALLING", "True"
+    ).lower()
+    in ("true", "1"),
     # Control the max chunk bytes (in MB) for the rpc message queue.
     # Object larger than this threshold will be broadcast to worker
     # processes via zmq.
@@ -1607,6 +1609,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_KV_CACHE_LAYOUT": env_with_choices(
         "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
     ),
+    # Opt into packed per-block KV cache allocation for multi-group
+    # attention-only HMA models (e.g. gpt-oss, Gemma 3/4).
+    "VLLM_USE_PACKED_HMA_KV_CACHE": lambda: bool(
+        int(os.getenv("VLLM_USE_PACKED_HMA_KV_CACHE", "0"))
+    ),
     # SSM conv state layout used for Mamba models.
     # - SD: (state_len, dim) — dim contiguous (default)
     # - DS: (dim, state_len) — TP-sharded dim on dim1,

vllm/lora/lora_model.py

@@ -17,7 +17,7 @@ from vllm.lora.utils import (
 )
 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
 from vllm.model_executor.models.utils import WeightsMapper
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 
 logger = init_logger(__name__)
 
@@ -126,7 +126,7 @@ class LoRAModel:
         skip_prefixes: list[str] | None = None,
     ) -> "LoRAModel":
         """Create a LoRAModel from a dictionary of tensors."""
-        pin_memory = str(device) == "cpu" and is_pin_memory_available()
+        pin_memory = str(device) == "cpu" and PIN_MEMORY
         loras: dict[str, LoRALayerWeights] = {}
         for tensor_name, tensor in tensors.items():
             if is_base_embedding_weights(tensor_name):

vllm/lora/lora_weights.py

@@ -7,7 +7,7 @@ import torch
 import torch.types
 
 from vllm.lora.peft_helper import PEFTHelper
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 
 
 class LoRALayerWeights:
@@ -79,7 +79,7 @@ class LoRALayerWeights:
         dtype: torch.dtype,
         device: torch.types.Device,
     ) -> "LoRALayerWeights":
-        pin_memory = str(device) == "cpu" and is_pin_memory_available()
+        pin_memory = str(device) == "cpu" and PIN_MEMORY
         lora_a = torch.zeros(
             [rank, input_dim], dtype=dtype, device=device, pin_memory=pin_memory
         )

vllm/lora/model_manager.py

@@ -42,7 +42,7 @@ from vllm.model_executor.models.utils import PPMissingLayer
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.encoder_budget import MultiModalBudget
 from vllm.utils.cache import LRUCache
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 
 logger = init_logger(__name__)
 
@@ -801,7 +801,7 @@ class LoRAModelManager:
         # 2. The weight packing above (e.g., pack_moe) may invalidate the
         # pin_memory allocation, so we execute it after packing.
 
-        pin_memory = str(lora_device) == "cpu" and is_pin_memory_available()
+        pin_memory = str(lora_device) == "cpu" and PIN_MEMORY
         if pin_memory:
             for lora in lora_model.loras.values():
                 if isinstance(lora.lora_a, list):

vllm/model_executor/layers/attention/mla_attention.py

@@ -1684,12 +1684,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                 #  [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
                 # Note(simon): this is done in CPU because of downstream's
                 # of `to_list`.
-                chunk_starts = (
+                chunk_starts = torch.empty(
+                    num_chunks, num_prefills, dtype=torch.int32, pin_memory=True
+                ).copy_(
                     torch.arange(num_chunks, dtype=torch.int32)
+                    .multiply_(max_context_chunk)
                     .unsqueeze(1)
-                    .expand(-1, num_prefills)
-                    * max_context_chunk
-                ).pin_memory()
+                )
                 chunk_ends = torch.min(
                     context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk
                 )
@@ -1746,12 +1747,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                         )
                         * self.dcp_local_block_size
                     )
-                    local_chunk_starts = (
+                    local_chunk_starts = torch.empty(
+                        num_chunks, num_prefills, dtype=torch.int32, pin_memory=True
+                    ).copy_(
                         torch.arange(num_chunks, dtype=torch.int32)
+                        .multiply_(padded_local_max_context_chunk_across_ranks)
                         .unsqueeze(1)
-                        .expand(-1, num_prefills)
-                        * padded_local_max_context_chunk_across_ranks
-                    ).pin_memory()
+                    )
                     local_chunk_ends = torch.min(
                         padded_local_context_lens_cpu.unsqueeze(0),
                         local_chunk_starts

vllm/model_executor/layers/attention/mm_encoder_attention.py

@@ -28,6 +28,7 @@ from vllm.utils.flashinfer import (
     is_flashinfer_cudnn_fp8_prefill_attn_supported,
 )
 from vllm.utils.math_utils import round_up
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.ops.vit_attn_wrappers import (
@@ -311,7 +312,7 @@ class MMEncoderAttention(CustomOp):
                 )
                 cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v])
 
-        cu_seqlens = torch.from_numpy(cu_seqlens).to(device, non_blocking=True)
+        cu_seqlens = async_tensor_h2d(cu_seqlens, device=device)
         return cu_seqlens
 
     def __init__(

vllm/model_executor/layers/fused_moe/experts/nvfp4_emulation_moe.py

@@ -21,7 +21,6 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
-from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
 from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
     dequantize_to_dtype,
 )
@@ -135,14 +134,6 @@ class Nvfp4QuantizationEmulationTritonExperts(TritonExperts):
             swizzle=False,
         )
 
-        hidden_states, _ = moe_kernel_quantize_input(
-            A=hidden_states,
-            A_scale=self.quant_config.a1_gscale,
-            quant_dtype="nvfp4",
-            per_act_token_quant=False,
-            quantization_emulation=True,
-        )
-
         # Activation quantization/dequantization is deferred to
         # `moe_kernel_quantize_input` in TritonExperts.apply.
         super().apply(

vllm/model_executor/layers/fused_moe/experts/ocp_mx_emulation_moe.py

@@ -21,7 +21,6 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
-from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
 from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
@@ -155,16 +154,6 @@ class OCP_MXQuantizationEmulationTritonExperts(TritonExperts):
             w2, self.w2_scale_val, hidden_states.dtype
         )
 
-        # Apply activation QDQ if needed by the OCP MX scheme
-        hidden_states, _ = moe_kernel_quantize_input(
-            A=hidden_states,
-            A_scale=None,
-            quant_dtype=self.quant_config.quant_dtype,
-            per_act_token_quant=False,
-            ocp_mx_scheme=self.ocp_mx_scheme,
-            quantization_emulation=True,
-        )
-
         # Activation quantization/dequantization is deferred to
         # `moe_kernel_quantize_input` in TritonExperts.apply.
         super().apply(

vllm/model_executor/layers/fused_moe/experts/triton_moe.py

@@ -245,7 +245,7 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
             lora_unquantized_hidden_states = hidden_states
             hidden_states, a1q_scale = moe_kernel_quantize_input(
                 hidden_states,
-                self.a1_scale,
+                self.a1_scale or self.a1_gscale,
                 self.quant_dtype,
                 self.per_act_token_quant,
                 self.block_shape,

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

@@ -56,35 +56,6 @@ class TrtLlmFp8ExpertsBase:
         self.moe_config = moe_config
         self.quant_config = quant_config
 
-        # Per-expert SwiGLU parameters from quant_config (MXFP8 + Swiglu only).
-        device = torch.accelerator.current_device_index()
-        if quant_config.gemm1_alpha is not None:
-            self.gemm1_alpha = torch.tensor(
-                [quant_config.gemm1_alpha] * self.local_num_experts,
-                dtype=torch.float32,
-                device=device,
-            )
-        else:
-            self.gemm1_alpha = None
-
-        if quant_config.gemm1_beta is not None:
-            self.gemm1_beta = torch.tensor(
-                [quant_config.gemm1_beta] * self.local_num_experts,
-                dtype=torch.float32,
-                device=device,
-            )
-        else:
-            self.gemm1_beta = None
-
-        if quant_config.gemm1_clamp_limit is not None:
-            self.gemm1_clamp_limit = torch.tensor(
-                [quant_config.gemm1_clamp_limit] * self.local_num_experts,
-                dtype=torch.float32,
-                device=device,
-            )
-        else:
-            self.gemm1_clamp_limit = None
-
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
@@ -106,12 +77,8 @@ class TrtLlmFp8ExpertsBase:
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        """Supports SiLU, SwiGLU-OAI (uninterleaved), and RELU^2 non-gated."""
-        return activation in [
-            MoEActivation.SILU,
-            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
-            MoEActivation.RELU2_NO_MUL,
-        ]
+        """Supports only SiLU and RELU^2 non-gated activation."""
+        return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
 
     @staticmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -231,9 +198,6 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             hidden_states_scale=hidden_states_scale,
             gemm1_weights=w1,
             gemm1_weights_scale=self.quant_config.w1_scale,
-            gemm1_alpha=self.gemm1_alpha,
-            gemm1_beta=self.gemm1_beta,
-            gemm1_clamp_limit=self.gemm1_clamp_limit,
             gemm2_weights=w2,
             gemm2_weights_scale=self.quant_config.w2_scale,
             num_experts=global_num_experts,
@@ -363,11 +327,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
         from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout
 
         assert not apply_router_weight_on_input
-        assert activation in [
-            MoEActivation.SILU,
-            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
-            MoEActivation.RELU2_NO_MUL,
-        ]
+        assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
         activation_type = activation_to_flashinfer_int(activation)
         assert self.topk <= global_num_experts
         assert global_num_experts % 4 == 0
@@ -402,9 +362,6 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
             hidden_states_scale=hidden_states_scale,
             gemm1_weights=w1,
             gemm1_weights_scale=self.quant_config.w1_scale,
-            gemm1_alpha=self.gemm1_alpha,
-            gemm1_beta=self.gemm1_beta,
-            gemm1_clamp_limit=self.gemm1_clamp_limit,
             gemm2_weights=w2,
             gemm2_weights_scale=self.quant_config.w2_scale,
             num_experts=global_num_experts,

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

@@ -66,47 +66,16 @@ class TrtLlmNvFp4ExpertsBase:
         else:
             self.g1_scale_c = self.quant_config.a2_gscale.clone()
 
-        # Fall back to moe_config.swiglu_* when quant_config doesn't carry them
-        # (ModelOpt NVFP4 checkpoints store these on moe_config, not quant_config).
-        device = torch.accelerator.current_device_index()
-
-        def _per_expert(val: float | None) -> torch.Tensor | None:
-            if val is None:
-                return None
-            return torch.full(
+        if moe_config.is_act_and_mul and quant_config.gemm1_clamp_limit is not None:
+            device = torch.accelerator.current_device_index()
+            self.gemm1_clamp_limit = torch.full(
                 (self.local_num_experts,),
-                float(val),
+                quant_config.gemm1_clamp_limit,
                 dtype=torch.float32,
                 device=device,
             )
-
-        clamp = quant_config.gemm1_clamp_limit
-        if clamp is None:
-            clamp = getattr(moe_config, "swiglu_limit", None)
-        alpha = quant_config.gemm1_alpha
-        if alpha is None:
-            alpha = getattr(moe_config, "swiglu_alpha", None)
-        beta = quant_config.gemm1_beta
-        if beta is None:
-            beta = getattr(moe_config, "swiglu_beta", None)
-
-
-        if moe_config.is_act_and_mul:
-            self.gemm1_clamp_limit = _per_expert(clamp)
-            self.gemm1_alpha = _per_expert(alpha)
-            self.gemm1_beta = _per_expert(beta)
         else:
             self.gemm1_clamp_limit = None
-            self.gemm1_alpha = None
-            self.gemm1_beta = None
-
-        logger.info_once(
-            "activation=%s, gemm1_alpha=%s, gemm1_beta=%s, gemm1_clamp_limit=%s",
-            moe_config.activation,
-            alpha,
-            beta,
-            clamp,
-        )
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -140,25 +109,6 @@ class TrtLlmNvFp4ExpertsBase:
             )
             self.gemm1_clamp_limit = layer.gemm1_clamp_limit
 
-        # beta shifts the raw GEMM1 accumulator, so fold by g1_alphas like the
-        # clamp limit. alpha is applied to the dequantized gate, so it stays
-        # raw. Register both on the layer so EPLB rearranges them with the
-        # other per-expert tensors.
-        if self.gemm1_beta is not None:
-            gemm1_beta = self.gemm1_beta / self.quant_config.g1_alphas
-            layer.register_parameter(
-                "gemm1_beta",
-                torch.nn.Parameter(gemm1_beta, requires_grad=False),
-            )
-            self.gemm1_beta = layer.gemm1_beta
-
-        if self.gemm1_alpha is not None:
-            layer.register_parameter(
-                "gemm1_alpha",
-                torch.nn.Parameter(self.gemm1_alpha, requires_grad=False),
-            )
-            self.gemm1_alpha = layer.gemm1_alpha
-
     @staticmethod
     def _supports_current_device() -> bool:
         """Supports only Blackwell-family GPUs."""
@@ -187,14 +137,12 @@ class TrtLlmNvFp4ExpertsBase:
 
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
-        """Supports SiLU, RELU^2 non-gated, GELU, and clamped SwiGLU-OAI."""
+        """Supports only SiLU, RELU^2 non-gated and GELU activation."""
         return activation in [
             MoEActivation.SILU,
             MoEActivation.RELU2_NO_MUL,
             MoEActivation.GELU,
             MoEActivation.GELU_TANH,
-            MoEActivation.SWIGLUOAI,
-            MoEActivation.SWIGLUOAI_UNINTERLEAVE,
         ]
 
     @staticmethod
@@ -300,8 +248,8 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
             gemm1_weights=w1,
             gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn),
             gemm1_bias=None,
-            gemm1_alpha=self.gemm1_alpha,
-            gemm1_beta=self.gemm1_beta,
+            gemm1_alpha=None,
+            gemm1_beta=None,
             gemm1_clamp_limit=self.gemm1_clamp_limit,
             gemm2_weights=w2,
             gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),
@@ -461,8 +409,8 @@ class TrtLlmNvFp4ExpertsMonolithic(
             gemm1_weights=w1,
             gemm1_weights_scale=self.quant_config.w1_scale.view(torch.float8_e4m3fn),
             gemm1_bias=None,
-            gemm1_alpha=self.gemm1_alpha,
-            gemm1_beta=self.gemm1_beta,
+            gemm1_alpha=None,
+            gemm1_beta=None,
             gemm1_clamp_limit=self.gemm1_clamp_limit,
             gemm2_weights=w2,
             gemm2_weights_scale=self.quant_config.w2_scale.view(torch.float8_e4m3fn),

vllm/model_executor/layers/fused_moe/layer.py

@@ -140,7 +140,6 @@ def FusedMoE(
     apply_routed_scale_to_output: bool = False,
     zero_expert_type: str | None = None,
     hash_indices_table: torch.Tensor | None = None,
-    reduce_results: bool = True,
     runner_cls: type[MoERunner] | None = None,
     runner_args: dict[str, Any] | None = None,
     routed_experts_cls: type[RoutedExperts] | None = None,
@@ -199,9 +198,6 @@ def FusedMoE(
                                       output instead of topk_weights
         zero_expert_type: Type of zero expert handling
         hash_indices_table: Hash table for expert indices
-        reduce_results: Whether to all-reduce the final output across TP/EP
-            ranks. Set to False to defer the all-reduce (e.g. to fuse it into
-            a subsequent GemmaRMSNorm).
         runner_cls: Custom MoERunner class (None = use default MoERunner)
         runner_args: Additional arguments for runner constructor
         routed_experts_cls: Custom RoutedExperts class (None = use default)
@@ -389,7 +385,6 @@ def FusedMoE(
         routed_scaling_factor=routed_scaling_factor
         if apply_routed_scale_to_output
         else 1.0,
-        reduce_results=reduce_results,
         **runner_args if runner_args is not None else {},
     )
 

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -254,7 +254,6 @@ class MoERunner(MoERunnerInterface):
         routed_input_transform: torch.nn.Module | None = None,
         routed_output_transform: torch.nn.Module | None = None,
         routed_scaling_factor: float = 1.0,
-        reduce_results: bool = True,
     ):
         super().__init__()
         self.moe_config = moe_config
@@ -266,7 +265,6 @@ class MoERunner(MoERunnerInterface):
         self.shared_expert_gate = shared_expert_gate
         self.routed_experts = routed_experts
         self.enable_dbo = enable_dbo
-        self.reduce_results = reduce_results
 
         # When both gates are present and FSE is enabled, fuse their
         # weight matrices into [num_experts + num_shared, hidden] so one
@@ -422,15 +420,6 @@ class MoERunner(MoERunnerInterface):
         * If we have SP (TP=N, DP=M, EP), there is a separate AG step handled
           in the model.
         """
-        # A combine kernel that already reduces the fused output is
-        # incompatible with deferring the all-reduce (reduce_results=False,
-        # e.g. fusing it into a subsequent GemmaRMSNorm): the deferred
-        # all-reduce would double-reduce the fused output.
-        assert not (self._fused_output_is_reduced and not self.reduce_results), (
-            "reduce_results=False is incompatible with a combine kernel that "
-            "already reduces the fused output (e.g. DeepEP/Mori/NIXL/"
-            "FlashInfer-NVLink all2all backends)."
-        )
         if (
             shared_output is not None
             and not self.moe_config.is_sequence_parallel
@@ -458,7 +447,6 @@ class MoERunner(MoERunnerInterface):
             not self.moe_config.is_sequence_parallel
             and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
             and not self._fused_output_is_reduced
-            and self.reduce_results
         ):
             states = tensor_model_parallel_all_reduce(states)
 

vllm/model_executor/layers/fused_moe/utils.py

@@ -296,6 +296,7 @@ def moe_kernel_quantize_input(
         if not quantization_emulation:
             return _nvfp4_quantize(A, A_scale, is_sf_swizzled_layout=is_scale_swizzled)
         else:
+            assert A_scale is not None
             A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
             return A, None
     elif quant_dtype == "mxfp4":

vllm/model_executor/layers/pooler/seqwise/methods.py

@@ -10,6 +10,7 @@ import torch.nn as nn
 from vllm.config.pooler import SequencePoolingType
 from vllm.model_executor.layers.pooler import PoolingParamsUpdate
 from vllm.tasks import PoolingTask
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.pool.metadata import PoolingMetadata
 
 SequencePoolingMethodOutput: TypeAlias = torch.Tensor | list[torch.Tensor]
@@ -74,15 +75,14 @@ class MeanPool(SequencePoolingMethod):
             # early return for empty batch
             return hidden_states.new_empty((0, hidden_size), dtype=torch.float32)
 
-        # Build segment_ids on CPU so repeat_interleave doesn't need to sync
-        # GPU->CPU to learn its data-dependent output length, then upload
-        # non-blocking. eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2]
+        prompt_lens = async_tensor_h2d(
+            prompt_lens_cpu, device=hidden_states.device, dtype=torch.int64
+        )
+        # eg. [2, 1, 3] -> [0, 0, 1, 2, 2, 2]
         segment_ids = torch.repeat_interleave(
-            torch.arange(num_seqs, dtype=torch.long),
-            prompt_lens_cpu,
-        ).to(hidden_states.device, non_blocking=True)
-        prompt_lens = prompt_lens_cpu.to(
-            hidden_states.device, dtype=torch.int64, non_blocking=True
+            torch.arange(num_seqs, device=hidden_states.device, dtype=torch.long),
+            prompt_lens,
+            output_size=int(prompt_lens_cpu.sum()),
         )
         segment_sums = torch.zeros(
             (num_seqs, hidden_size),

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py

@@ -153,8 +153,6 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
             a2_scale=layer.w2_input_scale,
             block_shape=self.weight_block_size,
             swiglu_limit=getattr(layer, "swiglu_limit", None),
-            gemm1_alpha=getattr(layer, "swiglu_alpha", None),
-            gemm1_beta=getattr(layer, "swiglu_beta", None),
         )
 
     def maybe_make_prepare_finalize(

vllm/model_executor/layers/quantization/modelopt.py

@@ -2283,7 +2283,6 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
         fp8_config: ModelOptFp8Config,
         nvfp4_config: ModelOptNvFp4Config,
         w4a16_nvfp4_config: ModelOptNvFp4Config,
-        mxfp8_config: ModelOptMxFp8Config,
     ) -> None:
         super().__init__(exclude_modules)
         self.kv_cache_quant_method = kv_cache_quant_method
@@ -2291,7 +2290,6 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
         self.fp8_config = fp8_config
         self.nvfp4_config = nvfp4_config
         self.w4a16_nvfp4_config = w4a16_nvfp4_config
-        self.mxfp8_config = mxfp8_config
 
     def get_name(self) -> QuantizationMethods:
         return "modelopt_mixed"
@@ -2381,12 +2379,6 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
             group_size=group_size,
         )
 
-        mxfp8_config = ModelOptMxFp8Config(
-            is_checkpoint_mxfp8_serialized=True,
-            kv_cache_quant_algo=kv_cache_quant_method,
-            exclude_modules=[],
-        )
-
         return cls(
             kv_cache_quant_method=kv_cache_quant_method,
             exclude_modules=exclude_modules,
@@ -2394,7 +2386,6 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
             fp8_config=fp8_config,
             nvfp4_config=nvfp4_config,
             w4a16_nvfp4_config=w4a16_nvfp4_config,
-            mxfp8_config=mxfp8_config,
         )
 
     def _resolve_quant_algo(self, prefix: str) -> str | None:
@@ -2449,17 +2440,6 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
                 if key.startswith(parent_dot):
                     return info["quant_algo"].upper()
 
-        # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj).
-        for candidate in self._quantized_layer_prefix_candidates(prefix):
-            parent_dot = candidate.rsplit(".", 1)[0] + "."
-            algos = {
-                info["quant_algo"].upper()
-                for key, info in self.quantized_layers.items()
-                if key.startswith(parent_dot) and "." not in key[len(parent_dot):]
-            }
-            if len(algos) == 1:
-                return algos.pop()
-
         return None
 
     @staticmethod
@@ -2505,8 +2485,6 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
                 return ModelOptNvFp4LinearMethod(self.nvfp4_config)
             if quant_algo == "W4A16_NVFP4":
                 return ModelOptNvFp4W4A16LinearMethod(self.w4a16_nvfp4_config)
-            if quant_algo == "MXFP8":
-                return ModelOptMxFp8LinearMethod(self.mxfp8_config)
             # Layer not in quantized_layers — leave unquantized
             return UnquantizedLinearMethod()
 
@@ -2526,11 +2504,6 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
                     quant_config=self.w4a16_nvfp4_config,
                     moe_config=layer.moe_config,
                 )
-            if quant_algo == "MXFP8":
-                return ModelOptMxFp8FusedMoE(
-                    quant_config=self.mxfp8_config,
-                    moe_config=layer.moe_config,
-                )
             return None
 
         return None

vllm/model_executor/layers/quantization/online/mxfp8.py

@@ -224,8 +224,6 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
             w2_bias=getattr(layer, "w2_bias", None),
             block_shape=self.weight_block_size,
             swiglu_limit=getattr(layer, "swiglu_limit", None),
-            gemm1_alpha=getattr(layer, "swiglu_alpha", None),
-            gemm1_beta=getattr(layer, "swiglu_beta", None),
         )
 
     def process_weights_after_loading(self, layer: Module) -> None:

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

@@ -36,13 +36,6 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType"
         MoEActivation.GELU: ActivationType.Geglu,
         MoEActivation.GELU_TANH: ActivationType.Geglu,
         MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
-        # Both OAI variants map to Swiglu: FlashInfer has no SwigluOAI enum;
-        # the clamped/biased behavior is driven by the per-expert gemm1_alpha/
-        # gemm1_beta/gemm1_clamp_limit tensors (see trtllm_nvfp4_moe.py).
-        # The interleaved-vs-contiguous row layout difference between the two
-        # is resolved in process_weights_after_loading, not here.
-        MoEActivation.SWIGLUOAI: ActivationType.Swiglu,
-        MoEActivation.SWIGLUOAI_UNINTERLEAVE: ActivationType.Swiglu,
     }
     return ACTIVATION_TO_FI_ACTIVATION[activation]
 

vllm/model_executor/models/deepseek_v2.py

@@ -1001,24 +1001,30 @@ class DeepseekV2MLAAttention(nn.Module):
         # IndexCache config
         # Refer: https://arxiv.org/abs/2603.12201 for more details.
         _skip_topk = False
-        _index_topk_freq = getattr(config, "index_topk_freq", 1)
-        _index_topk_pattern = getattr(config, "index_topk_pattern", None)
-        _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
-        layer_id = extract_layer_index(prefix)
+        is_mtp_layer = False
+        if self.is_v32:
+            _index_topk_freq = getattr(config, "index_topk_freq", 1)
+            _index_topk_pattern = getattr(config, "index_topk_pattern", None)
+            _index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)
+            layer_id = extract_layer_index(prefix)
 
-        if _index_topk_pattern is None:
-            _skip_topk = (
-                max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq != 0
+            if _index_topk_pattern is None:
+                _skip_topk = (
+                    max(layer_id - _index_skip_topk_offset + 1, 0) % _index_topk_freq
+                    != 0
+                )
+            elif 0 <= layer_id < len(_index_topk_pattern):
+                _skip_topk = _index_topk_pattern[layer_id] == "S"
+
+            # The skip pattern only governs backbone layers. MTP/nextn
+            # layers (layer_id >= num_hidden_layers) always build a full
+            # indexer: they compute indices at draft step 0 and toggle
+            # at runtime via set_skip_topk
+            # (index_share_for_mtp_iteration).
+            _num_hidden_layers = getattr(config, "num_hidden_layers", None)
+            is_mtp_layer = (
+                _num_hidden_layers is not None and layer_id >= _num_hidden_layers
             )
-        elif 0 <= layer_id < len(_index_topk_pattern):
-            _skip_topk = _index_topk_pattern[layer_id] == "S"
-
-        # The skip pattern only governs backbone layers. MTP/nextn layers
-        # (layer_id >= num_hidden_layers) always build a full indexer: they
-        # compute indices at draft step 0 and toggle at runtime via
-        # set_skip_topk (index_share_for_mtp_iteration).
-        _num_hidden_layers = getattr(config, "num_hidden_layers", None)
-        is_mtp_layer = _num_hidden_layers is not None and layer_id >= _num_hidden_layers
 
         if self.is_v32 and (not _skip_topk or is_mtp_layer):
             self.indexer_rope_emb = get_rope(

vllm/model_executor/models/moonvit.py

@@ -66,6 +66,7 @@ from vllm.model_executor.models.utils import maybe_prefix
 from vllm.model_executor.models.vision import is_vit_use_data_parallel
 from vllm.platforms import current_platform
 from vllm.transformers_utils.configs.moonvit import MoonViTConfig
+from vllm.utils.torch_utils import async_tensor_h2d
 
 
 def _apply_rope_input_validation(x, freqs_cis):
@@ -758,7 +759,7 @@ class MoonVitPretrainedModel(PreTrainedModel):
                         ),
                     ]
                 )
-        metadata["cu_seqlens"] = torch.from_numpy(cu_seqlens_np).to(device)
+        metadata["cu_seqlens"] = async_tensor_h2d(cu_seqlens_np, device=device)
 
         if max_seqlen_override is not None:
             max_seqlen_val = int(max_seqlen_override)
@@ -770,7 +771,7 @@ class MoonVitPretrainedModel(PreTrainedModel):
         metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)
 
         gather_idx_np = _build_merge_gather_idx(grid_pairs, self.merge_kernel_size)
-        metadata["merge_gather_idx"] = torch.from_numpy(gather_idx_np).to(device)
+        metadata["merge_gather_idx"] = async_tensor_h2d(gather_idx_np, device=device)
 
         return metadata
 

vllm/model_executor/models/qwen2_5_vl.py

@@ -83,9 +83,8 @@ from vllm.multimodal.parse import MultiModalDataItems
 from vllm.multimodal.processing import PromptReplacement, PromptUpdate
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
-from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
 
@@ -825,7 +824,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
     @staticmethod
     def invert_permutation(perm: torch.Tensor) -> torch.Tensor:
         # building the inverse permutation in O(n) time
-        inv = torch.empty_like(perm, pin_memory=is_pin_memory_available())
+        inv = torch.empty_like(perm, pin_memory=PIN_MEMORY)
         inv[perm] = torch.arange(perm.numel(), device=perm.device, dtype=perm.dtype)
         return inv
 

vllm/models/deepseek_v4/sparse_mla.py

@@ -13,6 +13,7 @@ from vllm.config.cache import CacheDType
 from vllm.platforms.interface import DeviceCapability
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
+from vllm.utils.torch_utils import np_to_pinned_tensor
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -207,7 +208,7 @@ class DeepseekV4FlashMLAMetadataBuilder(
         # Zero-fill for cudagraphs
         self.req_id_per_token_buffer.fill_(0)
         self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
-            torch.from_numpy(req_id_per_token), non_blocking=True
+            np_to_pinned_tensor(req_id_per_token), non_blocking=True
         )
         req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
 

vllm/models/minimax_m3/amd/model.py

@@ -457,7 +457,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
         cache_config: CacheConfig | None = None,
-        topk_indices_buffer: torch.Tensor | None = None,
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -536,9 +535,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
             self.kv_cache_dtype, vllm_config.model_config
         )
 
-        # Shared top-k buffer: the indexer writes the selected blocks into it and
-        # the attend impl reads them back (no Python value crosses the break).
-        self.topk_indices_buffer = topk_indices_buffer
         self.attn_backend = MiniMaxM3SparseBackend
         # Indexer and main attention are separate impls. On ROCm the SM100 gate
         # is always False, so both pick Triton and the index cache stays bf16.
@@ -569,7 +565,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
             local_blocks=sparse_cfg.get("sparse_local_block", 0),
             score_type=sparse_cfg.get("sparse_score_type", "max"),
             cache_config=cache_config,
-            topk_indices_buffer=topk_indices_buffer,
         )
 
         # Register the main K/V cache so the KV-cache manager allocates it.
@@ -662,10 +657,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         output: torch.Tensor,
     ) -> torch.Tensor:
         # Single eager break around both: their split-K kernels read per-request
-        # metadata and can't be captured into a cudagraph. The indexer writes its
-        # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
-        self.indexer(index_query)
-        return self.impl.forward(self, query, self.kv_cache, output)
+        # metadata and can't be captured into a cudagraph.
+        topk_idx = self.indexer(index_query)
+        return self.impl.forward(self, query, self.kv_cache, topk_idx, output)
 
 
 class MiniMaxM3DecoderLayer(nn.Module):
@@ -677,7 +671,6 @@ class MiniMaxM3DecoderLayer(nn.Module):
         quant_config: QuantizationConfig | None = None,
         force_sparse_attn: bool = False,
         force_moe: bool = False,
-        topk_indices_buffer: torch.Tensor | None = None,
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -697,7 +690,6 @@ class MiniMaxM3DecoderLayer(nn.Module):
                 quant_config=quant_config,
                 prefix=f"{prefix}.self_attn",
                 cache_config=cache_config,
-                topk_indices_buffer=topk_indices_buffer,
             )
         else:
             self.self_attn = MiniMaxM3Attention(
@@ -779,22 +771,6 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
             prefix=f"{prefix}.embed_tokens",
         )
 
-        # Reserved top-k indices buffer shared by all sparse-attention indexer
-        # layers (mirrors DeepseekV4); the indexer writes its per-head decode/
-        # prefill block selection into it, the attend reads it back.
-        sparse_cfg = getattr(config, "sparse_attention_config", None)
-        if sparse_cfg is not None:
-            tp_size = get_tensor_model_parallel_world_size()
-            num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
-            self.topk_indices_buffer = torch.empty(
-                num_index_heads,
-                vllm_config.scheduler_config.max_num_batched_tokens,
-                sparse_cfg["sparse_topk_blocks"],
-                dtype=torch.int32,
-            )
-        else:
-            self.topk_indices_buffer = None
-
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
             lambda prefix: MiniMaxM3DecoderLayer(
@@ -802,7 +778,6 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
                 prefix,
                 cache_config=cache_config,
                 quant_config=quant_config,
-                topk_indices_buffer=self.topk_indices_buffer,
             ),
             prefix=f"{prefix}.layers",
         )

vllm/models/minimax_m3/common/indexer.py

@@ -25,14 +25,12 @@ from vllm.config.attention import IndexerKVDType
 from vllm.config.cache import CacheDType
 from vllm.distributed import get_tensor_model_parallel_world_size
 from vllm.forward_context import get_forward_context
-from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.models.minimax_m3.common.ops.index_topk import (
     minimax_m3_index_decode,
     minimax_m3_index_score,
     minimax_m3_index_topk,
 )
-from vllm.platforms import current_platform
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -48,8 +46,6 @@ from vllm.v1.kv_cache_interface import (
     MLAAttentionSpec,
 )
 
-logger = init_logger(__name__)
-
 
 class MiniMaxM3IndexerBackend(AttentionBackend):
     """Indexer side-cache backend (key-only)."""
@@ -124,20 +120,16 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
         backend_cls: type[AttentionBackend] = MiniMaxM3IndexerBackend,
     ) -> None:
         super().__init__()
-        if indexer_kv_dtype in ("fp8", "fp8_e4m3"):
-            cache_dtype = torch.float8_e4m3fn
-        elif indexer_kv_dtype == "bf16":
-            cache_dtype = torch.bfloat16
-        else:
+        if indexer_kv_dtype != "bf16":
             raise NotImplementedError(
-                f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
-                "MiniMax M3 indexer cache (only 'bf16' or 'fp8'/'fp8_e4m3')."
+                f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported yet "
+                "for the MiniMax M3 indexer cache (only 'bf16')."
             )
         self.kv_cache = torch.tensor([])
         self.head_dim = head_dim
         self.indexer_kv_dtype = indexer_kv_dtype
-        # Side-cache storage dtype: bf16, or e4m3 for the fp8 score path.
-        self.dtype = cache_dtype
+        # Storage dtype for the side cache (bf16 today; quantized layouts later).
+        self.dtype = torch.bfloat16
         self.prefix = prefix
         self.cache_config = cache_config
         # Impl-chosen backend -> each impl gets its own builder (get_attn_backend).
@@ -352,7 +344,6 @@ class MiniMaxM3IndexerImpl(nn.Module):
         score_type: str = "max",
         cache_config: CacheConfig | None = None,
         indexer_kv_dtype: IndexerKVDType = "bf16",
-        topk_indices_buffer: torch.Tensor | None = None,
     ) -> None:
         super().__init__()
         self.num_kv_heads = num_kv_heads
@@ -365,9 +356,6 @@ class MiniMaxM3IndexerImpl(nn.Module):
         self.num_index_heads = num_index_heads
         self.index_head_dim = index_head_dim
         self.indexer_kv_dtype = indexer_kv_dtype
-        # Shared, stable-address top-k output buffer (set by the model for the
-        # cudagraph-safe MSA impl); None -> impl allocates fresh (eager).
-        self.topk_indices_buffer = topk_indices_buffer
         # Owns the side cache (registers itself in the static forward context).
         self.index_cache = MiniMaxM3IndexerCache(
             head_dim=index_head_dim,
@@ -404,10 +392,6 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
         )
         kv = self.index_cache.kv_cache
 
-        # Both sides write into the single shared persistent topk_indices_buffer
-        # (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the
-        # kernels' out= writes out[:, :total_q]. None -> allocate fresh.
-        buf = self.topk_indices_buffer
         decode_topk: torch.Tensor | None = None
         prefill_topk: torch.Tensor | None = None
         if index_md.num_decodes > 0:
@@ -425,7 +409,6 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
                 self.num_kv_heads,
                 d.decode_query_len,
                 d.max_decode_query_len,
-                out=buf,
             )
         if index_md.num_prefills > 0:
             p = index_md.prefill
@@ -449,61 +432,29 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
                 self.topk_blocks,
                 self.init_blocks,
                 self.local_blocks,
-                out=buf[:, nd:, :] if buf is not None else None,
             )
         return decode_topk, prefill_topk
 
 
 def select_indexer_impl_cls(
     *,
-    topk_blocks: int,
     indexer_kv_dtype: IndexerKVDType = "bf16",
 ) -> type[MiniMaxM3IndexerImpl]:
-    """Pick the indexer impl off the platform, top-k count, and cache dtype.
+    """Pick the indexer impl off the index-cache dtype.
 
-    On Blackwell (SM100) with ``topk_blocks`` in ``(4, 8, 16, 32)`` (matching the
-    main MSA attend), the fmha_sm100 score path + Triton top-k is used for both
-    bf16 and fp8 index caches. Everything else falls back to the Triton indexer
-    (bf16 only).
+    The SM100 MSA indexer score path is disabled for now; use the local Triton
+    indexer. If re-enabled, add a NVIDIA-specific ``MiniMaxM3IndexerImpl`` here.
     """
     if indexer_kv_dtype in ("mxfp4", "nvfp4"):
         raise NotImplementedError(
             f"indexer_kv_dtype={indexer_kv_dtype!r} needs the (not-yet-added) "
             "CuteDSL indexer impl."
         )
-    is_sm100 = (
-        current_platform.is_cuda() and current_platform.is_device_capability_family(100)
-    )
-    use_msa = (
-        is_sm100
-        and topk_blocks in (4, 8, 16, 32)
-        and indexer_kv_dtype in ("bf16", "fp8", "fp8_e4m3")
-    )
-    if use_msa:
-        # Lazy import so AMD / non-SM100 never import fmha_sm100.
-        from vllm.models.minimax_m3.nvidia.indexer_msa import (
-            MiniMaxM3IndexerMSAImpl,
-        )
-
-        logger.info_once(
-            "MiniMax M3 indexer: selected MSA (fmha_sm100 score + Triton top-k) "
-            "[topk_blocks=%d, indexer_kv_dtype=%s]",
-            topk_blocks,
-            indexer_kv_dtype,
-        )
-        return MiniMaxM3IndexerMSAImpl
     if indexer_kv_dtype != "bf16":
         raise NotImplementedError(
             f"indexer_kv_dtype={indexer_kv_dtype!r} is not supported by the "
             "Triton indexer impl."
         )
-    logger.info_once(
-        "MiniMax M3 indexer: selected Triton (no fmha_sm100) "
-        "[topk_blocks=%d, indexer_kv_dtype=%s, sm100=%s]",
-        topk_blocks,
-        indexer_kv_dtype,
-        is_sm100,
-    )
     return MiniMaxM3IndexerTritonImpl
 
 
@@ -529,11 +480,9 @@ class MiniMaxM3Indexer(nn.Module):
         score_type: str = "max",
         cache_config: CacheConfig | None = None,
         indexer_kv_dtype: IndexerKVDType = "bf16",
-        topk_indices_buffer: torch.Tensor | None = None,
     ) -> None:
         super().__init__()
         impl_cls = select_indexer_impl_cls(
-            topk_blocks=topk_blocks,
             indexer_kv_dtype=indexer_kv_dtype,
         )
         self.impl = impl_cls(
@@ -549,7 +498,6 @@ class MiniMaxM3Indexer(nn.Module):
             score_type=score_type,
             cache_config=cache_config,
             indexer_kv_dtype=indexer_kv_dtype,
-            topk_indices_buffer=topk_indices_buffer,
         )
 
     @property

vllm/models/minimax_m3/common/ops/index_topk.py

@@ -373,10 +373,7 @@ def _decode_index_score_kernel(
             + off_k[:, None] * stride_ik_pos
             + off_d * stride_ik_d,
         )  # [N,D]
-        # fp32 accumulation is required for the fp8 (e4m3) index cache: q/k are
-        # loaded in their stored dtype (bf16 or e4m3) and the MMA accumulates in
-        # fp32 so the per-block max score is exact for the fp8 indexer too.
-        kq = tl.dot(k, q, out_dtype=tl.float32)  # [N,HQ]
+        kq = tl.dot(k, q)  # [N,HQ]
         kq = tl.where(pos_mask & q_mask[None, :], kq, float("-inf"))
         score = tl.max(kq, axis=0)  # [HQ]
         is_visible_block = blk < num_blocks_q
@@ -712,25 +709,16 @@ def minimax_m3_index_topk(
     topk: int,
     init_blocks: int,
     local_blocks: int,
-    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """Select index top-k from a precomputed score tensor.
-
-    When ``out`` is provided (a ``[num_idx_heads, >=total_q, topk]`` buffer), the
-    result is written into ``out[:, :total_q, :]`` instead of a fresh tensor --
-    used to keep the top-k output at a stable address for cudagraph capture.
-    """
+    """Select index top-k from a precomputed score tensor."""
     num_idx_heads = score.shape[0]
     batch = cu_seqlens_q.shape[0] - 1
     total_q = score.shape[1]
-    if out is not None:
-        topk_idx = out[:, :total_q, :]
-    else:
-        topk_idx = torch.empty(
-            (num_idx_heads, total_q, topk),
-            dtype=torch.int32,
-            device=score.device,
-        )
+    topk_idx = torch.empty(
+        (num_idx_heads, total_q, topk),
+        dtype=torch.int32,
+        device=score.device,
+    )
     # block_size_q == 1 -> query blocks coincide with query tokens.
     grid_topk = (max_query_len, batch, num_idx_heads)
     _topk_index_kernel[grid_topk](
@@ -769,13 +757,10 @@ def minimax_m3_index_decode(
     num_kv_heads: int,
     decode_query_len: int,
     max_decode_query_len: int,
-    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Decode index block-score + top-k, both split-K (cudagraph-safe).
 
     Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).
-    When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into
-    ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.
     """
     total_q, num_idx_heads, head_dim = idx_q.shape
     assert num_idx_heads == num_kv_heads, (
@@ -809,7 +794,7 @@ def minimax_m3_index_decode(
     )
     # split-K over seq blocks; chunk count depends only on shape constants so
     # the grid is fixed within a cuda graph.
-    TARGET_GRID = 4096
+    TARGET_GRID = 512
     MAX_NUM_KV_CHUNKS = 256
     # Use the configured max decode length to avoid Triton recompiles when
     # switching between qlen=1 and spec-decode verification batches.
@@ -849,17 +834,14 @@ def minimax_m3_index_decode(
         **score_kwargs,
     )
 
-    if out is not None:
-        topk_idx = out[:, :total_q, :]
-    else:
-        topk_idx = torch.empty(
-            (num_idx_heads, total_q, topk),
-            dtype=torch.int32,
-            device=idx_q.device,
-        )
+    topk_idx = torch.empty(
+        (num_idx_heads, total_q, topk),
+        dtype=torch.int32,
+        device=idx_q.device,
+    )
     # Chunk count is shape-constant (cudagraph-safe), capped so the merge sorts
     # pow2(num_topk_chunks * pow2(topk)) candidates.
-    TOPK_TARGET_GRID = 512
+    TOPK_TARGET_GRID = 64
     MAX_NUM_TOPK_CHUNKS = 16
     topk_target = max(
         1, min(MAX_NUM_TOPK_CHUNKS, TOPK_TARGET_GRID // max(1, batch * num_idx_heads))

vllm/models/minimax_m3/common/sparse_attention.py

@@ -2,11 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Main block-sparse GQA attention for MiniMax M3 sparse layers.
 
-The lightning indexer (``indexer.py``) selects the top-k KV blocks (written into
-the shared ``layer.topk_indices_buffer``); this module holds the main attention
-that attends only to those blocks: the paged K/V cache backend, its metadata +
-builder, and the impl that reads the indexer's top-k from that buffer. The Triton
-attend kernel lives here; the SM100 (MSA)
+The lightning indexer (``indexer.py``) selects the top-k KV blocks; this module
+holds the main attention that attends only to those blocks: the paged K/V cache
+backend, its metadata + builder, and the impl that consumes the indexer's
+``topk_idx``. The Triton attend kernel lives here; the SM100 (MSA)
 ``build_k2q_csr`` + ``sparse_atten_func`` attend lives in
 ``nvidia/sparse_attention_msa.py``.
 
@@ -273,10 +272,9 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
     """Abstract base for block-sparse GQA over the indexer-selected blocks.
 
     Inherits ``AttentionImplBase`` for a custom forward signature (the layer
-    pre-inserts K/V and runs the indexer, which writes the selected blocks into
-    the shared ``layer.topk_indices_buffer``; the attend reads them back from
-    there). The Triton and MSA subclasses each own a full ``forward`` -- no
-    shared forward code.
+    pre-inserts K/V and runs the indexer, so forward takes the queries +
+    ``topk_idx``). The Triton and MSA subclasses each own a full ``forward`` --
+    no shared forward code.
     """
 
     def __init__(
@@ -313,14 +311,10 @@ class MiniMaxM3SparseImpl(AttentionImplBase[MiniMaxM3SparseMetadata]):
         layer: AttentionLayer,
         query: torch.Tensor,
         kv_cache: torch.Tensor,
+        topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
         output: torch.Tensor,
     ) -> torch.Tensor:
-        """Attend the queries to the indexer-selected blocks. Per kernel.
-
-        The indexer has already written the top-k block ids into
-        ``layer.topk_indices_buffer`` (decode at ``[:, :nd]``, prefill at
-        ``[:, nd:num_tokens]``); the attend reads them from there.
-        """
+        """Attend the queries to the indexer-selected blocks. Per kernel."""
         raise NotImplementedError
 
 
@@ -332,6 +326,7 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
         layer: AttentionLayer,
         query: torch.Tensor,
         kv_cache: torch.Tensor,
+        topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
         output: torch.Tensor,
     ) -> torch.Tensor:
         attn_metadata = get_forward_context().attn_metadata
@@ -339,12 +334,10 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
             return output  # profiling run; caches unbound
         main_md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
         assert isinstance(main_md, MiniMaxM3SparseMetadata)
+        decode_topk, prefill_topk = topk_idx
 
         nd = main_md.num_decode_tokens
         num_tokens = main_md.num_actual_tokens
-        # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].
-        topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
-        assert topk is not None
         hd = self.head_size
         q = query[:num_tokens].view(-1, self.num_heads, hd)
         out = output[:num_tokens].view(-1, self.num_heads, hd)
@@ -355,11 +348,11 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
         # Decode [:nd]: split-K over the selected blocks (request-major chunks).
         if main_md.num_decodes > 0:
             d = main_md.decode
-            assert d is not None
+            assert d is not None and decode_topk is not None
             minimax_m3_sparse_attn_decode(
                 q[:nd],
                 kv_cache,
-                topk[:, :nd, :],
+                decode_topk,
                 d.block_table,
                 d.seq_lens,
                 self.num_kv_heads,
@@ -371,11 +364,11 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
         # Prefill [nd:]: cu_seqlens_q already rebased to 0.
         if main_md.num_prefills > 0:
             p = main_md.prefill
-            assert p is not None
+            assert p is not None and prefill_topk is not None
             minimax_m3_sparse_attn(
                 q[nd:],
                 kv_cache,
-                topk[:, nd:num_tokens, :],
+                prefill_topk,
                 p.block_table,
                 p.cu_seqlens_q,
                 p.seq_lens,

vllm/models/minimax_m3/nvidia/indexer_msa.py

@@ -1,251 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""MSA (SM100/Blackwell) indexer impl for MiniMax M3.
-
-Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path then selects
-top-k blocks with the Triton ``minimax_m3_index_topk`` kernel -- fmha is much
-faster than Triton for the wide prefill score (benchmarked ~3-5x).
-
-Decode uses the Triton fused ``minimax_m3_index_decode`` (the same kernel the
-Triton indexer impl uses): for q_len==1 it is a purpose-built vector x matrix
-score (no wasted tensor-core tiles) with a 256-way split-K and a fused split-K
-top-k, which beats fmha's OnlyScore (wasted MMA on a single query, 64-split cap)
-by ~1.1-3.7x. It is cudagraph-safe by construction (shape-constant split grids)
-and writes the shared ``topk_indices_buffer`` via ``out=``.
-
-``fmha_sm100`` imports are function-local so this module is import-safe on
-AMD / non-SM100.
-"""
-
-from dataclasses import dataclass
-from typing import ClassVar
-
-import torch
-
-from vllm.forward_context import get_forward_context
-from vllm.models.minimax_m3.common.indexer import (
-    MiniMaxM3IndexerBackend,
-    MiniMaxM3IndexerDecodeMetadata,
-    MiniMaxM3IndexerImpl,
-    MiniMaxM3IndexerMetadata,
-    MiniMaxM3IndexerMetadataBuilder,
-)
-from vllm.models.minimax_m3.common.ops.index_topk import (
-    minimax_m3_index_decode,
-    minimax_m3_index_topk,
-)
-from vllm.v1.attention.backend import (
-    AttentionBackend,
-    AttentionCGSupport,
-    CommonAttentionMetadata,
-)
-from vllm.v1.attention.backends.utils import split_decodes_and_prefills
-
-# Page size == sparse block size == index-K block; fmha tile id == M3 block id.
-PAGE_SIZE = 128
-
-
-class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend):
-    """Indexer side-cache backend selecting the MSA builder."""
-
-    @staticmethod
-    def get_builder_cls() -> type["MiniMaxM3IndexerMSAMetadataBuilder"]:
-        return MiniMaxM3IndexerMSAMetadataBuilder
-
-
-@dataclass
-class MiniMaxM3IndexerMSAPrefillMetadata:
-    """fmha score plan + Triton top-k inputs for the prefill side (eager)."""
-
-    plan: dict  # fmha_sm100 PlanInfo
-    cu_seqlens_q: torch.Tensor  # [num_prefills + 1] int32, rebased to 0
-    prefix_lens: torch.Tensor  # [num_prefills] int32, context tokens
-    max_query_len: int
-    page_table: torch.Tensor  # flat physical page indices for the prefill side
-
-
-@dataclass
-class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata):
-    """Decode reuses the inherited base ``decode`` field (the Triton decode
-    metadata); ``prefill_msa`` carries the fmha score plan for the prefill side
-    (the base ``prefill`` field is unused on this path)."""
-
-    prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None
-
-
-class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):
-    """Decode metadata is the cudagraph-safe Triton decode metadata; the prefill
-    fmha plan is built eagerly (prefill batches are not captured)."""
-
-    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
-
-    def build(
-        self,
-        common_prefix_len: int,
-        common_attn_metadata: CommonAttentionMetadata,
-        fast_build: bool = False,
-    ) -> MiniMaxM3IndexerMSAMetadata:
-        num_reqs = common_attn_metadata.num_reqs
-        num_tokens = common_attn_metadata.num_actual_tokens
-        seq_lens = common_attn_metadata.seq_lens
-        block_table = common_attn_metadata.block_table_tensor
-        query_start_loc = common_attn_metadata.query_start_loc
-
-        num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
-            split_decodes_and_prefills(
-                common_attn_metadata,
-                decode_threshold=self.reorder_batch_threshold,
-                require_uniform=True,
-            )
-        )
-        assert num_decodes + num_prefills == num_reqs
-        assert num_decode_tokens + num_prefill_tokens == num_tokens
-
-        # Context (prefix) lengths into the stable cudagraph buffer.
-        context_lens = self.context_len_buffer[:num_reqs]
-        context_lens.copy_(
-            common_attn_metadata.compute_num_computed_tokens(), non_blocking=True
-        )
-
-        decode: MiniMaxM3IndexerDecodeMetadata | None = None
-        if num_decodes > 0:
-            qsl_cpu = common_attn_metadata.query_start_loc_cpu
-            query_lens_cpu = qsl_cpu[1 : num_decodes + 1] - qsl_cpu[:num_decodes]
-            decode_query_len = int(query_lens_cpu[0].item())
-            assert decode_query_len > 0
-            assert torch.all(
-                (query_lens_cpu == decode_query_len) | (query_lens_cpu == 0)
-            )
-            decode = MiniMaxM3IndexerDecodeMetadata(
-                seq_lens=seq_lens[:num_decodes],
-                block_table=block_table[:num_decodes],
-                max_seq_len=common_attn_metadata.max_seq_len,
-                decode_query_len=decode_query_len,
-                max_decode_query_len=self.max_decode_query_len,
-            )
-
-        prefill: MiniMaxM3IndexerMSAPrefillMetadata | None = None
-        if num_prefills > 0:
-            # Prefill is eager (not captured); the host lengths it needs (and the
-            # _fmha_sm100_plan .tolist() inside) make the D->H sync acceptable.
-            from vllm.third_party.fmha_sm100.api import _fmha_sm100_plan
-
-            lo, hi = num_decodes, num_reqs
-            qsl_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1]
-            qo_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32)
-            kv_lens_cpu = seq_lens[:num_reqs].cpu().to(torch.int32)
-            nvp = (kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE
-            side_qo = qo_lens_cpu[lo:hi]
-            side_kv = kv_lens_cpu[lo:hi]
-            plan = _fmha_sm100_plan(
-                side_qo,
-                side_kv,
-                self.num_index_heads,
-                num_kv_heads=1,
-                qo_offset=side_kv - side_qo,  # bottom-right causal
-                page_size=PAGE_SIZE,
-                output_maxscore=True,
-                causal=True,
-                num_kv_splits=1,
-            )
-            cols = torch.arange(block_table.shape[1], device=block_table.device)
-            valid = cols[None, :] < nvp[lo:hi].to(block_table.device)[:, None]
-            prefill = MiniMaxM3IndexerMSAPrefillMetadata(
-                plan=plan,
-                cu_seqlens_q=(query_start_loc[lo : hi + 1] - query_start_loc[lo]).to(
-                    torch.int32
-                ),
-                prefix_lens=context_lens[lo:hi],
-                max_query_len=int(side_qo.max()),
-                page_table=block_table[lo:hi][valid].to(torch.int32),
-            )
-
-        return MiniMaxM3IndexerMSAMetadata(
-            seq_lens=seq_lens,
-            max_seq_len=common_attn_metadata.max_seq_len,
-            slot_mapping=common_attn_metadata.slot_mapping,
-            num_actual_tokens=num_tokens,
-            num_decodes=num_decodes,
-            num_decode_tokens=num_decode_tokens,
-            num_prefills=num_prefills,
-            num_prefill_tokens=num_prefill_tokens,
-            decode=decode,
-            prefill_msa=prefill,
-        )
-
-
-class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):
-    """Decode: Triton fused score+top-k. Prefill: fmha_sm100 OnlyScore + top-k."""
-
-    indexer_backend_cls: ClassVar[type[AttentionBackend]] = MiniMaxM3IndexerMSABackend
-
-    def forward(
-        self,
-        index_query: torch.Tensor,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-        attn_metadata = get_forward_context().attn_metadata
-        if not isinstance(attn_metadata, dict):
-            return None, None  # profiling run; caches unbound
-        md = attn_metadata[self.index_cache.prefix]
-        assert isinstance(md, MiniMaxM3IndexerMSAMetadata)
-
-        num_tokens = md.num_actual_tokens
-        nd = md.num_decode_tokens
-        index_q = index_query[:num_tokens].view(
-            -1, self.num_index_heads, self.index_head_dim
-        )
-        kv = self.index_cache.kv_cache
-        # Both sides write into the single shared persistent topk_indices_buffer:
-        # decode at [:, :nd], prefill at [:, nd:] (each kernel writes [:, :total_q]).
-        buf = self.topk_indices_buffer
-
-        decode_topk: torch.Tensor | None = None
-        if md.decode is not None:
-            d = md.decode
-            decode_topk = minimax_m3_index_decode(
-                index_q[:nd],
-                kv,
-                d.block_table,
-                d.seq_lens,
-                d.max_seq_len,
-                self.topk_blocks,
-                self.init_blocks,
-                self.local_blocks,
-                self.num_kv_heads,
-                d.decode_query_len,
-                d.max_decode_query_len,
-                out=buf,
-            )
-
-        prefill_topk: torch.Tensor | None = None
-        if md.prefill_msa is not None:
-            from vllm.third_party.fmha_sm100.api import _fmha_sm100
-
-            p = md.prefill_msa
-            # Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks,1,128,D).
-            k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim)
-            _, max_score = _fmha_sm100(
-                index_q[nd:],
-                k_pages,
-                k_pages,  # V placeholder; not read in OnlyScore
-                p.plan,
-                kv_indices=p.page_table,
-                output_o=False,
-                output_maxscore=True,
-                sm_scale=self.scale,
-            )
-            # Triton top-k wants [num_index_heads, num_tokens, max_block]; the
-            # transpose is a strided view (the kernel reads via strides).
-            out = buf[:, nd:, :] if buf is not None else None
-            prefill_topk = minimax_m3_index_topk(
-                max_score.transpose(1, 2),
-                p.cu_seqlens_q,
-                p.prefix_lens,
-                p.max_query_len,
-                self.topk_blocks,
-                self.init_blocks,
-                self.local_blocks,
-                out=out,
-            )
-
-        return decode_topk, prefill_topk

vllm/models/minimax_m3/nvidia/model.py

@@ -193,7 +193,6 @@ class MiniMaxM3MoE(nn.Module):
         layer_id: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
-        reduce_results: bool = True,
     ) -> None:
         super().__init__()
         self.tp_size = get_tensor_model_parallel_world_size()
@@ -260,7 +259,6 @@ class MiniMaxM3MoE(nn.Module):
             shared_experts=self.shared_experts,
             quant_config=quant_config,
             prefix=f"{prefix}.experts",
-            reduce_results=reduce_results,
         )
 
     @staticmethod
@@ -404,7 +402,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
         cache_config: CacheConfig | None = None,
-        topk_indices_buffer: torch.Tensor | None = None,
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -492,10 +489,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         # cache (--attention-config '{"indexer_kv_dtype": ...}').
         self.indexer_kv_dtype = vllm_config.attention_config.indexer_kv_dtype
 
-        # Shared top-k buffer: the indexer writes the selected blocks into it and
-        # the attend impl reads them back (so nothing crosses the eager break as a
-        # Python value, which would freeze at capture).
-        self.topk_indices_buffer = topk_indices_buffer
         self.attn_backend = MiniMaxM3SparseBackend
         # Indexer (top-k selection) and main attention are separate impls, each
         # picking Triton vs MSA off its cache dtype. impl is AttentionImplBase
@@ -526,7 +519,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
             score_type=sparse_cfg.get("sparse_score_type", "max"),
             cache_config=cache_config,
             indexer_kv_dtype=self.indexer_kv_dtype,
-            topk_indices_buffer=topk_indices_buffer,
         )
 
         # Register the main K/V cache so the KV-cache manager allocates it.
@@ -584,12 +576,7 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         main_slot_mapping = fwd_slot_mapping[self.layer_name]
         index_slot_mapping = fwd_slot_mapping[self.indexer.index_cache.prefix]
         q = qkv.new_empty((num_tokens, self.q_size))
-        # index_q matches the index-K cache dtype (e4m3 for the fp8 score path);
-        # the fused kernel emits fp8 directly when this buffer is e4m3.
-        index_q = qkv.new_empty(
-            (num_tokens, self.index_q_size),
-            dtype=self.indexer.index_cache.dtype,
-        )
+        index_q = qkv.new_empty((num_tokens, self.index_q_size))
         ops.fused_minimax_m3_qknorm_rope_kv_insert(
             qkv,
             self.q_norm.weight,
@@ -626,10 +613,9 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
         output: torch.Tensor,
     ) -> torch.Tensor:
         # Single eager break around both: their split-K kernels read per-request
-        # metadata and can't be captured into a cudagraph. The indexer writes its
-        # top-k into the shared ``topk_indices_buffer``; the attend reads it back.
-        self.indexer(index_query)
-        return self.impl.forward(self, query, self.kv_cache, output)
+        # metadata and can't be captured into a cudagraph.
+        topk_idx = self.indexer(index_query)
+        return self.impl.forward(self, query, self.kv_cache, topk_idx, output)
 
 
 class MiniMaxM3DecoderLayer(nn.Module):
@@ -641,7 +627,6 @@ class MiniMaxM3DecoderLayer(nn.Module):
         force_sparse_attn: bool = False,
         force_moe: bool = False,
         is_mtp_block: bool = False,
-        topk_indices_buffer: torch.Tensor | None = None,
     ) -> None:
         super().__init__()
         if is_mtp_block:
@@ -657,12 +642,13 @@ class MiniMaxM3DecoderLayer(nn.Module):
         layer_id = int(prefix.split(sep=".")[-1])
         self.layer_id = layer_id
 
-        # Complete the preceding FFN's deferred all-reduce (its down_proj / MoE
-        # combine ran with reduce_results=False), fused into this layer's
-        # input_layernorm. Both dense and MoE FFNs defer under PP==1, so every
-        # non-first layer fuses; disable when PP>1 (FFNs reduce themselves).
+        # Complete the preceding dense MLP's deferred all-reduce
+        # (reduce_results=False), fused into this layer's input_layernorm.
+        # Disable this fusion when PP is set
         self.fuse_input_allreduce = (
-            layer_id > 0 and vllm_config.parallel_config.pipeline_parallel_size == 1
+            layer_id > 0
+            and not _is_moe_layer(config, layer_id - 1)
+            and vllm_config.parallel_config.pipeline_parallel_size == 1
         )
 
         is_sparse_attention_layer = (
@@ -676,7 +662,6 @@ class MiniMaxM3DecoderLayer(nn.Module):
                 quant_config=quant_config,
                 prefix=f"{prefix}.self_attn",
                 cache_config=cache_config,
-                topk_indices_buffer=topk_indices_buffer,
             )
         else:
             self.self_attn = MiniMaxM3Attention(
@@ -696,12 +681,6 @@ class MiniMaxM3DecoderLayer(nn.Module):
                 layer_id=layer_id,
                 quant_config=quant_config,
                 prefix=f"{prefix}.block_sparse_moe",
-                # Defer the MoE all-reduce only when it can be fused into a
-                # following GemmaRMSNorm
-                reduce_results=(
-                    vllm_config.parallel_config.pipeline_parallel_size > 1
-                    or is_mtp_block
-                ),
             )
         else:
             self.mlp = MiniMaxM3MLP(
@@ -768,43 +747,17 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
             prefix=f"{prefix}.embed_tokens",
         )
 
-        # Reserved top-k indices buffer shared by all sparse-attention indexer
-        # layers (mirrors DeepseekV4); kept at a stable address so the indexer's
-        # top-k output survives cudagraph capture/replay. Shape matches the
-        # per-head index top-k output [num_index_heads, total_q, topk].
-        sparse_cfg = getattr(config, "sparse_attention_config", None)
-        if sparse_cfg is not None:
-            tp_size = get_tensor_model_parallel_world_size()
-            num_index_heads = max(1, sparse_cfg["sparse_num_index_heads"] // tp_size)
-            max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
-            padded_num_tokens = (max_num_batched_tokens + 3) // 4 * 4
-            self.topk_indices_buffer = torch.empty(
-                num_index_heads,
-                padded_num_tokens,
-                sparse_cfg["sparse_topk_blocks"],
-                dtype=torch.int32,
-            )
-        else:
-            self.topk_indices_buffer = None
-
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
             lambda prefix: MiniMaxM3DecoderLayer(
                 vllm_config=vllm_config,
                 prefix=prefix,
-                topk_indices_buffer=self.topk_indices_buffer,
             ),
             prefix=f"{prefix}.layers",
         )
 
         self.norm = MiniMAXGemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        # The final decoder layer has no next layer, so its deferred all-reduce is
-        # completed here in the model norm.
-        self.fuse_final_allreduce = (
-            vllm_config.parallel_config.pipeline_parallel_size == 1
-        )
-
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         return self.embed_tokens(input_ids)
 
@@ -828,12 +781,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
                 aux_hidden_states, idx + 1, hidden_states, residual
             )
 
-        if self.fuse_final_allreduce and residual is not None:
-            hidden_states, _ = fused_allreduce_gemma_rms_norm(
-                hidden_states, residual, self.norm
-            )
-        else:
-            hidden_states, _ = self.norm(hidden_states, residual)
+        hidden_states, _ = self.norm(hidden_states, residual)
 
         if len(aux_hidden_states) > 0:
             return hidden_states, aux_hidden_states

vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

@@ -29,6 +29,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
         layer: AttentionLayer,
         query: torch.Tensor,
         kv_cache: torch.Tensor,
+        topk_idx: tuple[torch.Tensor | None, torch.Tensor | None],
         output: torch.Tensor,
     ) -> torch.Tensor:
         attn_metadata = get_forward_context().attn_metadata
@@ -36,12 +37,10 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
             return output  # profiling run; caches unbound
         main_md = attn_metadata[layer.layer_name]  # type: ignore[attr-defined]
         assert isinstance(main_md, MiniMaxM3SparseMetadata)
+        decode_topk, prefill_topk = topk_idx
 
         nd = main_md.num_decode_tokens
         num_tokens = main_md.num_actual_tokens
-        # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].
-        topk = layer.topk_indices_buffer  # type: ignore[attr-defined]
-        assert topk is not None
         hd = self.head_size
         q = query[:num_tokens].view(-1, self.num_heads, hd)
         out = output[:num_tokens].view(-1, self.num_heads, hd)
@@ -52,11 +51,11 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
         # Decode [:nd]: Triton split-K placeholder (no MSA decode yet).
         if main_md.num_decodes > 0:
             d = main_md.decode
-            assert d is not None
+            assert d is not None and decode_topk is not None
             minimax_m3_sparse_attn_decode(
                 q[:nd],
                 kv_cache,
-                topk[:, :nd, :],
+                decode_topk,
                 d.block_table,
                 d.seq_lens,
                 self.num_kv_heads,
@@ -73,9 +72,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):
             )
 
             p = main_md.prefill
-            assert p is not None
-            # build_k2q_csr() doesn't support strided topk buffer
-            prefill_topk = topk[:, nd:num_tokens, :]
+            assert p is not None and prefill_topk is not None
             qp = q[nd:]
             k_cache = kv_cache[:, 0].transpose(1, 2)
             v_cache = kv_cache[:, 1].transpose(1, 2)

vllm/multimodal/inputs.py

@@ -488,7 +488,13 @@ class MultiModalBatchedField(BaseMultiModalField):
                 # An optimization when `batch` contains only one tensor:
                 # - produce exactly same result as `torch.stack(batch)`
                 # - will achieve zero-copy if the tensor is contiguous
-                return batch[0].unsqueeze(0).contiguous()
+                out = batch[0].unsqueeze(0)
+                if not pin_memory:
+                    return out.contiguous()
+                # Avoid extra copy - pinning unpinned memory will make it contiguous
+                if not out.is_contiguous() and out.is_pinned():
+                    out = out.contiguous()
+                return out.pin_memory()
             first_shape = batch[0].shape
             if all(elem.shape == first_shape for elem in batch):
                 out = torch.empty(
@@ -538,7 +544,13 @@ class MultiModalFlatField(BaseMultiModalField):
                 # An optimization when `batch` contains only one tensor:
                 # - produce exactly same result as `torch.concat(batch)`
                 # - will achieve zero-copy if the tensor is contiguous
-                return batch[0].contiguous()
+                out = batch[0]
+                if not pin_memory:
+                    return out.contiguous()
+                # Avoid extra copy - pinning unpinned memory will make it contiguous
+                if not out.is_contiguous() and out.is_pinned():
+                    out = out.contiguous()
+                return out.pin_memory()
 
             dim = self.dim + (self.dim < 0) * len(batch[0].shape)
 

vllm/parser/minimax_m3.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MiniMax M3 parser for reasoning markers."""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
+
+from vllm.parser.engine.events import EventType
+from vllm.parser.engine.parser_engine import ParserEngine
+from vllm.parser.engine.parser_engine_config import (
+    ParserEngineConfig,
+    ParserState,
+    Transition,
+)
+
+if TYPE_CHECKING:
+    from vllm.tokenizers import TokenizerLike
+    from vllm.tool_parsers.abstract_tool_parser import Tool
+
+THINK_START = "<mm:think>"
+THINK_END = "</mm:think>"
+
+
+@functools.cache
+def minimax_m3_config(thinking: bool = False) -> ParserEngineConfig:
+    return ParserEngineConfig(
+        name="minimax_m3",
+        initial_state=ParserState.REASONING if thinking else ParserState.CONTENT,
+        terminals={
+            "THINK_START": THINK_START,
+            "THINK_END": THINK_END,
+        },
+        transitions={
+            (ParserState.CONTENT, "THINK_START"): Transition(
+                ParserState.REASONING,
+                (EventType.REASONING_START,),
+            ),
+            (ParserState.REASONING, "THINK_START"): Transition(
+                ParserState.REASONING,
+                (),
+            ),
+            (ParserState.REASONING, "THINK_END"): Transition(
+                ParserState.CONTENT,
+                (EventType.REASONING_END,),
+            ),
+            (ParserState.CONTENT, "THINK_END"): Transition(
+                ParserState.CONTENT,
+                (),
+            ),
+        },
+    )
+
+
+class MiniMaxM3Parser(ParserEngine):
+    """MiniMax M3 parser backed by the declarative parser engine."""
+
+    def __init__(
+        self,
+        tokenizer: TokenizerLike,
+        tools: list[Tool] | None = None,
+        **kwargs,
+    ) -> None:
+        chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
+        self._initial_in_reasoning = chat_kwargs.get("thinking_mode") == "enabled"
+        kwargs.setdefault(
+            "parser_engine_config",
+            minimax_m3_config(thinking=self._initial_in_reasoning),
+        )
+        super().__init__(tokenizer, tools, **kwargs)
+        self._start_token_ids = self._encode_marker(THINK_START)
+        self._end_token_ids = self._encode_marker(THINK_END)
+
+    def _encode_marker(self, marker: str) -> tuple[int, ...]:
+        try:
+            token_ids = self.model_tokenizer.encode(marker, add_special_tokens=False)
+        except TypeError:
+            token_ids = self.model_tokenizer.encode(marker)
+        return tuple(token_ids)
+
+    @staticmethod
+    def _contains_token_sequence(
+        token_ids: Sequence[int], marker_ids: Sequence[int]
+    ) -> bool:
+        if not marker_ids or len(marker_ids) > len(token_ids):
+            return False
+        marker_len = len(marker_ids)
+        return any(
+            tuple(token_ids[i : i + marker_len]) == tuple(marker_ids)
+            for i in range(len(token_ids) - marker_len + 1)
+        )
+
+    @staticmethod
+    def _rfind_token_sequence(
+        token_ids: Sequence[int], marker_ids: Sequence[int]
+    ) -> int:
+        if not marker_ids or len(marker_ids) > len(token_ids):
+            return -1
+        marker_len = len(marker_ids)
+        for i in range(len(token_ids) - marker_len, -1, -1):
+            if tuple(token_ids[i : i + marker_len]) == tuple(marker_ids):
+                return i
+        return -1
+
+    def is_reasoning_end(self, input_ids: list[int]) -> bool:
+        start_index = self._rfind_token_sequence(input_ids, self._start_token_ids)
+        end_index = self._rfind_token_sequence(input_ids, self._end_token_ids)
+        if end_index < 0:
+            return False
+        if start_index < 0:
+            return True
+        return end_index > start_index
+
+    def is_reasoning_end_streaming(
+        self, input_ids: Sequence[int], delta_ids: Iterable[int]
+    ) -> bool:
+        if self.reasoning_ended:
+            return True
+        if self._engine._lexer.buffer:
+            return False
+        if self._initial_in_reasoning:
+            return False
+        if self._engine.state == ParserState.CONTENT:
+            return bool(input_ids)
+        return False
+
+    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
+        end_index = self._rfind_token_sequence(input_ids, self._end_token_ids)
+        if end_index >= 0:
+            return input_ids[end_index + len(self._end_token_ids) :]
+
+        has_start = self._contains_token_sequence(input_ids, self._start_token_ids)
+        if self._initial_in_reasoning and not has_start:
+            return []
+
+        if not has_start:
+            return input_ids
+        return []
+
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        count = 0
+        depth = 1 if self._initial_in_reasoning else 0
+        i = 0
+        while i < len(token_ids):
+            if tuple(token_ids[i : i + len(self._start_token_ids)]) == (
+                self._start_token_ids
+            ):
+                depth += 1
+                i += len(self._start_token_ids)
+                continue
+            if tuple(token_ids[i : i + len(self._end_token_ids)]) == (
+                self._end_token_ids
+            ):
+                if depth > 0:
+                    depth -= 1
+                i += len(self._end_token_ids)
+                continue
+            if depth > 0:
+                count += 1
+            i += 1
+        return count

vllm/platforms/__init__.py

@@ -9,7 +9,6 @@ from typing import TYPE_CHECKING
 from vllm import envs
 from vllm.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group
 from vllm.utils.import_utils import resolve_obj_by_qualname
-from vllm.utils.torch_utils import supports_xccl
 
 from .interface import CpuArchEnum, Platform, PlatformEnum
 
@@ -135,7 +134,7 @@ def xpu_platform_plugin() -> str | None:
     try:
         import torch
 
-        if supports_xccl():
+        if torch.distributed.is_xccl_available():
             dist_backend = "xccl"
             from vllm.platforms.xpu import XPUPlatform
 

vllm/platforms/cuda.py

@@ -23,7 +23,6 @@ import vllm._C_stable_libtorch  # noqa
 import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.utils.import_utils import import_pynvml
-from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 from .interface import DeviceCapability, Platform, PlatformEnum, in_wsl
@@ -88,6 +87,8 @@ def _get_backend_priorities(
     kv_cache_dtype: CacheDType | None = None,
 ) -> list[AttentionBackendEnum]:
     """Get backend priorities with lazy import to avoid circular dependency."""
+    from vllm.utils.torch_utils import is_quantized_kv_cache
+
     if use_mla:
         if device_capability.major == 10:
             # Sparse MLA backend priorities

vllm/platforms/xpu.py

@@ -14,7 +14,6 @@ import vllm_xpu_kernels._xpu_C  # noqa
 
 import vllm.envs as envs
 from vllm.logger import init_logger
-from vllm.utils.torch_utils import supports_xpu_graph
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
 from .interface import DeviceCapability, Platform, PlatformEnum
@@ -178,8 +177,6 @@ class XPUPlatform(Platform):
 
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
-        parallel_config = vllm_config.parallel_config
-
         # lazy import to avoid circular import
         from vllm.config import CUDAGraphMode
 
@@ -190,6 +187,10 @@ class XPUPlatform(Platform):
         attention_config = vllm_config.attention_config
         if attention_config.backend is None:
             attention_config.backend = AttentionBackendEnum.FLASH_ATTN
+
+        # lazy import to avoid circular import
+        from vllm.utils.torch_utils import supports_xpu_graph
+
         if not supports_xpu_graph():
             compilation_config.cudagraph_mode = CUDAGraphMode.NONE
             logger.warning(
@@ -324,9 +325,8 @@ class XPUPlatform(Platform):
 
     @classmethod
     def get_device_communicator_cls(cls) -> str:
-        from vllm.utils.torch_utils import supports_xccl
-
-        if not supports_xccl():
+        if not torch.distributed.is_xccl_available():
+            # Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
             logger.warning(
                 "xccl is not enabled in this torch build, communication"
                 " is not available."

vllm/reasoning/minimax_m3_reasoning_parser.py

@@ -2,170 +2,19 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Iterable, Sequence
-from typing import TYPE_CHECKING
 
-from vllm.entrypoints.openai.engine.protocol import DeltaMessage
-from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
-
-if TYPE_CHECKING:
-    from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
-    from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.parser.engine.adapters import ParserEngineReasoningAdapter
+from vllm.parser.minimax_m3 import MiniMaxM3Parser
 
 
-class MiniMaxM3ReasoningParser(BaseThinkingReasoningParser):
-    """Reasoning parser for MiniMax M3 explicit thinking blocks.
+class MiniMaxM3ReasoningParser(ParserEngineReasoningAdapter):
+    """Reasoning parser adapter for MiniMax M3 explicit thinking blocks."""
 
-    MiniMax M3 emits reasoning as:
-
-        <mm:think>reasoning text</mm:think>assistant content
-
-    The M3 tokenizer exposes both markers as complete vocabulary tokens. The
-    chat template may also prefill the start marker when
-    ``thinking_mode="enabled"``, so generated text can begin directly inside a
-    reasoning block without emitting ``<mm:think>`` again.
-    """
-
-    @property
-    def start_token(self) -> str:
-        return "<mm:think>"
-
-    @property
-    def end_token(self) -> str:
-        return "</mm:think>"
-
-    def __init__(self, tokenizer, *args, **kwargs):
-        super().__init__(tokenizer, *args, **kwargs)
-        chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
-        self._initial_in_reasoning = chat_kwargs.get("thinking_mode") == "enabled"
-        self._at_response_start = True
-
-    def extract_reasoning(
-        self,
-        model_output: str,
-        request: "ChatCompletionRequest | ResponsesRequest",
-    ) -> tuple[str | None, str | None]:
-        # MiniMax M3 can start a response with a stray closer. Drop that first
-        # token only; later unmatched closers stay visible as content.
-        if not self._initial_in_reasoning and model_output.startswith(self.end_token):
-            content = model_output[len(self.end_token) :]
-            return None, content or None
-
-        if self._initial_in_reasoning and self.start_token not in model_output:
-            reasoning, end, content = model_output.partition(self.end_token)
-            if not end:
-                return model_output, None
-            return reasoning, content or None
-
-        if self.start_token not in model_output:
-            return None, model_output
-
-        content_before, _, after_start = model_output.partition(self.start_token)
-        reasoning, end, content_after = after_start.partition(self.end_token)
-        if not end:
-            return reasoning, content_before or None
-
-        return reasoning, (content_before + content_after) or None
+    _parser_engine_cls = MiniMaxM3Parser
 
     def is_reasoning_end_streaming(
         self, input_ids: Sequence[int], delta_ids: Iterable[int]
     ) -> bool:
-        delta_ids = tuple(delta_ids)
-        if self.end_token_id in delta_ids:
-            return True
-        if self.end_token_id in input_ids:
-            return True
-        if self._initial_in_reasoning:
-            return False
-        if self.start_token_id not in input_ids:
-            return bool(input_ids)
-        return False
-
-    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
-        if self.end_token_id in input_ids:
-            end_index = len(input_ids) - 1 - input_ids[::-1].index(self.end_token_id)
-            return input_ids[end_index + 1 :]
-
-        if self._initial_in_reasoning and self.start_token_id not in input_ids:
-            return []
-
-        if self.start_token_id not in input_ids:
-            return input_ids
-        return []
-
-    def extract_reasoning_streaming(
-        self,
-        previous_text: str,
-        current_text: str,
-        delta_text: str,
-        previous_token_ids: Sequence[int],
-        current_token_ids: Sequence[int],
-        delta_token_ids: Sequence[int],
-    ) -> DeltaMessage | None:
-        if not delta_text:
-            return None
-
-        if self._at_response_start and not self._initial_in_reasoning:
-            # Apply the leading-closer tolerance once. Later unmatched closers
-            # stay visible as content.
-            self._at_response_start = False
-            if delta_text.startswith(self.end_token):
-                delta_text = delta_text[len(self.end_token) :]
-                if not delta_text:
-                    return None
-                if delta_token_ids and delta_token_ids[0] == self.end_token_id:
-                    delta_token_ids = delta_token_ids[1:]
-
-        if self.end_token_id in previous_token_ids:
-            return DeltaMessage(content=delta_text)
-
-        if (
-            self._initial_in_reasoning
-            and self.start_token_id not in previous_token_ids
-            and self.start_token_id not in delta_token_ids
-        ):
-            if self.end_token_id in delta_token_ids:
-                reasoning, _, content = delta_text.partition(self.end_token)
-                return DeltaMessage(
-                    reasoning=reasoning or None,
-                    content=content or None,
-                )
-            return DeltaMessage(reasoning=delta_text)
-
-        if (
-            self.start_token_id not in previous_token_ids
-            and self.start_token_id not in delta_token_ids
-        ):
-            return DeltaMessage(content=delta_text)
-
-        if self.end_token_id in delta_token_ids:
-            reasoning_text, _, content = delta_text.partition(self.end_token)
-            if self.start_token_id in delta_token_ids:
-                _, _, reasoning_text = reasoning_text.partition(self.start_token)
-            return DeltaMessage(
-                reasoning=reasoning_text or None,
-                content=content or None,
-            )
-
-        if self.start_token_id in delta_token_ids:
-            _, _, reasoning = delta_text.partition(self.start_token)
-            return DeltaMessage(reasoning=reasoning) if reasoning else None
-
-        return DeltaMessage(reasoning=delta_text)
-
-    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
-        if not self._initial_in_reasoning:
-            return super().count_reasoning_tokens(token_ids)
-
-        count = 0
-        depth = 1
-        for token_id in token_ids:
-            if token_id == self.start_token_id:
-                depth += 1
-                continue
-            if token_id == self.end_token_id:
-                if depth > 0:
-                    depth -= 1
-                continue
-            if depth > 0:
-                count += 1
-        return count
+        return self._parser_engine.is_reasoning_end_streaming(
+            list(input_ids), tuple(delta_ids)
+        )

vllm/tool_parsers/__init__.py

@@ -15,7 +15,7 @@ Register a lazy module mapping.
 Example:
     ToolParserManager.register_lazy_module(
         name="kimi_k2",
-        module_path="vllm.tool_parsers.kimi_k2_parser",
+        module_path="vllm.tool_parsers.kimi_k2_tool_parser",
         class_name="KimiK2ToolParser",
     )
 """

vllm/utils/torch_utils.py

@@ -3,7 +3,6 @@
 import contextlib
 import importlib.metadata
 import os
-import platform
 import random
 import threading
 from collections.abc import Callable, Collection
@@ -18,6 +17,7 @@ from torch.library import Library, infer_schema
 
 import vllm.envs as envs
 from vllm.logger import init_logger
+from vllm.utils.platform_utils import is_pin_memory_available
 
 if TYPE_CHECKING:
     from vllm.config import ModelConfig
@@ -68,9 +68,7 @@ MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {
 T = TypeVar("T")
 
 
-# Pin memory in non-WSL case.
-# Logic duplicated here for now to avoid circular import.
-PIN_MEMORY = "microsoft" not in " ".join(platform.uname()).lower()
+PIN_MEMORY = is_pin_memory_available()
 
 
 def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
@@ -606,14 +604,24 @@ def create_kv_caches_with_random(
 
 
 def async_tensor_h2d(
-    data: list,
-    dtype: torch.dtype,
+    data: list | np.ndarray | torch.Tensor,
     device: str | torch.device,
-    pin_memory: bool = PIN_MEMORY,
+    dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
-    """Asynchronously create a tensor and copy it from host to device."""
-    t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu")
-    return t.to(device=device, non_blocking=True)
+    """Copy list/numpy array/tensor async from host to device."""
+    if isinstance(data, np.ndarray):
+        data = torch.from_numpy(data)
+    if isinstance(data, torch.Tensor):
+        t = data.pin_memory() if PIN_MEMORY else data
+    else:
+        t = torch.tensor(data, dtype=dtype, pin_memory=PIN_MEMORY, device="cpu")
+    assert t.is_cpu
+    return t.to(device=device, dtype=dtype, non_blocking=True)
+
+
+def np_to_pinned_tensor(array: np.ndarray) -> torch.Tensor:
+    t = torch.from_numpy(array)
+    return t.pin_memory() if PIN_MEMORY else t
 
 
 def make_ndarray_with_pad(
@@ -914,11 +922,6 @@ def _encode_layer_name(layer_name: str) -> str | LayerName:
     return LayerName(layer_name) if _USE_LAYERNAME else layer_name
 
 
-# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
-def supports_xccl() -> bool:
-    return torch.distributed.is_xccl_available()
-
-
 # Supports XPU Graph with PyTorch versions >= 2.11.0.dev for XPU platform
 def supports_xpu_graph() -> bool:
     return is_torch_equal_or_newer("2.11.0.dev")

vllm/v1/attention/backends/flashinfer.py

@@ -41,8 +41,8 @@ from vllm.utils.flashinfer import (
     use_trtllm_attention,
 )
 from vllm.utils.math_utils import cdiv
-from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.utils.torch_utils import (
+    PIN_MEMORY,
     canonicalize_singleton_dim_strides,
     is_quantized_kv_cache,
     is_strictly_contiguous,
@@ -708,9 +708,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
         # Since we do not have explicit synchronization in ModelRunnerV2, we do not pin
         # reused CPU buffers to avoid a race condition between step N async copies to
         # GPU and step N+1 buffer updates.
-        self.pin_memory = (
-            not vllm_config.use_v2_model_runner and is_pin_memory_available()
-        )
+        self.pin_memory = not vllm_config.use_v2_model_runner and PIN_MEMORY
         self.paged_kv_indptr = self._make_buffer(max_num_reqs + 1)
         self.paged_kv_indptr_cpu_buffer = torch.zeros_like(
             self.paged_kv_indptr.cpu, pin_memory=self.pin_memory

vllm/v1/attention/backends/flex_attention.py

@@ -28,7 +28,11 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
-from vllm.utils.torch_utils import is_quantized_kv_cache, is_torch_equal_or_newer
+from vllm.utils.torch_utils import (
+    async_tensor_h2d,
+    is_quantized_kv_cache,
+    is_torch_equal_or_newer,
+)
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -58,7 +62,7 @@ def _offsets_to_doc_ids_tensor(
     doc_ids = torch.repeat_interleave(
         torch.arange(len(counts), dtype=torch.int32), counts
     )
-    return doc_ids.to(device, non_blocking=True)
+    return async_tensor_h2d(doc_ids, device=device)
 
 
 def pad_to_multiple(x: torch.Tensor, multiple: int, dim: int):

vllm/v1/attention/backends/gdn_attn.py

@@ -8,6 +8,7 @@ from typing import Literal
 import torch
 
 from vllm.config import VllmConfig
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -203,8 +204,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
                 spec_sequence_masks = None
                 spec_sequence_masks_cpu = None
             else:
-                spec_sequence_masks = spec_sequence_masks_cpu.to(
-                    query_start_loc.device, non_blocking=True
+                spec_sequence_masks = async_tensor_h2d(
+                    spec_sequence_masks_cpu, device=query_start_loc.device
                 )
 
         if spec_sequence_masks is None:
@@ -376,12 +377,14 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
                 )
 
                 assert prefill_query_start_loc_cpu is not None
-                chunk_indices = prepare_chunk_indices(
-                    prefill_query_start_loc_cpu, FLA_CHUNK_SIZE
-                ).to(device=gpu_device, non_blocking=True)
-                chunk_offsets = prepare_chunk_offsets(
-                    prefill_query_start_loc_cpu, FLA_CHUNK_SIZE
-                ).to(device=gpu_device, non_blocking=True)
+                chunk_indices = async_tensor_h2d(
+                    prepare_chunk_indices(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
+                    device=gpu_device,
+                )
+                chunk_offsets = async_tensor_h2d(
+                    prepare_chunk_offsets(prefill_query_start_loc_cpu, FLA_CHUNK_SIZE),
+                    device=gpu_device,
+                )
 
         if num_prefills > 0:
             has_initial_state = context_lens_tensor > 0

vllm/v1/attention/backends/mamba2_attn.py

@@ -7,6 +7,7 @@ from typing import Any
 import torch
 
 from vllm.config import VllmConfig
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import (
     AttentionBackend,
     CommonAttentionMetadata,
@@ -68,22 +69,22 @@ def compute_varlen_chunk_metadata(
 
     # Exclusive prefix sum over logical-chunk lengths
     if chunk_lens:
-        cu_chunk_seqlens = torch.tensor(
-            [0] + list(itertools.accumulate(chunk_lens)),
-            device=device,
-            dtype=torch.int32,
-        )
-        # Final boundary must equal total tokens
-        assert int(cu_chunk_seqlens[-1].item()) == total
+        cu_chunk_seqlens_list = [0] + list(itertools.accumulate(chunk_lens))
+        # Final boundary must equal total tokens (check on host to avoid a sync)
+        assert cu_chunk_seqlens_list[-1] == total
     else:
-        cu_chunk_seqlens = torch.tensor([0], device=device, dtype=torch.int32)
+        cu_chunk_seqlens_list = [0]
+    cu_chunk_seqlens = async_tensor_h2d(
+        cu_chunk_seqlens_list, dtype=torch.int32, device=device
+    )
 
-    last_chunk_indices_t = (
-        torch.tensor(last_chunk_indices, device=device, dtype=torch.int32)
-        if len(starts) > 0
-        else torch.empty((0,), device=device, dtype=torch.int32)
+    # last_chunk_indices is empty when there are no sequences (len(starts) == 0).
+    last_chunk_indices_t = async_tensor_h2d(
+        last_chunk_indices, dtype=torch.int32, device=device
+    )
+    seq_idx_chunks_t = async_tensor_h2d(
+        seq_idx_chunks, dtype=torch.int32, device=device
     )
-    seq_idx_chunks_t = torch.tensor(seq_idx_chunks, device=device, dtype=torch.int32)
     return cu_chunk_seqlens, last_chunk_indices_t, seq_idx_chunks_t
 
 

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -26,7 +26,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
     get_mla_dims,
 )
 from vllm.platforms.interface import DeviceCapability
-from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -217,7 +217,7 @@ class FlashInferMLASparseMetadataBuilder(
         # Zero-fill for cudagraphs
         self.req_id_per_token_buffer.fill_(0)
         self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
-            torch.from_numpy(req_id_per_token), non_blocking=True
+            np_to_pinned_tensor(req_id_per_token), non_blocking=True
         )
         req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens]
 

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -16,7 +16,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
 from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
 from vllm.utils.platform_utils import num_compute_units
-from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -503,7 +503,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
         # Zero-fill for cudagraphs
         self.req_id_per_token_buffer.fill_(0)
         self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
-            torch.from_numpy(req_id_per_token), non_blocking=True
+            np_to_pinned_tensor(req_id_per_token), non_blocking=True
         )
         req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
 

vllm/v1/attention/backends/utils.py

@@ -17,7 +17,7 @@ from typing_extensions import runtime_checkable
 
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.utils.math_utils import cdiv
-from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor
 from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
 
 if TYPE_CHECKING:
@@ -364,8 +364,8 @@ def make_local_attention_virtual_batches(
     # tensor first, which recovers perf.
     # Upload the index tensors to the block_table's device up-front so that the
     # fancy indexing below doesn't implicitly force a synchronous H2D copy.
-    batch_indices_torch = torch.from_numpy(batch_indices).to(device, non_blocking=True)
-    block_indices_torch = torch.from_numpy(block_indices).to(device, non_blocking=True)
+    batch_indices_torch = async_tensor_h2d(batch_indices, device=device)
+    block_indices_torch = async_tensor_h2d(block_indices, device=device)
 
     # Save as a lambda so we can return this for update_block_table
     make_block_table = lambda block_table: block_table[
@@ -379,8 +379,8 @@ def make_local_attention_virtual_batches(
 
     return CommonAttentionMetadata(
         query_start_loc_cpu=query_start_loc_cpu,
-        query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True),
-        seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),
+        query_start_loc=async_tensor_h2d(query_start_loc_cpu, device=device),
+        seq_lens=async_tensor_h2d(seq_lens_cpu, device=device),
         num_reqs=len(seq_lens_cpu),
         num_actual_tokens=common_attn_metadata.num_actual_tokens,
         max_query_len=seqlens_q_local.max(),
@@ -808,14 +808,12 @@ def create_fast_prefill_custom_backend(
 
 
 def compute_causal_conv1d_metadata(
-    query_start_loc_p_cpu: torch.Tensor,
-    *,
-    device: torch.device,
-):
+    query_start_loc_p_cpu: torch.Tensor, *, device: torch.device
+) -> tuple[dict[int, dict[str, Any]], torch.Tensor, torch.Tensor]:
     # Needed for causal_conv1d. Use the CPU query_start_loc to avoid DtoH sync.
     assert query_start_loc_p_cpu.device.type == "cpu"
     seqlens = query_start_loc_p_cpu.diff()
-    nums_dict = {}  # type: ignore
+    nums_dict: dict[int, dict[str, Any]] = {}
     batch_ptr = None
     token_chunk_offset_ptr = None
     for BLOCK_M in [8]:  # cover all BLOCK_M values
@@ -823,7 +821,7 @@ def compute_causal_conv1d_metadata(
         nums_dict[BLOCK_M] = {}
         nums_dict[BLOCK_M]["nums"] = nums
         nums_dict[BLOCK_M]["tot"] = nums.sum().item()
-        mlist = torch.from_numpy(np.repeat(np.arange(len(nums)), nums))
+        mlist = np_to_pinned_tensor(np.repeat(np.arange(len(nums)), nums))
         nums_dict[BLOCK_M]["mlist"] = mlist
         mlist_len = len(nums_dict[BLOCK_M]["mlist"])
         nums_dict[BLOCK_M]["mlist_len"] = mlist_len
@@ -831,7 +829,7 @@ def compute_causal_conv1d_metadata(
         offsetlist = []  # type: ignore
         for idx, num in enumerate(nums):
             offsetlist.extend(range(num))
-        offsetlist = torch.tensor(offsetlist, dtype=torch.int32)
+        offsetlist = torch.tensor(offsetlist, dtype=torch.int32, pin_memory=PIN_MEMORY)
         nums_dict[BLOCK_M]["offsetlist"] = offsetlist
 
         if batch_ptr is None:
@@ -845,16 +843,15 @@ def compute_causal_conv1d_metadata(
         else:
             if batch_ptr.nelement() < MAX_NUM_PROGRAMS:
                 batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)
-                token_chunk_offset_ptr.resize_(  # type: ignore
-                    MAX_NUM_PROGRAMS
-                ).fill_(PAD_SLOT_ID)
+                assert token_chunk_offset_ptr is not None
+                token_chunk_offset_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)
 
+        assert batch_ptr is not None
         batch_ptr[0:mlist_len].copy_(mlist, non_blocking=True)
-        token_chunk_offset_ptr[  # type: ignore
-            0:mlist_len
-        ].copy_(offsetlist, non_blocking=True)
+        assert token_chunk_offset_ptr is not None
+        token_chunk_offset_ptr[0:mlist_len].copy_(offsetlist, non_blocking=True)
         nums_dict[BLOCK_M]["batch_ptr"] = batch_ptr
-        nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr  # type: ignore
+        nums_dict[BLOCK_M]["token_chunk_offset_ptr"] = token_chunk_offset_ptr
 
     return nums_dict, batch_ptr, token_chunk_offset_ptr
 

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

@@ -58,7 +58,9 @@ def _indexer_k_quant_and_cache_kernel(
     slot_id = tl.load(slot_mapping_ptr + tid)
     if slot_id < 0:
         return
-    block_id = slot_id // block_size
+    # The packed KV layout makes per-block strides large
+    # enough that block_id * stride can exceed 32-bit range.
+    block_id = (slot_id // block_size).to(tl.int64)
     block_offset = slot_id % block_size
     tile_block_id = block_offset // BLOCK_TILE_SIZE
     tile_block_offset = block_offset % BLOCK_TILE_SIZE
@@ -179,7 +181,9 @@ def _cp_gather_indexer_quant_cache_kernel(
         block_table_ptr + block_table_offset, mask=valid_block_table, other=-1
     )
     valid_block = valid_block_table & (block_id >= 0) & (block_id < NUM_BLOCKS)
-    safe_block_id = tl.where(valid_block, block_id, 0)
+    # The packed KV layout makes per-block strides large
+    # enough that block_id * stride can exceed 32-bit range.
+    safe_block_id = tl.where(valid_block, block_id, 0).to(tl.int64)
     safe_block_offset = tl.where(valid_block, block_offset, 0)
     tiled_block_offset = safe_block_offset % BLOCK_TILE_SIZE
     if LAYOUT == "SHUFFLE":

vllm/v1/core/kv_cache_utils.py

@@ -938,9 +938,7 @@ def _pool_bytes_per_block(kv_cache_groups: list[KVCacheGroupSpec]) -> int:
         kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
     ):
         return kv_cache_groups[0].kv_cache_spec.page_size_bytes
-    if all(
-        isinstance(g.kv_cache_spec, UniformTypeKVCacheSpecs) for g in kv_cache_groups
-    ):
+    if _use_packed_kv_cache_groups(kv_cache_groups):
         # buckets = {page_size: [[layer_names], [layer_names], ...]}
         buckets = _bucket_layers_by_page_size(kv_cache_groups)
         return sum(ps * len(slots) for ps, slots in buckets.items())
@@ -1218,16 +1216,29 @@ def _bucket_layers_by_page_size(
     return buckets
 
 
-def _get_kv_cache_config_deepseek_v4(
+def _use_packed_kv_cache_groups(
+    kv_cache_groups: list[KVCacheGroupSpec],
+) -> bool:
+    is_dsv4 = all(
+        isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
+        for group in kv_cache_groups
+    )
+    return is_dsv4 or (
+        bool(envs.VLLM_USE_PACKED_HMA_KV_CACHE) and len(kv_cache_groups) > 1
+    )
+
+
+def _get_kv_cache_config_packed(
     vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
     available_memory: int,
 ) -> tuple[int, list[KVCacheTensor]]:
-    """DeepseekV4 KV cache tensor layout planning.
+    """Plan a packed per-block KV cache tensor layout.
 
     Emit one KVCacheTensor per (slot_idx, page_size). Layers from different
     groups at the same slot share a tensor (they have independent block
-    tables so block-id namespaces never collide).
+    tables so block-id namespaces never collide). Each emitted tensor aliases
+    one physical backing allocation, with per-block data laid out contiguously.
     """
     # buckets = {page_size: [[layer_names], [layer_names], ...]}
     buckets = _bucket_layers_by_page_size(kv_cache_groups)
@@ -1255,6 +1266,9 @@ def _get_kv_cache_config_deepseek_v4(
     return num_blocks, kv_cache_tensors
 
 
+_get_kv_cache_config_deepseek_v4 = _get_kv_cache_config_packed
+
+
 def get_kv_cache_config_from_groups(
     vllm_config: VllmConfig,
     kv_cache_groups: list[KVCacheGroupSpec],
@@ -1299,13 +1313,11 @@ def get_kv_cache_config_from_groups(
             )
             for layer_name in kv_cache_groups[0].layer_names
         ]
-    elif all(
-        isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
-        for group in kv_cache_groups
-    ):
-        # DeepseekV4: UniformTypeKVCacheSpecs but multiple groups.
-        # Delegate to the DeepseekV4-specific allocator.
-        num_blocks, kv_cache_tensors = _get_kv_cache_config_deepseek_v4(
+    elif _use_packed_kv_cache_groups(kv_cache_groups):
+        # DeepSeek V4 keeps the existing packed layout. Other multi-group
+        # attention-only HMA layouts can opt in with
+        # VLLM_USE_PACKED_HMA_KV_CACHE=1.
+        num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
             vllm_config, kv_cache_groups, available_memory
         )
     else:

vllm/v1/core/sched/output.py

@@ -118,8 +118,8 @@ class CachedRequestData:
     # NOTE(woosuk): new_token_ids is only used for pipeline parallelism.
     # When PP is not used, new_token_ids will be empty.
     new_token_ids: list[list[int]]
-    # For requests not scheduled in the last step, propagate the token ids to the
-    # connector. Won't contain requests that were scheduled in the prior step.
+    # MRV1-only: For requests not scheduled in the last step, propagate the token ids
+    # to the connector. Won't contain requests scheduled in the prior step.
     all_token_ids: dict[str, list[int]]
     new_block_ids: list[tuple[list[int], ...] | None]
     num_computed_tokens: list[int]

vllm/v1/core/sched/scheduler.py

@@ -101,6 +101,7 @@ class Scheduler(SchedulerInterface):
         self.finished_req_ids_dict: dict[int, set[str]] | None = (
             defaultdict(set) if include_finished_set else None
         )
+        # Track requests scheduled in prior step (MRV1-only).
         self.prev_step_scheduled_req_ids: set[str] = set()
 
         # Scheduling constraints.
@@ -1010,8 +1011,8 @@ class Scheduler(SchedulerInterface):
 
         # Construct the scheduler output.
         if self.use_v2_model_runner:
-            scheduled_new_reqs = scheduled_new_reqs + scheduled_resumed_reqs
-            scheduled_resumed_reqs = []
+            scheduled_new_reqs.extend(scheduled_resumed_reqs)
+            scheduled_resumed_reqs.clear()
             new_reqs_data = [
                 NewRequestData.from_request(
                     req,
@@ -1037,9 +1038,10 @@ class Scheduler(SchedulerInterface):
                 req_to_new_blocks,
             )
 
-        # Record the request ids that were scheduled in this step.
-        self.prev_step_scheduled_req_ids.clear()
-        self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
+        # Record the request ids that were scheduled in this step (MRV1-only).
+        if not self.use_v2_model_runner:
+            self.prev_step_scheduled_req_ids.clear()
+            self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
 
         new_block_ids_to_zero = (
             (self.kv_cache_manager.take_new_block_ids() or None)
@@ -1252,12 +1254,11 @@ class Scheduler(SchedulerInterface):
                     req.num_computed_tokens : req.num_computed_tokens + num_tokens
                 ]
                 new_token_ids.append(token_ids)
-            scheduled_in_prev_step = req_id in self.prev_step_scheduled_req_ids
             if idx >= num_running_reqs:
-                assert not scheduled_in_prev_step
                 resumed_req_ids.add(req_id)
-            if not scheduled_in_prev_step:
-                all_token_ids[req_id] = req.all_token_ids.copy()
+            if not self.use_v2_model_runner:  # noqa: SIM102
+                if req_id not in self.prev_step_scheduled_req_ids:
+                    all_token_ids[req_id] = req.all_token_ids.copy()
             new_block_ids.append(
                 req_to_new_blocks[req_id].get_block_ids(allow_none=True)
             )

vllm/v1/kv_offload/cpu/common.py

@@ -4,7 +4,10 @@ from typing_extensions import override
 
 from vllm.v1.kv_offload.base import BlockIDsLoadStoreSpec
 
-METRIC_STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
+
+class CPUOffloadingMetrics:
+    STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
+    CPU_CACHE_USAGE_PERC = "vllm:kv_offload_cpu_cache_usage_perc"
 
 
 class CPULoadStoreSpec(BlockIDsLoadStoreSpec):

vllm/v1/kv_offload/cpu/gpu_worker.py

@@ -14,7 +14,7 @@ from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.triton_utils import HAS_TRITON, triton
 from vllm.utils.math_utils import cdiv
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.kv_offload.base import (
     BlockIDsLoadStoreSpec,
     CanonicalKVCacheRef,
@@ -156,7 +156,7 @@ def pin_mmap_region(region: SharedOffloadRegion) -> None:
 def _new_descriptor_buffers(
     num_copy_ops: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    pin = is_pin_memory_available()
+    pin = PIN_MEMORY
     # CUDA cache_kernels.cu requires int64; XPU DMA engine requires uint64.
     ptr_dtype = torch.uint64 if current_platform.is_xpu() else torch.int64
     return (
@@ -482,7 +482,7 @@ class CpuGpuOffloadingHandlers:
         num_cpu_blocks: int,
         mmap_region: SharedOffloadRegion | None = None,
     ):
-        pin_memory = is_pin_memory_available()
+        pin_memory = PIN_MEMORY
         logger.info("Allocating %d CPU tensors...", len(kv_caches.tensors))
         self._mmap_region = mmap_region
         if mmap_region is not None and pin_memory:

vllm/v1/kv_offload/cpu/manager.py

@@ -18,7 +18,10 @@ from vllm.v1.kv_offload.base import (
     ReqContext,
     RequestOffloadingContext,
 )
-from vllm.v1.kv_offload.cpu.common import METRIC_STORES_SKIPPED, CPULoadStoreSpec
+from vllm.v1.kv_offload.cpu.common import (
+    CPULoadStoreSpec,
+    CPUOffloadingMetrics,
+)
 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
 from vllm.v1.kv_offload.cpu.policies.base import BlockStatus, CachePolicy
 from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy
@@ -282,13 +285,21 @@ class CPUOffloadingManager(OffloadingManager):
             self.events.clear()
 
     def get_stats(self) -> OffloadingConnectorStats | None:
-        if self.store_threshold < 2:
-            return None
-
         stats = OffloadingConnectorStats()
-        stats.increase_counter(
-            METRIC_STORES_SKIPPED,
-            self.stores_skipped_in_current_batch,
+
+        # Compute cache usage.
+        num_used = (
+            self._num_allocated_blocks
+            - len(self._free_list)
+            - self._num_evictable_cache_blocks
         )
-        self.stores_skipped_in_current_batch = 0
+        usage = num_used / self._num_blocks if self._num_blocks > 0 else 0.0
+        stats.set_gauge(CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC, usage)
+
+        if self.store_threshold >= 2:
+            stats.increase_counter(
+                CPUOffloadingMetrics.STORES_SKIPPED,
+                self.stores_skipped_in_current_batch,
+            )
+            self.stores_skipped_in_current_batch = 0
         return stats

vllm/v1/kv_offload/cpu/spec.py

@@ -14,11 +14,15 @@ from vllm.v1.kv_offload.base import (
     GPULoadStoreSpec,
     LoadStoreSpec,
     OffloadingCounterMetadata,
+    OffloadingGaugeMetadata,
     OffloadingManager,
     OffloadingMetricMetadata,
     OffloadingSpec,
 )
-from vllm.v1.kv_offload.cpu.common import METRIC_STORES_SKIPPED, CPULoadStoreSpec
+from vllm.v1.kv_offload.cpu.common import (
+    CPULoadStoreSpec,
+    CPUOffloadingMetrics,
+)
 from vllm.v1.kv_offload.cpu.gpu_worker import CpuGpuOffloadingHandlers
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
 from vllm.v1.kv_offload.worker.worker import OffloadingHandler
@@ -31,17 +35,27 @@ class CPUOffloadingSpec(OffloadingSpec):
     def build_metric_definitions(
         cls, extra_config: dict[str, Any]
     ) -> dict[str, OffloadingMetricMetadata]:
-        store_threshold = int(extra_config.get("store_threshold", 0))
-        if store_threshold < 2:
-            return {}
-        return {
-            METRIC_STORES_SKIPPED: OffloadingCounterMetadata(
+        definitions: dict[str, OffloadingMetricMetadata] = {
+            CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC: OffloadingGaugeMetadata(
                 documentation=(
-                    "Number of KV offload stores skipped because the reuse "
-                    "threshold was not reached."
+                    "Fraction of CPU KV-cache space currently pinned by active "
+                    "transfers (0.0 = idle, 1.0 = saturated). Sustained high "
+                    "values indicate transfers (stores or promotions) may be "
+                    "dropped due to insufficient capacity."
                 ),
             )
         }
+        store_threshold = int(extra_config.get("store_threshold", 0))
+        if store_threshold >= 2:
+            definitions[CPUOffloadingMetrics.STORES_SKIPPED] = (
+                OffloadingCounterMetadata(
+                    documentation=(
+                        "Number of KV offload stores skipped because the reuse "
+                        "threshold was not reached."
+                    ),
+                )
+            )
+        return definitions
 
     def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
         super().__init__(vllm_config, kv_cache_config)
@@ -58,7 +72,15 @@ class CPUOffloadingSpec(OffloadingSpec):
         self.cpu_page_size_per_worker = 0
         assert kv_cache_config is not None
         if kv_cache_config.num_blocks > 0 and world_size > 0:
-            total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
+            is_packed = any(t.block_stride for t in kv_cache_config.kv_cache_tensors)
+            assert not is_packed or all(
+                t.block_stride for t in kv_cache_config.kv_cache_tensors
+            )
+            total_gpu_kv_bytes = (
+                kv_cache_config.kv_cache_tensors[0].size
+                if is_packed
+                else sum(t.size for t in kv_cache_config.kv_cache_tensors)
+            )
             kv_bytes_per_block = (
                 total_gpu_kv_bytes // kv_cache_config.num_blocks
             ) * world_size

vllm/v1/pool/metadata.py

@@ -7,9 +7,7 @@ import torch
 
 from vllm.pooling_params import PoolingParams
 from vllm.tasks import PoolingTask
-from vllm.utils.platform_utils import is_pin_memory_available
-
-pin_memory = is_pin_memory_available()
+from vllm.utils.torch_utils import PIN_MEMORY
 
 
 @dataclass
@@ -134,7 +132,7 @@ class PoolingMetadata:
         num_scheduled_tokens_cpu = torch.from_numpy(num_scheduled_tokens_np)
         if query_start_loc_gpu is None:
             cumsum = torch.zeros(
-                n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu"
+                n_seq + 1, dtype=torch.int64, pin_memory=PIN_MEMORY, device="cpu"
             )
             torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
             cumsum = cumsum.to(device, non_blocking=True)

vllm/v1/sample/logits_processor/builtin.py

@@ -7,6 +7,7 @@ import numpy as np
 import torch
 
 from vllm import SamplingParams
+from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.sample.logits_processor.interface import (
     BatchUpdate,
     LogitsProcessor,
@@ -118,7 +119,6 @@ class MinPLogitsProcessor(LogitsProcessor):
 class LogitBiasLogitsProcessor(LogitsProcessor):
     def __init__(self, _, device: torch.device, is_pin_memory: bool):
         self.device = device
-        self.pin_memory = is_pin_memory
         self.biases: dict[int, dict[int, float]] = {}
 
         self.bias_tensor: torch.Tensor = torch.tensor(())
@@ -154,9 +154,7 @@ class LogitBiasLogitsProcessor(LogitsProcessor):
             )
 
     def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor:
-        return torch.tensor(
-            data, device="cpu", dtype=dtype, pin_memory=self.pin_memory
-        ).to(device=self.device, non_blocking=True)
+        return async_tensor_h2d(data, device=self.device, dtype=dtype)
 
     def apply(self, logits: torch.Tensor) -> torch.Tensor:
         if self.biases:
@@ -170,7 +168,6 @@ class MinTokensLogitsProcessor(LogitsProcessor):
     ):
         # index -> (min_toks, output_token_ids, stop_token_ids)
         self.device = device
-        self.pin_memory = is_pin_memory
         self.min_toks: dict[int, tuple[int, Sequence[int], set[int]]] = {}
 
         # (req_idx_tensor,eos_tok_id_tensor)
@@ -227,9 +224,7 @@ class MinTokensLogitsProcessor(LogitsProcessor):
             )
 
     def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor:
-        return torch.tensor(
-            data, device="cpu", dtype=dtype, pin_memory=self.pin_memory
-        ).to(device=self.device, non_blocking=True)
+        return async_tensor_h2d(data, device=self.device, dtype=dtype)
 
     def apply(self, logits: torch.Tensor) -> torch.Tensor:
         if self.min_toks:
@@ -283,8 +278,8 @@ class MinTokensLogitsProcessor(LogitsProcessor):
             toks_arr = np.concatenate(all_toks)
             # (row_indices, token_indices) for index_put_ to set -inf.
             logits_slice = (
-                torch.from_numpy(rows_arr).to(self.device, non_blocking=True),
-                torch.from_numpy(toks_arr).to(self.device, non_blocking=True),
+                async_tensor_h2d(rows_arr, device=self.device),
+                async_tensor_h2d(toks_arr, device=self.device),
             )
             logits.index_put_(logits_slice, self.neg_inf_tensor)
 

vllm/v1/sample/ops/penalties.py

@@ -4,8 +4,7 @@
 import torch
 
 from vllm.model_executor.layers.utils import apply_penalties
-from vllm.utils.platform_utils import is_pin_memory_available
-from vllm.utils.torch_utils import make_tensor_with_pad
+from vllm.utils.torch_utils import PIN_MEMORY, make_tensor_with_pad
 
 
 def apply_all_penalties(
@@ -52,6 +51,6 @@ def _convert_to_tensors(
         pad=vocab_size,
         device="cpu",
         dtype=torch.int64,
-        pin_memory=is_pin_memory_available(),
+        pin_memory=PIN_MEMORY,
     )
     return output_tokens_tensor.to(device, non_blocking=True)

vllm/v1/sample/sampler.py

@@ -6,7 +6,7 @@ import torch
 import torch.nn as nn
 
 from vllm.config.model import LogprobsMode
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.outputs import LogprobsTensors, SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.ops.bad_words import apply_bad_words
@@ -65,7 +65,7 @@ class Sampler(nn.Module):
     ):
         super().__init__()
         self.topk_topp_sampler = TopKTopPSampler(logprobs_mode, use_fp64_gumbel)
-        self.pin_memory = is_pin_memory_available()
+        self.pin_memory = PIN_MEMORY
         self.logprobs_mode = logprobs_mode
         self.use_fp64_gumbel = use_fp64_gumbel
 

vllm/v1/sample/thinking_budget_state.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 
 from vllm.platforms import current_platform
-from vllm.utils.torch_utils import async_tensor_h2d
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
 from vllm.v1.sample.logits_processor.interface import (
     BatchUpdate,
     MoveDirectionality,
@@ -22,12 +22,11 @@ def maybe_create_thinking_budget_state_holder(
     max_num_seqs: int,
     num_spec_tokens: int,
     device: torch.device,
-    is_pin_memory: bool,
 ) -> "ThinkingBudgetStateHolder | None":
     if reasoning_config is None:
         return None
     return ThinkingBudgetStateHolder(
-        reasoning_config, max_num_seqs, num_spec_tokens, device, is_pin_memory
+        reasoning_config, max_num_seqs, num_spec_tokens, device, PIN_MEMORY
     )
 
 

vllm/v1/serial_utils.py

@@ -33,7 +33,7 @@ from vllm.multimodal.inputs import (
     MultiModalSharedField,
     NestedTensors,
 )
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.utils import tensor_data
 
 logger = init_logger(__name__)
@@ -327,7 +327,7 @@ class MsgpackDecoder:
         oob_tensor_provider: OOBTensorProvider | None = None,
     ):
         self.share_mem = share_mem
-        self.pin_tensors = is_pin_memory_available()
+        self.pin_tensors = PIN_MEMORY
         args = () if t is None else (t,)
         self.decoder = msgpack.Decoder(
             *args, ext_hook=self.ext_hook, dec_hook=self.dec_hook

vllm/v1/simple_kv_offload/manager.py

@@ -82,6 +82,9 @@ class SimpleCPUOffloadScheduler:
             vllm_config.kv_events_config is not None
             and vllm_config.kv_events_config.enable_kv_cache_events
         )
+        dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
+        pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
+        self.cp_world_size = dcp_world_size * pcp_world_size
         self.block_size = scheduler_block_size
         self.hash_block_size = hash_block_size
         assert self.block_size % self.hash_block_size == 0
@@ -113,9 +116,6 @@ class SimpleCPUOffloadScheduler:
         )
 
         # TODO (yifan): maybe need to enable kv_cache_events and metrics_collector here.
-        dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
-        pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
-        assert dcp_world_size == 1 and pcp_world_size == 1
         self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator(
             kv_cache_config=self.cpu_kv_cache_config,
             max_model_len=vllm_config.model_config.max_model_len,
@@ -155,6 +155,7 @@ class SimpleCPUOffloadScheduler:
             self._target_free = self._estimate_lazy_target_blocks(
                 kv_cache_config,
                 vllm_config.scheduler_config.max_num_batched_tokens,
+                self.cp_world_size,
             )
         else:
             self._target_free = 0
@@ -187,7 +188,13 @@ class SimpleCPUOffloadScheduler:
 
         assert len(gpu_config.kv_cache_tensors) > 0
 
-        gpu_total_bytes = sum(t.size for t in gpu_config.kv_cache_tensors)
+        is_packed = any(t.block_stride for t in gpu_config.kv_cache_tensors)
+        assert not is_packed or all(t.block_stride for t in gpu_config.kv_cache_tensors)
+        gpu_total_bytes = (
+            gpu_config.kv_cache_tensors[0].size
+            if is_packed
+            else sum(t.size for t in gpu_config.kv_cache_tensors)
+        )
         num_gpu_blocks = gpu_config.num_blocks
         num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes)
         # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally.
@@ -195,6 +202,8 @@ class SimpleCPUOffloadScheduler:
             KVCacheTensor(
                 size=t.size // num_gpu_blocks * num_cpu_blocks,
                 shared_by=list(t.shared_by),
+                offset=t.offset,
+                block_stride=t.block_stride,
             )
             for t in gpu_config.kv_cache_tensors
         ]
@@ -207,19 +216,22 @@ class SimpleCPUOffloadScheduler:
 
     @staticmethod
     def _estimate_lazy_target_blocks(
-        kv_cache_config: "KVCacheConfig", max_num_batched_tokens: int
+        kv_cache_config: "KVCacheConfig",
+        max_num_batched_tokens: int,
+        cp_world_size: int = 1,
     ) -> int:
         """GPU blocks to keep available (free/offloaded) per step in lazy mode."""
         WATERMARK_RATIO = 1.0  # Reserve larger space to avoid running out of GPU blocks
         target = 0
         for g in kv_cache_config.kv_cache_groups:
             spec = g.kv_cache_spec
+            block_size = spec.block_size * cp_world_size
             if isinstance(spec, MambaSpec):
                 target += 2
             elif isinstance(spec, SlidingWindowSpec):
-                target += cdiv(spec.sliding_window, spec.block_size) + 1
+                target += cdiv(spec.sliding_window, block_size) + 1
             else:
-                target += cdiv(max_num_batched_tokens, spec.block_size)
+                target += cdiv(max_num_batched_tokens, block_size)
         return int(target * (1 + WATERMARK_RATIO))
 
     def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:
@@ -355,7 +367,9 @@ class SimpleCPUOffloadScheduler:
                 continue
 
             # Number of blocks in the computed range for this group.
-            g_block_size = kv_cache_groups[g].kv_cache_spec.block_size
+            g_block_size = (
+                kv_cache_groups[g].kv_cache_spec.block_size * self.cp_world_size
+            )
             n_computed_g = cdiv(total_computed_tokens, g_block_size)
 
             # Back-trace: ext blocks sit at the tail of the computed range.

vllm/v1/simple_kv_offload/worker.py

@@ -8,7 +8,7 @@ import torch
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
 from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
 from vllm.v1.simple_kv_offload.metadata import (
@@ -149,7 +149,7 @@ class SimpleCPUOffloadWorker:
             (self.num_cpu_blocks * total_bytes_per_block) / (1024**3),
         )
 
-        pin_memory = is_pin_memory_available()
+        pin_memory = PIN_MEMORY
         if not pin_memory:
             logger.warning(
                 "Pinned memory not available. CPU offload performance may be degraded."

vllm/v1/spec_decode/extract_hidden_states.py

@@ -12,7 +12,7 @@ from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.model_loader import get_model
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata
 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
 from vllm.v1.utils import CpuGpuBuffer
@@ -58,7 +58,7 @@ class ExtractHiddenStatesProposer:
         self.backup_next_token_ids = CpuGpuBuffer(
             max_batch_size,
             dtype=torch.int32,
-            pin_memory=is_pin_memory_available(),
+            pin_memory=PIN_MEMORY,
             device=device,
             with_numpy=True,
         )
@@ -317,7 +317,6 @@ class ExtractHiddenStatesProposer:
         (batch_size, 1). For each request we either use the sampled token
         (if valid and not discarded) or a backup token from the request state.
         """
-        num_reqs = gpu_input_batch.num_reqs
 
         # Precompute backup token IDs for discarded requests.
         num_reqs = gpu_input_batch.num_reqs

vllm/v1/spec_decode/llm_base_proposer.py

@@ -26,7 +26,7 @@ from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
 from vllm.model_executor.models.qwen3_dflash import DFlashQwen3ForCausalLM
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.platforms import current_platform
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
@@ -228,7 +228,7 @@ class SpecDecodeBaseProposer:
         self.backup_next_token_ids = CpuGpuBuffer(
             self.max_batch_size,
             dtype=torch.int32,
-            pin_memory=is_pin_memory_available(),
+            pin_memory=PIN_MEMORY,
             device=device,
             with_numpy=True,
         )
@@ -239,9 +239,7 @@ class SpecDecodeBaseProposer:
         self._last_draft_probs: torch.Tensor | None = None
 
         self._slot_mapping_buffer = torch.zeros(
-            self.max_positions,
-            dtype=torch.int64,
-            device=device,
+            self.max_positions, dtype=torch.int64, device=device
         )
 
         # Determine allowed attention backends once during initialization.
@@ -1127,7 +1125,7 @@ class SpecDecodeBaseProposer:
         new_query_start_loc_cpu = torch.zeros(
             query_start_loc_cpu.shape,
             dtype=torch.int32,
-            pin_memory=is_pin_memory_available(),
+            pin_memory=PIN_MEMORY,
         )
         new_query_start_loc_np = new_query_start_loc_cpu.numpy()
         np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:])
@@ -1160,11 +1158,11 @@ class SpecDecodeBaseProposer:
         #  q1 + 0, q1 + 1, q1 + 2, q1 + 3,       // req 2
         #  q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3
         token_indices_np = token_offsets + old_query_start_locs_expanded
-        token_indices = torch.from_numpy(token_indices_np).to(device, non_blocking=True)
+        token_indices = async_tensor_h2d(token_indices_np, device=device)
 
         spec_common_attn_metadata = CommonAttentionMetadata(
-            query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True),
-            seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),
+            query_start_loc=async_tensor_h2d(new_query_start_loc_cpu, device=device),
+            seq_lens=async_tensor_h2d(new_seq_lens_cpu, device=device),
             query_start_loc_cpu=new_query_start_loc_cpu,
             _seq_lens_cpu=new_seq_lens_cpu,
             _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,

vllm/v1/spec_decode/ngram_proposer_gpu.py

@@ -545,7 +545,7 @@ def update_ngram_gpu_tensors_incremental(
             num_tokens = input_batch.num_tokens_no_spec[idx]
             if num_tokens > 0:
                 token_ids_gpu_tensor[idx, :num_tokens].copy_(
-                    input_batch.token_ids_cpu_tensor[idx, :num_tokens],
+                    input_batch.token_ids_cpu_tensor[idx, :num_tokens].pin_memory(),
                     non_blocking=True,
                 )
 
@@ -591,7 +591,7 @@ def update_ngram_gpu_tensors_incremental(
         num_tokens = input_batch.num_tokens_no_spec[new_req_idx]
         if num_tokens > 0:
             token_ids_gpu_tensor[new_req_idx, :num_tokens].copy_(
-                input_batch.token_ids_cpu_tensor[new_req_idx, :num_tokens],
+                input_batch.token_ids_cpu_tensor[new_req_idx, :num_tokens].pin_memory(),
                 non_blocking=True,
             )
 

vllm/v1/structured_output/backend_lm_format_enforcer.py

@@ -11,7 +11,7 @@ from transformers import PreTrainedTokenizerBase
 
 from vllm.sampling_params import SamplingParams
 from vllm.utils.import_utils import LazyLoader
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY
 from vllm.v1.structured_output.backend_types import (
     StructuredOutputBackend,
     StructuredOutputGrammar,
@@ -139,7 +139,7 @@ class LMFormatEnforcerBackend(StructuredOutputBackend):
             (max_num_seqs, (self.vocab_size + 31) // 32),
             -1,
             dtype=torch.int32,
-            pin_memory=is_pin_memory_available(),
+            pin_memory=PIN_MEMORY,
         )
 
     def destroy(self):