[Bugfix][MultiConnector] Fix MultiConnector for SupportsHMA sub-connectors (#36549)

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>

.github/mergify.yml

@@ -27,7 +27,7 @@ pull_request_rules:
         Hi @{{author}}, the pre-commit checks have failed. Please run:
 
         ```bash 
-        uv pip install pre-commit
+        uv pip install pre-commit>=4.5.1
         pre-commit install
         pre-commit run --all-files
         ```

benchmarks/attention_benchmarks/runner.py

@@ -418,8 +418,8 @@ def _run_single_benchmark(
     mem_stats = {}
     if config.profile_memory:
         mem_stats = {
-            "allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
-            "reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
+            "allocated_mb": torch.accelerator.memory_allocated(device) / 1024**2,
+            "reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
         }
 
     return times, mem_stats

benchmarks/benchmark_topk_topp.py

@@ -95,13 +95,16 @@ def create_logits(
 def measure_memory() -> tuple[int, int]:
     """Return (allocated, reserved) memory in bytes."""
     torch.accelerator.synchronize()
-    return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
+    return (
+        torch.accelerator.memory_allocated(),
+        torch.accelerator.max_memory_allocated(),
+    )
 
 
 def reset_memory_stats():
     """Reset peak memory statistics."""
     reset_buffer_cache()
-    torch.cuda.reset_peak_memory_stats()
+    torch.accelerator.reset_peak_memory_stats()
     torch.accelerator.empty_cache()
     gc.collect()
 

csrc/cuda_vec_utils.cuh

@@ -196,6 +196,7 @@ __forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
   return val;
 #else
   assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
+  return u32x8_t{};
 #endif
 }
 

csrc/cumem_allocator.cpp

@@ -109,16 +109,18 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
 
 #ifndef USE_ROCM
   int flag = 0;
-  CUDA_CHECK(cuDeviceGetAttribute(
+  CUresult rdma_result = cuDeviceGetAttribute(
       &flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
-      device));
-  if (flag) {  // support GPUDirect RDMA if possible
+      device);
+  if (rdma_result == CUDA_SUCCESS &&
+      flag) {  // support GPUDirect RDMA if possible
     prop.allocFlags.gpuDirectRDMACapable = 1;
   }
   int fab_flag = 0;
-  CUDA_CHECK(cuDeviceGetAttribute(
-      &fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device));
-  if (fab_flag) {  // support fabric handle if possible
+  CUresult fab_result = cuDeviceGetAttribute(
+      &fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
+  if (fab_result == CUDA_SUCCESS &&
+      fab_flag) {  // support fabric handle if possible
     prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
   }
 #endif

csrc/moe/moe_permute_unpermute_op.cu

@@ -73,10 +73,9 @@ void moe_permute(
   MOE_DISPATCH(input.scalar_type(), [&] {
     expandInputRowsKernelLauncher<scalar_t>(
         get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
-        get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
-        get_ptr<int>(inv_permuted_idx), get_ptr<int>(permuted_idx),
-        get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
-        n_hidden, topk, n_local_expert, stream);
+        get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
+        get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
+        n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
   });
 }
 

csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h

@@ -57,7 +57,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
 
 template <typename T>
 void expandInputRowsKernelLauncher(
-    T const* unpermuted_input, T* permuted_output, int* sorted_experts,
+    T const* unpermuted_input, T* permuted_output,
     int const* expanded_dest_row_to_expanded_source_row,
     int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
     int64_t const* expert_first_token_offset, int64_t const num_rows,

csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.inl

@@ -2,7 +2,7 @@
 
 template <typename T, bool CHECK_SKIPPED>
 __global__ void expandInputRowsKernel(
-    T const* unpermuted_input, T* permuted_output, int* sorted_experts,
+    T const* unpermuted_input, T* permuted_output,
     int const* expanded_dest_row_to_expanded_source_row,
     int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
     int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -16,7 +16,6 @@ __global__ void expandInputRowsKernel(
   int64_t expanded_dest_row = blockIdx.x;
   int64_t const expanded_source_row =
       expanded_dest_row_to_expanded_source_row[expanded_dest_row];
-  int expert_id = sorted_experts[expanded_dest_row];
 
   if (threadIdx.x == 0) {
     assert(expanded_dest_row <= INT32_MAX);
@@ -54,7 +53,7 @@ __global__ void expandInputRowsKernel(
 
 template <typename T>
 void expandInputRowsKernelLauncher(
-    T const* unpermuted_input, T* permuted_output, int* sorted_experts,
+    T const* unpermuted_input, T* permuted_output,
     int const* expanded_dest_row_to_expanded_source_row,
     int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
     int64_t const* expert_first_token_offset, int64_t const num_rows,
@@ -70,12 +69,12 @@ void expandInputRowsKernelLauncher(
   bool is_check_skip = num_valid_tokens_ptr != nullptr;
   auto func = func_map[is_check_skip];
 
-  func<<<blocks, threads, 0, stream>>>(
-      unpermuted_input, permuted_output, sorted_experts,
-      expanded_dest_row_to_expanded_source_row,
-      expanded_source_row_to_expanded_dest_row, permuted_idx,
-      expert_first_token_offset, num_rows, num_valid_tokens_ptr, cols, k,
-      num_local_experts);
+  func<<<blocks, threads, 0, stream>>>(unpermuted_input, permuted_output,
+                                       expanded_dest_row_to_expanded_source_row,
+                                       expanded_source_row_to_expanded_dest_row,
+                                       permuted_idx, expert_first_token_offset,
+                                       num_rows, num_valid_tokens_ptr, cols, k,
+                                       num_local_experts);
 }
 
 template <class T, class U>

csrc/sampler.cu

@@ -575,7 +575,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
   // The range of logits within the row.
   int rowStart = 0;
   int seq_len = seqLens[rowIdx / next_n];
-  int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1;
+  int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1);
 
   // Local pointers to this block
   if constexpr (!multipleBlocksPerRow && !mergeBlocks) {

docker/Dockerfile

@@ -620,7 +620,7 @@ RUN set -eux; \
 ARG BITSANDBYTES_VERSION_X86=0.46.1
 ARG BITSANDBYTES_VERSION_ARM64=0.42.0
 ARG TIMM_VERSION=">=1.0.17"
-ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.3"
+ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"
 RUN --mount=type=cache,target=/root/.cache/uv \
     if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
         BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_ARM64}"; \
@@ -628,7 +628,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
         BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
     fi; \
     uv pip install --system accelerate hf_transfer modelscope \
-        "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs]${RUNAI_MODEL_STREAMER_VERSION}"
+        "bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
 
 # ============================================================
 # VLLM INSTALLATION (depends on build stage)

docker/Dockerfile.cpu

@@ -9,6 +9,7 @@
 #
 # Build targets:
 #   vllm-openai (default): used for serving deployment
+#   vllm-openai-zen: vLLM from source + zentorch from PyPI via vllm[zen]
 #   vllm-test: used for CI tests
 #   vllm-dev: used for development
 #
@@ -222,3 +223,19 @@ LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
 LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
 
 ENTRYPOINT ["vllm", "serve"]
+
+
+######################### ZEN CPU PYPI IMAGE #########################
+FROM vllm-openai AS vllm-openai-zen
+
+ARG TARGETARCH
+
+RUN if [ "$TARGETARCH" != "amd64" ]; then \
+        echo "ERROR: vllm-openai-amd only supports --platform=linux/amd64"; \
+        exit 1; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install "vllm[zen]"
+
+ENTRYPOINT ["vllm", "serve"]

docker/versions.json

@@ -83,7 +83,7 @@
       "default": ">=1.0.17"
     },
     "RUNAI_MODEL_STREAMER_VERSION": {
-      "default": ">=0.15.3"
+      "default": ">=0.15.7"
     }
   }
 }

docs/contributing/README.md

@@ -75,7 +75,7 @@ For an optimized workflow when iterating on C++/CUDA kernels, see the [Increment
 vLLM uses `pre-commit` to lint and format the codebase. See <https://pre-commit.com/#usage> if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:
 
 ```bash
-uv pip install pre-commit
+uv pip install pre-commit>=4.5.1
 pre-commit install
 ```
 

docs/design/attention_backends.md

@@ -164,18 +164,18 @@ Priority is **1 = highest** (tried first).
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
-| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
-| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
-| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
+| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
+| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
-| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
+| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
-| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
-| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
-| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
+| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
+| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
+| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
 
 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
 >
@@ -204,14 +204,14 @@ configuration.
 
 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
-| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
-| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
+| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
-| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
-| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
-| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
+| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
+| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
-| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
-| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
+| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
+| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |

docs/design/moe_kernel_features.md

@@ -35,7 +35,8 @@ th {
 | naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
 | deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
 | deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
-| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |
+| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
+| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize.FlashInferNVLinkOneSidedPrepareAndFinalize] |
 
 !!! info "Table key"
     1. All types: mxfp4, nvfp4, int4, int8, fp8

docs/design/torch_compile_multimodal.md

@@ -34,9 +34,6 @@ relies on caching artifacts to reduce start time, we must properly propagate the
 with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder
 components (see Compile Range Integration).
 
-3. `with set_forward_context` context manager should be used around the nn.Module's forward call. This will properly forward the vllm_config which is needed
-for torch.compile integration.
-
 ### CompilationConfig
 
 With the exception of `compile_mm_encoder: true`, the multimodal encoder will inherit from the same compilation config as the text LLM. We may extend

docs/features/tool_calling.md

@@ -219,7 +219,7 @@ Supported models:
 
 * `ibm-granite/granite-4.0-h-small` and other Granite 4.0 models
 
-    Recommended flags: `--tool-call-parser hermes`
+    Recommended flags: `--tool-call-parser granite4`
 
 * `ibm-granite/granite-3.0-8b-instruct`
 

docs/getting_started/installation/README.md

@@ -16,4 +16,6 @@ vLLM supports the following hardware platforms:
 
 vLLM supports third-party hardware plugins that live **outside** the main `vllm` repository. These follow the [Hardware-Pluggable RFC](../../design/plugin_system.md).
 
-A list of all supported hardware can be found on the [vllm.ai website](https://vllm.ai/#compatibility). If you want to add new hardware, please contact us on [Slack](https://slack.vllm.ai/) or [Email](mailto:collaboration@vllm.ai).
+A list of all supported hardware can be found on the vLLM website, see [Universal Compatibility - Hardware](https://vllm.ai/#compatibility).
+
+If you want to add new hardware, please contact us on [Slack](https://slack.vllm.ai/) or [Email](mailto:collaboration@vllm.ai).

docs/models/extensions/runai_model_streamer.md

@@ -31,6 +31,16 @@ vllm serve gs://core-llm/Llama-3-8b \
     --load-format runai_streamer
 ```
 
+To run model from Azure Blob Storage run:
+
+```bash
+AZURE_STORAGE_ACCOUNT_NAME=<account> \
+vllm serve az://<container>/<model-path> \
+    --load-format runai_streamer
+```
+
+Authentication uses `DefaultAzureCredential`, which supports `az login`, managed identity, environment variables (`AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_CLIENT_SECRET`), and other methods.
+
 To run model from a S3 compatible object store run:
 
 ```bash

docs/models/supported_models.md

@@ -418,6 +418,7 @@ th {
 | `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
+| `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
 | `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
 | `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |

docs/serving/expert_parallel_deployment.md

@@ -21,7 +21,8 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
 | `allgather_reducescatter` | Default backend | Standard all2all using allgather/reducescatter primitives | General purpose, works with any EP+DP configuration |
 | `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios |
 | `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios |
-| `flashinfer_all2allv` | MNNVL systems | FlashInfer alltoallv kernels for multi-node NVLink | Systems with NVLink across nodes |
+| `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads |
+| `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
 | `naive` | Testing/debugging | Simple broadcast-based implementation | Debugging, not recommended for production |
 
 ## Single Node Deployment

examples/online_serving/disaggregated_serving/moriio_toy_proxy_server.py

@@ -14,6 +14,10 @@ import regex as re
 import zmq
 from quart import Quart, make_response, request
 
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
+    MoRIIOConstants,
+)
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 prefill_instances: list[dict] = []
@@ -213,6 +217,8 @@ async def handle_request():
 
         dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
 
+        transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"
+
         req_data_to_prefill = copy.deepcopy(req_data)
         req_data_to_prefill["kv_transfer_params"] = {}
         req_data["kv_transfer_params"] = {}
@@ -222,6 +228,7 @@ async def handle_request():
         req_data_to_prefill["kv_transfer_params"]["remote_tp_size"] = (
             decode_instance_endpoint["tp_size"]
         )
+        req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id
 
         send_prefill_task = asyncio.create_task(
             send_request_to_prefill(
@@ -267,6 +274,7 @@ async def handle_request():
 
         if selected_prefill_dp_rank is not None:
             req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
+        req_data["kv_transfer_params"]["transfer_id"] = transfer_id
 
         decode_request_task = asyncio.create_task(
             start_decode_request(

requirements/common.txt

@@ -24,7 +24,7 @@ outlines_core == 0.2.11
 # required for outlines backend disk cache
 diskcache == 5.6.3
 lark == 1.2.2
-xgrammar == 0.1.29; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
+xgrammar >= 0.1.32, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
 typing_extensions >= 4.10
 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
 partial-json-parser # used for parsing partial JSON outputs

requirements/lint.txt

@@ -1,2 +1,2 @@
 # formatting
-pre-commit==4.0.1
+pre-commit>=4.5.1

requirements/nightly_torch_test.txt

@@ -42,7 +42,7 @@ tritonclient>=2.51.0
 
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
-runai-model-streamer[s3,gcs]==0.15.3
+runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2
 instanttensor>=0.1.5
 pydantic>=2.12 # 2.11 leads to error on python 3.13

requirements/rocm-test.txt

@@ -45,6 +45,8 @@ pystemmer==3.0.0
     # via mteb
 
 # Multi-modal processing
+av==16.1.0
+    # required for audio_in_video tests
 blobfile==3.0.0
     # Multi-Modal Models Test
 decord==0.6.0

requirements/rocm.txt

@@ -15,7 +15,7 @@ tensorizer==2.10.1
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8
-runai-model-streamer[s3,gcs]==0.15.3
+runai-model-streamer[s3,gcs,azure]==0.15.7
 conch-triton-kernels==1.2.1
 timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 

requirements/test.in

@@ -10,6 +10,7 @@ pytest-cov
 
 # testing utils
 albumentations # required for Nemotron Parse in test_common.py
+av  # required for audio_in_video tests
 backoff # required for phi4mm test
 blobfile # required for kimi-vl test
 einops # required for MPT, qwen-vl
@@ -55,7 +56,7 @@ grpcio-reflection==1.78.0
 arctic-inference == 0.1.1 # Required for suffix decoding test
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
-runai-model-streamer[s3,gcs]==0.15.3
+runai-model-streamer[s3,gcs,azure]==0.15.7
 fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
 instanttensor>=0.1.5
 pydantic>=2.12 # 2.11 leads to error on python 3.13

requirements/test.txt

@@ -62,6 +62,16 @@ attrs==24.2.0
     #   referencing
 audioread==3.0.1
     # via librosa
+av==16.1.0
+    # via -r requirements/test.in
+azure-core==1.38.2
+    # via
+    #   azure-identity
+    #   azure-storage-blob
+azure-identity==1.25.2
+    # via runai-model-streamer-azure
+azure-storage-blob==12.28.0
+    # via runai-model-streamer-azure
 backoff==2.2.1
     # via
     #   -r requirements/test.in
@@ -101,8 +111,10 @@ certifi==2024.8.30
     #   rasterio
     #   requests
     #   sentry-sdk
-cffi==1.17.1
-    # via soundfile
+cffi==2.0.0
+    # via
+    #   cryptography
+    #   soundfile
 chardet==5.2.0
     # via mbstrdecoder
 charset-normalizer==3.4.0
@@ -146,6 +158,12 @@ coverage==7.10.6
     # via pytest-cov
 cramjam==2.9.0
     # via fastparquet
+cryptography==46.0.5
+    # via
+    #   azure-identity
+    #   azure-storage-blob
+    #   msal
+    #   pyjwt
 cuda-bindings==12.9.4
     # via torch
 cuda-pathfinder==1.3.3
@@ -377,6 +395,8 @@ iniconfig==2.0.0
     # via pytest
 instanttensor==0.1.5
     # via -r requirements/test.in
+isodate==0.7.2
+    # via azure-storage-blob
 isoduration==20.11.0
     # via jsonschema
 isort==5.13.2
@@ -490,6 +510,12 @@ more-itertools==10.5.0
     # via lm-eval
 mpmath==1.3.0
     # via sympy
+msal==1.34.0
+    # via
+    #   azure-identity
+    #   msal-extensions
+msal-extensions==1.3.1
+    # via azure-identity
 msgpack==1.1.0
     # via
     #   librosa
@@ -826,6 +852,8 @@ pydantic-extra-types==2.10.5
     # via mistral-common
 pygments==2.18.0
     # via rich
+pyjwt==2.11.0
+    # via msal
 pyogrio==0.11.0
     # via geopandas
 pyparsing==3.2.0
@@ -943,6 +971,7 @@ regex==2024.9.11
     #   transformers
 requests==2.32.3
     # via
+    #   azure-core
     #   buildkite-test-collector
     #   datasets
     #   diffusers
@@ -955,6 +984,7 @@ requests==2.32.3
     #   lightly
     #   lm-eval
     #   mistral-common
+    #   msal
     #   mteb
     #   pooch
     #   ray
@@ -991,11 +1021,13 @@ rsa==4.9.1
     # via google-auth
 rtree==1.4.0
     # via torchgeo
-runai-model-streamer==0.15.3
+runai-model-streamer==0.15.7
     # via -r requirements/test.in
-runai-model-streamer-gcs==0.15.3
+runai-model-streamer-azure==0.15.7
     # via runai-model-streamer
-runai-model-streamer-s3==0.15.3
+runai-model-streamer-gcs==0.15.7
+    # via runai-model-streamer
+runai-model-streamer-s3==0.15.7
     # via runai-model-streamer
 s3transfer==0.10.3
     # via boto3
@@ -1264,6 +1296,9 @@ typing-extensions==4.15.0
     #   aiosignal
     #   albumentations
     #   alembic
+    #   azure-core
+    #   azure-identity
+    #   azure-storage-blob
     #   chz
     #   fastapi
     #   grpcio

setup.py

@@ -657,13 +657,18 @@ class precompiled_wheel_utils:
     def get_base_commit_in_main_branch() -> str:
         try:
             # Get the latest commit hash of the upstream main branch.
-            resp_json = subprocess.check_output(
-                [
-                    "curl",
-                    "-s",
-                    "https://api.github.com/repos/vllm-project/vllm/commits/main",
+            curl_cmd = [
+                "curl",
+                "-s",
+                "https://api.github.com/repos/vllm-project/vllm/commits/main",
+            ]
+            github_token = os.getenv("GH_TOKEN", os.getenv("GITHUB_TOKEN"))
+            if github_token:
+                curl_cmd += [
+                    "-H",
+                    f"Authorization: token {github_token}",
                 ]
-            ).decode("utf-8")
+            resp_json = subprocess.check_output(curl_cmd).decode("utf-8")
             upstream_main_commit = json.loads(resp_json)["sha"]
             print(f"Upstream main branch latest commit: {upstream_main_commit}")
 
@@ -966,18 +971,19 @@ setup(
     ext_modules=ext_modules,
     install_requires=get_requirements(),
     extras_require={
+        # AMD Zen CPU optimizations via zentorch
+        "zen": ["zentorch"],
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.2.2"],
         "instanttensor": ["instanttensor >= 0.1.5"],
-        "runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
+        "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
         "audio": [
             "librosa",
             "scipy",
             "soundfile",
             "mistral_common[audio]",
             "av",
-            "torchcodec",
         ],  # Required for audio processing
         "video": [],  # Kept for backwards compatibility
         "flashinfer": [],  # Kept for backwards compatibility

tests/entrypoints/anthropic/test_anthropic_messages_conversion.py

@@ -4,6 +4,9 @@
 
 Tests the image source handling and tool_result content parsing in
 AnthropicServingMessages._convert_anthropic_to_openai_request().
+
+Also covers extended-thinking edge cases such as ``redacted_thinking``
+blocks echoed back by Anthropic clients.
 """
 
 from vllm.entrypoints.anthropic.protocol import (
@@ -373,3 +376,262 @@ class TestAttributionHeaderStripping:
         result = _convert(request)
         system_msg = result.messages[0]
         assert system_msg["content"] == "You are a helpful assistant."
+
+
+# ======================================================================
+# Thinking block conversion (Anthropic → OpenAI)
+# ======================================================================
+
+
+class TestThinkingBlockConversion:
+    """Verify that thinking blocks in assistant messages are correctly
+    moved to the ``reasoning`` field and stripped from ``content`` during
+    the Anthropic→OpenAI conversion.
+
+    This is the Anthropic-endpoint path: the client echoes back the full
+    assistant message (including thinking blocks emitted by vllm) in
+    subsequent requests.
+    """
+
+    def test_thinking_plus_text_in_assistant_message(self):
+        """thinking + text → reasoning field + plain-string content."""
+        request = _make_request(
+            [
+                {"role": "user", "content": "Write me some code."},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "thinking",
+                            "thinking": "I should write a simple example.",
+                            "signature": "sig_abc123",
+                        },
+                        {"type": "text", "text": "Sure! Here is the code."},
+                    ],
+                },
+                {"role": "user", "content": "Can you fix the bug?"},
+            ]
+        )
+        result = _convert(request)
+
+        # Find the assistant message in the converted output.
+        asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
+        assert len(asst_msgs) == 1
+        asst = asst_msgs[0]
+
+        # Thinking content must be in reasoning, NOT in content.
+        assert asst.get("reasoning") == "I should write a simple example."
+        assert asst.get("content") == "Sure! Here is the code."
+
+    def test_thinking_only_in_assistant_message(self):
+        """Assistant message with only a thinking block (no visible text).
+
+        This can happen when the model emits reasoning but no final answer
+        yet (e.g. a mid-turn reasoning step).  Content should be None.
+        """
+        request = _make_request(
+            [
+                {"role": "user", "content": "Hello"},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "thinking",
+                            "thinking": "Just thinking...",
+                            "signature": "sig_xyz",
+                        }
+                    ],
+                },
+                {"role": "user", "content": "Go on."},
+            ]
+        )
+        result = _convert(request)
+
+        asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
+        assert len(asst_msgs) == 1
+        asst = asst_msgs[0]
+
+        assert asst.get("reasoning") == "Just thinking..."
+        # No visible text → content should be absent or None.
+        assert asst.get("content") is None
+
+    def test_thinking_plus_tool_use_in_assistant_message(self):
+        """thinking + tool_use: reasoning field set, tool_calls populated."""
+        request = _make_request(
+            [
+                {"role": "user", "content": "What is 2+2?"},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "thinking",
+                            "thinking": "I need to call the calculator.",
+                            "signature": "sig_tool",
+                        },
+                        {
+                            "type": "tool_use",
+                            "id": "call_001",
+                            "name": "calculator",
+                            "input": {"expression": "2+2"},
+                        },
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "tool_result",
+                            "tool_use_id": "call_001",
+                            "content": "4",
+                        }
+                    ],
+                },
+            ]
+        )
+        result = _convert(request)
+
+        asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
+        assert len(asst_msgs) == 1
+        asst = asst_msgs[0]
+
+        assert asst.get("reasoning") == "I need to call the calculator."
+        tool_calls = list(asst.get("tool_calls", []))
+        assert len(tool_calls) == 1
+        assert tool_calls[0]["function"]["name"] == "calculator"
+        # No text content alongside reasoning + tool_use.
+        assert asst.get("content") is None
+
+    def test_multiple_thinking_blocks_concatenated(self):
+        """Multiple thinking blocks should be joined in order."""
+        request = _make_request(
+            [
+                {"role": "user", "content": "Think hard."},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "thinking",
+                            "thinking": "First thought. ",
+                            "signature": "s1",
+                        },
+                        {
+                            "type": "thinking",
+                            "thinking": "Second thought.",
+                            "signature": "s2",
+                        },
+                        {"type": "text", "text": "Done."},
+                    ],
+                },
+            ]
+        )
+        result = _convert(request)
+
+        asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
+        assert len(asst_msgs) == 1
+        asst = asst_msgs[0]
+
+        assert asst.get("reasoning") == "First thought. Second thought."
+        assert asst.get("content") == "Done."
+
+    def test_no_thinking_blocks_unchanged(self):
+        """Messages without thinking blocks must not be modified."""
+        request = _make_request(
+            [
+                {"role": "user", "content": "Hi"},
+                {"role": "assistant", "content": "Hello!"},
+            ]
+        )
+        result = _convert(request)
+
+        asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
+        assert len(asst_msgs) == 1
+        asst = asst_msgs[0]
+
+        assert asst.get("content") == "Hello!"
+        assert "reasoning" not in asst
+
+    def test_multi_turn_with_thinking_blocks(self):
+        """Full multi-turn conversation: previous assistant messages that
+        include thinking blocks must all be converted without a 400 error.
+
+        This is the primary regression scenario from the bug report:
+        upgrading vllm from v0.15.1 → v0.17.0 introduced thinking-block
+        support in responses, but echoing those responses back in subsequent
+        requests caused a Pydantic validation failure.
+        """
+        request = _make_request(
+            [
+                {"role": "user", "content": "Turn 1 question"},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "thinking",
+                            "thinking": "Reasoning for turn 1.",
+                            "signature": "s_t1",
+                        },
+                        {"type": "text", "text": "Answer for turn 1."},
+                    ],
+                },
+                {"role": "user", "content": "Turn 2 question"},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "thinking",
+                            "thinking": "Reasoning for turn 2.",
+                            "signature": "s_t2",
+                        },
+                        {"type": "text", "text": "Answer for turn 2."},
+                    ],
+                },
+                {"role": "user", "content": "Turn 3 question"},
+            ]
+        )
+        # Must not raise a ValidationError / 400.
+        result = _convert(request)
+
+        asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
+        assert len(asst_msgs) == 2
+
+        assert asst_msgs[0].get("reasoning") == "Reasoning for turn 1."
+        assert asst_msgs[0].get("content") == "Answer for turn 1."
+        assert asst_msgs[1].get("reasoning") == "Reasoning for turn 2."
+        assert asst_msgs[1].get("content") == "Answer for turn 2."
+
+    def test_redacted_thinking_block_is_accepted(self):
+        """Anthropic clients may echo back redacted thinking blocks.
+
+        vLLM should accept these blocks (to avoid 400 validation errors)
+        and ignore them when constructing the OpenAI-format prompt.
+        """
+        request = _make_request(
+            [
+                {"role": "user", "content": "Hello"},
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "type": "thinking",
+                            "thinking": "Thinking...",
+                            "signature": "sig_think",
+                        },
+                        {
+                            "type": "redacted_thinking",
+                            "data": "BASE64_OR_OTHER_OPAQUE_DATA",
+                        },
+                        {"type": "text", "text": "Hi!"},
+                    ],
+                },
+                {"role": "user", "content": "Continue"},
+            ]
+        )
+        result = _convert(request)
+
+        asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
+        assert len(asst_msgs) == 1
+        asst = asst_msgs[0]
+
+        # Redacted thinking is ignored, normal thinking still becomes reasoning.
+        assert asst.get("reasoning") == "Thinking..."
+        assert asst.get("content") == "Hi!"

tests/entrypoints/openai/responses/test_simple.py

@@ -137,6 +137,59 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str):
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+async def test_streaming_logprobs(client: OpenAI, model_name: str):
+    """Test that streaming with logprobs returns valid logprob data on
+    output_text.delta events and that top_logprobs has the requested count."""
+    response = await client.responses.create(
+        model=model_name,
+        input="Say hello.",
+        stream=True,
+        top_logprobs=3,
+        include=["message.output_text.logprobs"],
+    )
+
+    events = []
+    async for event in response:
+        events.append(event)
+
+    assert len(events) > 0
+
+    # Collect all output_text.delta events that carry logprobs
+    text_delta_events = [e for e in events if e.type == "response.output_text.delta"]
+    assert len(text_delta_events) > 0, "Expected at least one text delta event"
+
+    for delta_event in text_delta_events:
+        logprobs = delta_event.logprobs
+        assert logprobs is not None, "logprobs should be present on text delta events"
+        assert len(logprobs) > 0, "logprobs list should not be empty"
+        for lp in logprobs:
+            # Each logprob entry must have a token and a logprob value
+            assert lp.token is not None
+            assert isinstance(lp.logprob, float)
+            assert lp.logprob <= 0.0, f"logprob should be <= 0, got {lp.logprob}"
+            # top_logprobs should have up to 3 entries
+            assert lp.top_logprobs is not None
+            assert len(lp.top_logprobs) <= 3
+            for tl in lp.top_logprobs:
+                assert tl.token is not None
+                assert isinstance(tl.logprob, float)
+
+    # Verify that top_logprobs are actually populated, not always empty
+    all_top_logprobs = [
+        tl for e in text_delta_events for lp in e.logprobs for tl in lp.top_logprobs
+    ]
+    assert len(all_top_logprobs) > 0, (
+        "Expected at least one top_logprobs entry across all delta events"
+    )
+
+    # Verify the completed event still has valid output
+    completed = events[-1]
+    assert completed.type == "response.completed"
+    assert completed.response.status == "completed"
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_streaming_reasoning_tokens_e2e(client: OpenAI, model_name: str):

tests/entrypoints/openai/test_audio_in_video.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import base64
+import json
+
+import openai
+import pytest
+import pytest_asyncio
+
+from ...conftest import VideoTestAssets
+from ...utils import RemoteOpenAIServer
+
+MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
+
+
+@pytest.fixture
+def server():
+    args = [
+        "--max-model-len",
+        "16384",
+        "--enforce-eager",
+        "--limit-mm-per-prompt",
+        json.dumps({"audio": 3, "video": 3}),
+    ]
+
+    with RemoteOpenAIServer(
+        MODEL_NAME,
+        args,
+    ) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def client(server):
+    async with server.get_async_client() as async_client:
+        yield async_client
+
+
+@pytest.mark.core_model
+@pytest.mark.asyncio
+async def test_online_audio_in_video(
+    client: openai.AsyncOpenAI, video_assets: VideoTestAssets
+):
+    """Test video input with `audio_in_video=True`"""
+
+    # we don't use video_urls above because they missed audio stream.
+    video_path = video_assets[0].video_path
+    with open(video_path, "rb") as f:
+        video_base64 = base64.b64encode(f.read()).decode("utf-8")
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What's in this video?"},
+                {
+                    "type": "video_url",
+                    "video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
+                },
+            ],
+        }
+    ]
+
+    # multi-turn to test mm processor cache as well
+    for _ in range(2):
+        chat_completion = await client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=messages,
+            max_tokens=16,
+            extra_body={
+                "mm_processor_kwargs": {
+                    "use_audio_in_video": True,
+                }
+            },
+        )
+
+        assert len(chat_completion.choices) == 1
+        choice = chat_completion.choices[0]
+        assert choice.finish_reason == "length"
+
+
+@pytest.mark.core_model
+@pytest.mark.asyncio
+async def test_online_audio_in_video_multi_videos(
+    client: openai.AsyncOpenAI, video_assets: VideoTestAssets
+):
+    """Test multi-video input with `audio_in_video=True`"""
+
+    # we don't use video_urls above because they missed audio stream.
+    video_path = video_assets[0].video_path
+    with open(video_path, "rb") as f:
+        video_base64 = base64.b64encode(f.read()).decode("utf-8")
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What's in these two videos?"},
+                {
+                    "type": "video_url",
+                    "video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
+                },
+                {
+                    "type": "video_url",
+                    "video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
+                },
+            ],
+        }
+    ]
+
+    # multi-turn to test mm processor cache as well
+    for _ in range(2):
+        chat_completion = await client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=messages,
+            max_tokens=16,
+            extra_body={
+                "mm_processor_kwargs": {
+                    "use_audio_in_video": True,
+                }
+            },
+        )
+
+        assert len(chat_completion.choices) == 1
+        choice = chat_completion.choices[0]
+        assert choice.finish_reason == "length"
+
+
+@pytest.mark.core_model
+@pytest.mark.asyncio
+async def test_online_audio_in_video_interleaved(
+    client: openai.AsyncOpenAI, video_assets: VideoTestAssets
+):
+    """Test interleaved video/audio input with `audio_in_video=True`"""
+
+    # we don't use video_urls above because they missed audio stream.
+    video_path = video_assets[0].video_path
+    with open(video_path, "rb") as f:
+        video_base64 = base64.b64encode(f.read()).decode("utf-8")
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What's in these two videos?"},
+                {
+                    "type": "video_url",
+                    "video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
+                },
+                {
+                    "type": "audio_url",
+                    "audio_url": {"url": f"data:audio/mp4;base64,{video_base64}"},
+                },
+                {
+                    "type": "video_url",
+                    "video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
+                },
+            ],
+        }
+    ]
+    with pytest.raises(
+        openai.BadRequestError,
+        match="use_audio_in_video requires equal number of audio and video items",
+    ):
+        await client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=messages,
+            max_tokens=16,
+            extra_body={
+                "mm_processor_kwargs": {
+                    "use_audio_in_video": True,
+                }
+            },
+        )

tests/entrypoints/openai/test_gptoss_structural_tags_integration.py

@@ -1,279 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Integration tests for GPT-OSS structural tags functionality (PR #25515)."""
-
-import json
-from unittest.mock import Mock
-
-import pytest
-
-from vllm.entrypoints.mcp.tool_server import ToolServer
-from vllm.reasoning.gptoss_reasoning_parser import (
-    GptOssReasoningParser,
-)
-from vllm.sampling_params import StructuredOutputsParams
-
-
-class TestGptOssStructuralTagsIntegration:
-    """Integration tests for structural tags in GPT-OSS tool calls."""
-
-    @pytest.fixture
-    def mock_tokenizer(self):
-        """Create a mock tokenizer."""
-        tokenizer = Mock()
-        tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
-        tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
-        return tokenizer
-
-    @pytest.fixture
-    def gptoss_parser(self, mock_tokenizer):
-        """Create a real GptOssReasoningParser instance."""
-        return GptOssReasoningParser(mock_tokenizer)
-
-    @pytest.fixture
-    def tool_server_with_python(self):
-        """Create a tool server with Python tool enabled."""
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python")
-        return tool_server
-
-    @pytest.fixture
-    def tool_server_empty(self):
-        """Create a tool server with no tools."""
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(return_value=False)
-        return tool_server
-
-    def test_end_to_end_no_tools(self, gptoss_parser):
-        """Test end-to-end flow when no tools are available."""
-        # Test the parser directly
-        result = gptoss_parser.prepare_structured_tag(None, None)
-        parsed_result = json.loads(result)
-
-        # Verify basic structure
-        assert parsed_result["type"] == "structural_tag"
-        assert parsed_result["format"]["type"] == "triggered_tags"
-        assert len(parsed_result["format"]["tags"]) == 1
-
-        # Verify only analysis channel is allowed
-        analysis_tag = parsed_result["format"]["tags"][0]
-        assert analysis_tag["begin"] == "<|channel|>analysis<|message|>"
-        assert analysis_tag["content"]["type"] == "any_text"
-        assert analysis_tag["end"] == "<|end|>"
-
-        # Verify triggers
-        assert parsed_result["format"]["triggers"] == ["<|channel|>analysis"]
-        assert parsed_result["format"]["stop_after_first"] is False
-
-    def test_end_to_end_with_python_tool(self, gptoss_parser, tool_server_with_python):
-        """Test end-to-end flow with Python tool enabled."""
-        result = gptoss_parser.prepare_structured_tag(None, tool_server_with_python)
-        parsed_result = json.loads(result)
-
-        # Should have analysis tag + 2 python tags
-        assert len(parsed_result["format"]["tags"]) == 3
-
-        # Verify all expected tags are present
-        tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]]
-        expected_begins = [
-            "<|channel|>analysis<|message|>",
-            "<|channel|>commentary to=python",
-            "<|channel|>analysis to=python",
-        ]
-
-        for expected in expected_begins:
-            assert expected in tag_begins
-
-        # Verify triggers include commentary
-        assert "<|channel|>analysis" in parsed_result["format"]["triggers"]
-        assert "<|channel|>commentary to=" in parsed_result["format"]["triggers"]
-
-    def test_structured_outputs_params_integration(
-        self, gptoss_parser, tool_server_with_python
-    ):
-        """Test integration with StructuredOutputsParams."""
-        # Generate structural tag
-        structural_tag = gptoss_parser.prepare_structured_tag(
-            None, tool_server_with_python
-        )
-
-        # Create StructuredOutputsParams
-        params = StructuredOutputsParams(structural_tag=structural_tag)
-
-        # Verify the tag is properly stored and accessible
-        assert params.structural_tag == structural_tag
-
-        # Verify the tag is valid JSON
-        parsed_tag = json.loads(params.structural_tag)
-        assert parsed_tag["type"] == "structural_tag"
-
-    @pytest.mark.parametrize(
-        "browser, python, container, expected_tags",
-        [
-            # No tools
-            (False, False, False, 1),
-            # Single tool
-            (True, False, False, 3),
-            # Multiple tools
-            (True, True, False, 5),
-            # All tools
-            (True, True, True, 7),
-        ],
-    )
-    def test_tool_server_interaction_flow(
-        self, gptoss_parser, browser, python, container, expected_tags
-    ):
-        """Test the complete tool server interaction flow."""
-
-        # Create a mock ToolServer
-        tool_server = Mock(spec=ToolServer)
-
-        # Simulate tool availability based on parameters
-        tool_server.has_tool = Mock(
-            side_effect=lambda tool: {
-                "browser": browser,
-                "python": python,
-                "container": container,
-            }.get(tool, False)
-        )
-
-        # Run the parser and verify results
-        result = gptoss_parser.prepare_structured_tag(None, tool_server)
-        parsed_result = json.loads(result)
-
-        # Validate number of tags
-        assert len(parsed_result["format"]["tags"]) == expected_tags
-
-        # Verify tool-specific tags exist for enabled tools
-        tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]]
-        for tool, enabled in {
-            "browser": browser,
-            "python": python,
-            "container": container,
-        }.items():
-            if enabled:
-                assert f"<|channel|>commentary to={tool}" in tag_begins
-                assert f"<|channel|>analysis to={tool}" in tag_begins
-
-    def test_original_tag_preservation(self, gptoss_parser, tool_server_with_python):
-        """Test that original tags are preserved when provided."""
-        original_tag = '{"type": "custom_tag", "data": "preserved"}'
-
-        result = gptoss_parser.prepare_structured_tag(
-            original_tag, tool_server_with_python
-        )
-
-        # Should return original tag unchanged
-        assert result == original_tag
-
-    @pytest.mark.parametrize(
-        "tools",
-        [
-            [],
-            ["browser"],
-            ["python"],
-            ["container"],
-            ["browser", "python"],
-            ["browser", "container"],
-            ["python", "container"],
-            ["browser", "python", "container"],
-        ],
-    )
-    def test_json_validity_comprehensive(self, gptoss_parser, tools):
-        """Test JSON validity across all possible tool combinations."""
-
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools)
-
-        result = gptoss_parser.prepare_structured_tag(None, tool_server)
-
-        # Should be valid JSON
-        parsed_result = json.loads(result)
-
-        # Should have correct structure
-        assert parsed_result["type"] == "structural_tag"
-        assert "format" in parsed_result
-        assert "tags" in parsed_result["format"]
-        assert "triggers" in parsed_result["format"]
-
-        # Tag count should be: 1 (analysis) + 2 * len(tools)
-        expected_tag_count = 1 + (2 * len(tools))
-        assert len(parsed_result["format"]["tags"]) == expected_tag_count
-
-    def test_error_handling_invalid_tool_server(self, gptoss_parser):
-        """Test error handling with invalid tool server."""
-        # Tool server that raises exceptions
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(side_effect=Exception("Tool server error"))
-
-        # Should handle gracefully and still return a valid tag
-        with pytest.raises(Exception, match="Tool server error"):
-            gptoss_parser.prepare_structured_tag(None, tool_server)
-
-    def test_concurrent_requests_isolation(self, gptoss_parser):
-        """Test that concurrent requests don't interfere with each other."""
-        # Simulate concurrent requests with different tool servers
-        tool_server_1 = Mock(spec=ToolServer)
-        tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python")
-
-        tool_server_2 = Mock(spec=ToolServer)
-        tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser")
-
-        # Generate tags concurrently
-        result_1 = gptoss_parser.prepare_structured_tag(None, tool_server_1)
-        result_2 = gptoss_parser.prepare_structured_tag(None, tool_server_2)
-
-        # Parse results
-        parsed_1 = json.loads(result_1)
-        parsed_2 = json.loads(result_2)
-
-        # Verify they have different tool configurations
-        tags_1 = [tag["begin"] for tag in parsed_1["format"]["tags"]]
-        tags_2 = [tag["begin"] for tag in parsed_2["format"]["tags"]]
-
-        # Result 1 should have python tags
-        assert "<|channel|>commentary to=python" in tags_1
-        assert "<|channel|>commentary to=browser" not in tags_1
-
-        # Result 2 should have browser tags
-        assert "<|channel|>commentary to=browser" in tags_2
-        assert "<|channel|>commentary to=python" not in tags_2
-
-    def test_tag_format_consistency(self, gptoss_parser):
-        """Test that all generated tags follow consistent format."""
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(
-            side_effect=lambda tool: tool in ["python", "browser"]
-        )
-
-        result = gptoss_parser.prepare_structured_tag(None, tool_server)
-        parsed_result = json.loads(result)
-
-        # Verify all tags have required fields
-        for tag in parsed_result["format"]["tags"]:
-            assert "begin" in tag
-            assert "content" in tag
-            assert "end" in tag
-            assert tag["content"]["type"] == "any_text"
-            assert tag["end"] == "<|end|>"
-
-            # Verify begin format
-            assert tag["begin"].startswith("<|channel|>")
-
-    def test_trigger_configuration(self, gptoss_parser):
-        """Test trigger configuration for different tool setups."""
-        # Test with no tools
-        result_no_tools = gptoss_parser.prepare_structured_tag(None, None)
-        parsed_no_tools = json.loads(result_no_tools)
-        assert parsed_no_tools["format"]["triggers"] == ["<|channel|>analysis"]
-
-        # Test with tools
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python")
-
-        result_with_tools = gptoss_parser.prepare_structured_tag(None, tool_server)
-        parsed_with_tools = json.loads(result_with_tools)
-
-        expected_triggers = ["<|channel|>analysis", "<|channel|>commentary to="]
-        assert set(parsed_with_tools["format"]["triggers"]) == set(expected_triggers)

tests/entrypoints/openai/tool_parsers/test_granite4_tool_parser.py

@@ -0,0 +1,360 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+import random
+from typing import Any
+
+import openai
+import pytest
+from transformers import AutoTokenizer
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.engine.protocol import (
+    DeltaMessage,
+)
+from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
+
+from ....utils import RemoteOpenAIServer
+
+MODEL = "ibm-granite/granite-4.0-h-tiny"
+
+
+@pytest.fixture(scope="module")
+def server():
+    model = MODEL
+    args_for_model = [
+        "--enforce-eager",
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "granite4",
+        "--tokenizer",
+        "ibm-granite/granite-4.0-h-tiny",
+        "--max-model-len",
+        "4096",
+        "--max-num-seqs",
+        "2",
+    ]
+    with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server:
+        yield server
+
+
+def create_complex_input(create_string_args: bool):
+    coord_arg: dict | str = {
+        "coordinates": [[23.54, 43.1], [-12.2, 54.3], [4, 5]],
+        "coordinate_type": "latlong",
+    }
+    if create_string_args:
+        # test granite behavior
+        coord_arg = json.dumps(coord_arg)
+    return [
+        {"name": "find_bbox", "arguments": coord_arg},
+        {
+            "name": "get_stock_price",
+            "arguments": {
+                "symbol": "AAPL",
+                "start_date": "2021-01-01",
+                "end_date": "2021-12-31",
+            },
+        },
+        {"name": "find_bbox", "arguments": coord_arg},
+    ]
+
+
+def random_chunks(s: str, min_len: int, max_len: int):
+    chunks = []
+    i = 0
+    n = len(s)
+
+    while i < n:
+        size = random.randint(min_len, max_len)
+        chunks.append(s[i : i + size])
+        i += size
+
+    return chunks
+
+
+@pytest.fixture(scope="module")
+def tokenizer():
+    return AutoTokenizer.from_pretrained(MODEL)
+
+
+# create a variety of input chunk sizes
+@pytest.mark.parametrize(
+    "min_chunk, max_chunk",
+    [
+        (1, 1),
+        (1, 2),
+        (5, 7),
+        (6, 20),
+    ],
+)
+def test_tool_call_parser_complex(min_chunk: int, max_chunk: int, tokenizer):
+    input_dicts = create_complex_input(True)
+
+    formatted_tcs = [
+        "<tool_call> " + json.dumps(call) + " </tool_call>" for call in input_dicts
+    ]
+
+    text_messages = [
+        "Here goes the bbox call: \n",
+        " Now the stock price call: \n ",
+        " Now another bbox call: \n ",
+        " See? I'm a helpful assistant.",
+    ]
+
+    test_input = (
+        text_messages[0]
+        + formatted_tcs[0]
+        + text_messages[1]
+        + formatted_tcs[1]
+        + text_messages[2]
+        + formatted_tcs[2]
+        + text_messages[3]
+    )
+
+    any_chat_request = ChatCompletionRequest(
+        seed=42,
+        model=MODEL,
+        messages=[],
+    )
+
+    parser = Granite4ToolParser(tokenizer=tokenizer)
+
+    delta_messages = list[DeltaMessage]()
+    for text in random_chunks(test_input, min_chunk, max_chunk):
+        delta = parser.extract_tool_calls_streaming(
+            previous_text="",
+            current_text="",
+            delta_text=text,
+            previous_token_ids=[],
+            current_token_ids=[],
+            delta_token_ids=[],
+            request=any_chat_request,
+        )
+        if delta is not None:
+            delta_messages.append(delta)
+
+    content = ""
+    tool_calls = list[dict[str, Any]]()
+
+    current_name = "__start__"
+    current_args = ""
+
+    for msg in delta_messages:
+        if msg.content:
+            content += msg.content
+        for tool_call in msg.tool_calls:
+            if delta_func := tool_call.function:
+                if delta_func.name is not None:
+                    if current_name == "__start__":
+                        current_name = delta_func.name
+
+                    if delta_func.name != current_name:
+                        tool_calls.append(
+                            {
+                                "name": current_name,
+                                "arguments": json.loads(current_args),
+                            }
+                        )
+                        current_name = delta_func.name
+                        current_args = ""
+
+                if delta_func.arguments:
+                    current_args += delta_func.arguments
+
+    if current_name != "__start__":
+        tool_calls.append({"name": current_name, "arguments": json.loads(current_args)})
+
+    assert content == "".join(text_messages)
+    assert tool_calls == create_complex_input(False)
+
+
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_acme_region_name_for_transaction_id",
+            "description": "Returns ACME transaction/transaction ID information"
+            " including ACME regions\n\nArgs:\n    start_time "
+            "(str): Start date and time in datetime format "
+            '"%Y-%m-%dT%H:%M:%S.%f"\n    end_time (str): End '
+            "date and time in datetime format "
+            '"%Y-%m-%dT%H:%M:%S.%f"\n    size (int, optional): '
+            "Number of ACME Transaction IDs to return\n    "
+            "order (str, optional): Sort by most run "
+            "transaction IDs. The value can be 'asc' for "
+            "ascending or 'desc' for descending\n    "
+            "transaction_id (str, optional): ACME Transaction "
+            "ID to filter on\n    acme_region (str, optional): "
+            "ACME Region to filter on\nReturns:\n    - A "
+            "dictionary containing a list of ACME transaction "
+            "ids and the ACME regions they run in:\n        {\n"
+            '            "Number of transaction IDs"   : int,\n'
+            '            "Total transaction IDs available": int'
+            ',\n            "ACME Transaction IDs": [\n        '
+            '        {\n                    "Transaction ID": '
+            'str,\n                    "Number of runs": int,\n'
+            '                    "ACME Regions": [str],\n      '
+            "          },\n                ...\n            ],"
+            '\n            "Start time"         : datetime,\n '
+            '           "End time"           : datetime,\n    '
+            '        "Order"              : str\n        }\n  '
+            "  - If no ACME region found for transaction id, "
+            'returns:\n        {"Success": "No ACME region '
+            'found for transaction id."}\n    - If an error '
+            'occurs, returns:\n        {"Error": "{exception'
+            ' message}"}',
+            "parameters": {
+                "properties": {
+                    "start_time": {},
+                    "end_time": {},
+                    "size": {"default": 500},
+                    "order": {"default": "desc"},
+                    "transaction_id": {"default": None},
+                    "acme_region": {"default": None},
+                },
+                "required": ["start_time", "end_time"],
+                "type": "object",
+            },
+        },
+    }
+]
+
+tools2 = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_current_weather",
+            "description": "Get the current weather",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "description": "The city and state, e.g. San Francisco, CA",
+                        "type": "string",
+                    }
+                },
+                "required": ["location"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "get_stock_price",
+            "description": "Retrieves the current stock price for a given "
+            "ticker symbol. The ticker symbol must be a valid "
+            "symbol for a publicly traded company on a major US"
+            " stock exchange like NYSE or NASDAQ. The tool will"
+            " return the latest trade price in USD. It should "
+            "be used when the user asks about the current or "
+            "most recent price of a specific stock. It will not"
+            " provide any other information about the stock or"
+            " company.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "ticker": {
+                        "description": "The stock ticker symbol, e.g."
+                        " AAPL for Apple Inc.",
+                        "type": "string",
+                    }
+                },
+            },
+        },
+    },
+]
+
+messages = [
+    {
+        "content": "\n\nSystem: You are a helpful, precise, and methodical AI"
+        " assistant that uses tool outputs provided inline.\nAlways"
+        " assume the current datetime is 2026-01-29T13:59:09.238901"
+        "+00:00.\n\nIf you receive a ToolMessage with `tool_call_id"
+        '` equal to "get_time_range" (or "time_range_tool"), you '
+        "MUST:\n  1. Parse that JSON and use the values `start` and"
+        " `end` directly when calling other tools.\n  2. Do not "
+        "re-call or re-compute the time range.\n  3. Pass resolved "
+        "values (ISO strings) as arguments to any subsequent tool "
+        "(do not pass function metadata or placeholders).\n  4. If "
+        "a tool requires datetime objects rather than strings, "
+        "convert the ISO strings into language-native datetime "
+        "objects before invoking.\n\nAlways return fully resolved "
+        "arguments in correct types (e.g., ISO datetime strings or"
+        " datetime objects) and never include placeholders like "
+        '"<start>".\n\n',
+        "role": "system",
+    },
+    {
+        "content": "What are the transaction IDs that ran in the"
+        " ACME region A9345 over the last two months?",
+        "role": "user",
+    },
+    {
+        "content": '["2026-01-26T09: 51: 55.467722Z", "2026-01-27T09: 51: 55.467722Z"]',
+        "role": "tool",
+        "tool_call_id": "time_range_tool",
+    },
+]
+messages2 = [{"role": "user", "content": "What's stock price for IBM?"}]
+
+messages3 = [{"role": "user", "content": "What's the current weather in New York?"}]
+
+
+def get_args(client: openai.OpenAI, _tools, _messages, _stop):
+    response = client.chat.completions.create(
+        model=MODEL,
+        messages=_messages,
+        temperature=0,
+        tools=_tools,
+        max_tokens=200,
+        stop=_stop,
+        tool_choice="auto",
+    )
+
+    return response.choices[0].message.tool_calls[0].function.arguments
+
+
+async def get_args_streaming(
+    async_client: openai.AsyncOpenAI, _tools, _messages, _stop
+):
+    stream = await async_client.chat.completions.create(
+        model=MODEL,
+        messages=_messages,
+        temperature=0,
+        tools=_tools,
+        max_tokens=200,
+        stop=_stop,
+        tool_choice="auto",
+        stream=True,
+    )
+    full_call = []
+    async for chunk in stream:
+        tc = chunk.choices[0].delta.tool_calls
+        if tc and tc[0].function.arguments:
+            full_call.append(tc[0].function.arguments)
+    return "".join(full_call)
+
+
+async def run_scenario(server: RemoteOpenAIServer, _tools, _messages, _stop):
+    non_streaming = get_args(server.get_client(), _tools, _messages, _stop)
+    json.loads(non_streaming)  # verify that it is json loadable
+    streaming = await get_args_streaming(
+        server.get_async_client(), _tools, _messages, _stop
+    )
+    json.loads(streaming)
+    assert non_streaming == streaming, f"{non_streaming=}, {streaming=}"
+
+
+@pytest.mark.asyncio
+async def test_stop_sequence_interference(server: RemoteOpenAIServer):
+    print("Testing scenario 1")
+    await run_scenario(server, tools, messages, "veroniqueprattyushveroniqueprattyush")
+
+    print("Testing scenario 2")
+    await run_scenario(
+        server, tools2, messages2, "veroniqueprattyushveroniqueprattyush"
+    )
+
+    print("Testing scenario 3")
+    await run_scenario(server, tools2, messages3, "prattyush")

tests/entrypoints/openai/tool_parsers/test_hermes_tool_parser.py

@@ -3,29 +3,22 @@
 
 import json
 
+import openai
 import pytest
+import pytest_asyncio
+from huggingface_hub import snapshot_download
+from typing_extensions import TypedDict
 
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
 from vllm.tokenizers import TokenizerLike
+from vllm.tool_parsers.abstract_tool_parser import ToolParser
+from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
 from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
 
 from ....utils import RemoteOpenAIServer
 
-MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci"
 
-SERVER_ARGS = [
-    "--enforce-eager",
-    "--enable-auto-tool-choice",
-    "--tool-call-parser",
-    "hermes",
-    "--enable-lora",
-    "--lora-modules",
-    f"{LORA_MODEL}={LORA_MODEL}",
-    "--tokenizer",
-    f"{LORA_MODEL}",
-]
-
 TOOLS = [
     {
         "type": "function",
@@ -50,6 +43,75 @@ TOOLS = [
     }
 ]
 
+
+class ServerConfig(TypedDict, total=False):
+    model: str
+    arguments: list[str]
+    model_arg: str
+    tool_parser: ToolParser
+
+
+CONFIGS: dict[str, ServerConfig] = {
+    "llama": {
+        "model": "meta-llama/Llama-3.2-1B-Instruct",
+        "arguments": [
+            "--enforce-eager",
+            "--enable-auto-tool-choice",
+            "--tool-call-parser",
+            "hermes",
+            "--enable-lora",
+            "--lora-modules",
+            f"{LORA_MODEL}={LORA_MODEL}",
+            "--tokenizer",
+            f"{LORA_MODEL}",
+        ],
+        "model_arg": LORA_MODEL,
+        "tool_parser": Hermes2ProToolParser,
+    },
+    "granite4": {
+        "model": "ibm-granite/granite-4.0-h-tiny",
+        "arguments": [
+            "--enforce-eager",
+            "--enable-auto-tool-choice",
+            "--tool-call-parser",
+            "granite4",
+            "--tokenizer",
+            "ibm-granite/granite-4.0-h-tiny",
+            "--max-model-len",
+            "4096",
+            "--max-num-seqs",
+            "2",
+        ],
+        "model_arg": "ibm-granite/granite-4.0-h-tiny",
+        "tool_parser": Granite4ToolParser,
+    },
+}
+
+
+# for each server config, download the model and return the config
+@pytest.fixture(scope="session", params=CONFIGS.keys())
+def server_config(request):
+    config = CONFIGS[request.param]
+
+    # download model and tokenizer using transformers
+    snapshot_download(config["model"])
+    yield CONFIGS[request.param]
+
+
+@pytest.fixture(scope="module")
+def server(request, server_config: ServerConfig):
+    model = server_config["model"]
+    args_for_model = server_config["arguments"]
+    with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server:
+        yield server
+
+
+@pytest_asyncio.fixture
+async def client(server: RemoteOpenAIServer):
+    async with server.get_async_client() as async_client:
+        yield async_client
+
+
 PRODUCT_TOOLS = [
     {
         "type": "function",
@@ -87,186 +149,182 @@ PRODUCT_MESSAGES = [
 
 
 @pytest.mark.asyncio
-async def test_non_streaming_tool_call():
+async def test_non_streaming_tool_call(
+    client: openai.AsyncOpenAI, server_config: ServerConfig
+):
     """Test tool call in non-streaming mode."""
-    with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
-        client = server.get_async_client()
 
-        response = await client.chat.completions.create(
-            model=LORA_MODEL,
-            messages=MESSAGES,
-            tools=TOOLS,
-            tool_choice="auto",
-            temperature=0.0,
-        )
+    response = await client.chat.completions.create(
+        model=server_config["model_arg"],
+        messages=MESSAGES,
+        tools=TOOLS,
+        tool_choice="auto",
+        temperature=0.0,
+    )
 
-        assert response.choices
-        choice = response.choices[0]
-        message = choice.message
+    assert response.choices
+    choice = response.choices[0]
+    message = choice.message
 
-        assert choice.finish_reason == "tool_calls"
-        assert message.tool_calls is not None
+    assert choice.finish_reason == "tool_calls"
+    assert message.tool_calls is not None
 
-        tool_call = message.tool_calls[0]
-        assert tool_call.type == "function"
-        assert tool_call.function.name == "get_current_weather"
+    tool_call = message.tool_calls[0]
+    assert tool_call.type == "function"
+    assert tool_call.function.name == "get_current_weather"
 
-        arguments = json.loads(tool_call.function.arguments)
-        assert "location" in arguments
-        assert "Boston" in arguments["location"]
-        print("\n[Non-Streaming Test Passed]")
-        print(f"Tool Call: {tool_call.function.name}")
-        print(f"Arguments: {arguments}")
+    arguments = json.loads(tool_call.function.arguments)
+    assert "location" in arguments
+    assert "Boston" in arguments["location"]
+    print("\n[Non-Streaming Test Passed]")
+    print(f"Tool Call: {tool_call.function.name}")
+    print(f"Arguments: {arguments}")
 
 
 @pytest.mark.asyncio
-async def test_streaming_tool_call():
+async def test_streaming_tool_call(
+    client: openai.AsyncOpenAI, server_config: ServerConfig
+):
     """Test tool call in streaming mode."""
-    with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
-        client = server.get_async_client()
 
-        stream = await client.chat.completions.create(
-            model=LORA_MODEL,
-            messages=MESSAGES,
-            tools=TOOLS,
-            tool_choice="auto",
-            temperature=0.0,
-            stream=True,
-        )
+    stream = await client.chat.completions.create(
+        model=server_config["model_arg"],
+        messages=MESSAGES,
+        tools=TOOLS,
+        tool_choice="auto",
+        temperature=0.0,
+        stream=True,
+    )
 
-        tool_call_chunks = {}
-        async for chunk in stream:
-            if not chunk.choices:
-                continue
+    tool_call_chunks = {}
+    async for chunk in stream:
+        if not chunk.choices:
+            continue
 
-            delta = chunk.choices[0].delta
-            if not delta or not delta.tool_calls:
-                continue
+        delta = chunk.choices[0].delta
+        if not delta or not delta.tool_calls:
+            continue
 
-            for tool_chunk in delta.tool_calls:
-                index = tool_chunk.index
-                if index not in tool_call_chunks:
-                    tool_call_chunks[index] = {"name": "", "arguments": ""}
+        for tool_chunk in delta.tool_calls:
+            index = tool_chunk.index
+            if index not in tool_call_chunks:
+                tool_call_chunks[index] = {"name": "", "arguments": ""}
 
-                if tool_chunk.function.name:
-                    tool_call_chunks[index]["name"] += tool_chunk.function.name
-                if tool_chunk.function.arguments:
-                    tool_call_chunks[index]["arguments"] += (
-                        tool_chunk.function.arguments
-                    )
+            if tool_chunk.function.name:
+                tool_call_chunks[index]["name"] += tool_chunk.function.name
+            if tool_chunk.function.arguments:
+                tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments
 
-        assert len(tool_call_chunks) == 1
-        reconstructed_tool_call = tool_call_chunks[0]
+    assert len(tool_call_chunks) == 1
+    reconstructed_tool_call = tool_call_chunks[0]
 
-        assert reconstructed_tool_call["name"] == "get_current_weather"
+    assert reconstructed_tool_call["name"] == "get_current_weather"
 
-        arguments = json.loads(reconstructed_tool_call["arguments"])
-        assert "location" in arguments
-        assert "Boston" in arguments["location"]
-        print("\n[Streaming Test Passed]")
-        print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
-        print(f"Reconstructed Arguments: {arguments}")
+    arguments = json.loads(reconstructed_tool_call["arguments"])
+    assert "location" in arguments
+    assert "Boston" in arguments["location"]
+    print("\n[Streaming Test Passed]")
+    print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
+    print(f"Reconstructed Arguments: {arguments}")
 
 
 @pytest.mark.asyncio
-async def test_non_streaming_product_tool_call():
+async def test_non_streaming_product_tool_call(
+    client: openai.AsyncOpenAI, server_config: ServerConfig
+):
     """Test tool call integer and boolean parameters in non-streaming mode."""
-    with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
-        client = server.get_async_client()
 
-        response = await client.chat.completions.create(
-            model=LORA_MODEL,
-            messages=PRODUCT_MESSAGES,
-            tools=PRODUCT_TOOLS,
-            tool_choice="auto",
-            temperature=0.66,
-        )
+    response = await client.chat.completions.create(
+        model=server_config["model_arg"],
+        messages=PRODUCT_MESSAGES,
+        tools=PRODUCT_TOOLS,
+        tool_choice="auto",
+        temperature=0.66,
+    )
 
-        assert response.choices
-        choice = response.choices[0]
-        message = choice.message
+    assert response.choices
+    choice = response.choices[0]
+    message = choice.message
 
-        assert choice.finish_reason == "tool_calls"
-        assert message.tool_calls is not None
+    assert choice.finish_reason == "tool_calls"
+    assert message.tool_calls is not None
 
-        tool_call = message.tool_calls[0]
-        assert tool_call.type == "function"
-        assert tool_call.function.name == "get_product_info"
+    tool_call = message.tool_calls[0]
+    assert tool_call.type == "function"
+    assert tool_call.function.name == "get_product_info"
 
-        arguments = json.loads(tool_call.function.arguments)
-        assert "product_id" in arguments
-        assert "inserted" in arguments
+    arguments = json.loads(tool_call.function.arguments)
+    assert "product_id" in arguments
+    assert "inserted" in arguments
 
-        product_id = arguments.get("product_id")
-        inserted = arguments.get("inserted")
+    product_id = arguments.get("product_id")
+    inserted = arguments.get("inserted")
 
-        assert isinstance(product_id, int)
-        assert product_id == 7355608
-        assert isinstance(inserted, bool)
-        assert inserted is True
+    assert isinstance(product_id, int)
+    assert product_id == 7355608
+    assert isinstance(inserted, bool)
+    assert inserted is True
 
-        print("\n[Non-Streaming Product Test Passed]")
-        print(f"Tool Call: {tool_call.function.name}")
-        print(f"Arguments: {arguments}")
+    print("\n[Non-Streaming Product Test Passed]")
+    print(f"Tool Call: {tool_call.function.name}")
+    print(f"Arguments: {arguments}")
 
 
 @pytest.mark.asyncio
-async def test_streaming_product_tool_call():
+async def test_streaming_product_tool_call(
+    client: openai.AsyncOpenAI, server_config: ServerConfig
+):
     """Test tool call integer and boolean parameters in streaming mode."""
-    with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
-        client = server.get_async_client()
 
-        stream = await client.chat.completions.create(
-            model=LORA_MODEL,
-            messages=PRODUCT_MESSAGES,
-            tools=PRODUCT_TOOLS,
-            tool_choice="auto",
-            temperature=0.66,
-            stream=True,
-        )
+    stream = await client.chat.completions.create(
+        model=server_config["model_arg"],
+        messages=PRODUCT_MESSAGES,
+        tools=PRODUCT_TOOLS,
+        tool_choice="auto",
+        temperature=0.66,
+        stream=True,
+    )
 
-        tool_call_chunks = {}
-        async for chunk in stream:
-            if not chunk.choices:
-                continue
+    tool_call_chunks = {}
+    async for chunk in stream:
+        if not chunk.choices:
+            continue
 
-            delta = chunk.choices[0].delta
-            if not delta or not delta.tool_calls:
-                continue
+        delta = chunk.choices[0].delta
+        if not delta or not delta.tool_calls:
+            continue
 
-            for tool_chunk in delta.tool_calls:
-                index = tool_chunk.index
-                if index not in tool_call_chunks:
-                    tool_call_chunks[index] = {"name": "", "arguments": ""}
+        for tool_chunk in delta.tool_calls:
+            index = tool_chunk.index
+            if index not in tool_call_chunks:
+                tool_call_chunks[index] = {"name": "", "arguments": ""}
 
-                if tool_chunk.function.name:
-                    tool_call_chunks[index]["name"] += tool_chunk.function.name
-                if tool_chunk.function.arguments:
-                    tool_call_chunks[index]["arguments"] += (
-                        tool_chunk.function.arguments
-                    )
+            if tool_chunk.function.name:
+                tool_call_chunks[index]["name"] += tool_chunk.function.name
+            if tool_chunk.function.arguments:
+                tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments
 
-        assert len(tool_call_chunks) == 1
-        reconstructed_tool_call = tool_call_chunks[0]
+    assert len(tool_call_chunks) == 1
+    reconstructed_tool_call = tool_call_chunks[0]
 
-        assert reconstructed_tool_call["name"] == "get_product_info"
+    assert reconstructed_tool_call["name"] == "get_product_info"
 
-        arguments = json.loads(reconstructed_tool_call["arguments"])
-        assert "product_id" in arguments
-        assert "inserted" in arguments
+    arguments = json.loads(reconstructed_tool_call["arguments"])
+    assert "product_id" in arguments
+    assert "inserted" in arguments
 
-        # Handle type coercion for streaming test as well
-        product_id = arguments.get("product_id")
-        inserted = arguments.get("inserted")
+    # Handle type coercion for streaming test as well
+    product_id = arguments.get("product_id")
+    inserted = arguments.get("inserted")
 
-        assert isinstance(product_id, int)
-        assert product_id == 7355608
-        assert isinstance(inserted, bool)
-        assert inserted is True
+    assert isinstance(product_id, int)
+    assert product_id == 7355608
+    assert isinstance(inserted, bool)
+    assert inserted is True
 
-        print("\n[Streaming Product Test Passed]")
-        print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
-        print(f"Reconstructed Arguments: {arguments}")
+    print("\n[Streaming Product Test Passed]")
+    print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
+    print(f"Reconstructed Arguments: {arguments}")
 
 
 @pytest.fixture
@@ -276,9 +334,10 @@ def qwen_tokenizer() -> TokenizerLike:
     return get_tokenizer("Qwen/Qwen3-32B")
 
 
-@pytest.fixture
-def hermes_parser(qwen_tokenizer: TokenizerLike) -> Hermes2ProToolParser:
-    return Hermes2ProToolParser(qwen_tokenizer)
+@pytest.fixture(params=CONFIGS.keys())
+def hermes_parser(request, qwen_tokenizer: TokenizerLike) -> ToolParser:
+    config = CONFIGS[request.param]
+    return config["tool_parser"](qwen_tokenizer)
 
 
 @pytest.fixture
@@ -292,7 +351,7 @@ def any_chat_request() -> ChatCompletionRequest:
 
 def test_hermes_parser_streaming_just_forward_text(
     qwen_tokenizer: TokenizerLike,
-    hermes_parser: Hermes2ProToolParser,
+    hermes_parser: ToolParser,
     any_chat_request: ChatCompletionRequest,
 ) -> None:
     text = """This is some prior text that has nothing to do with tool calling."""
@@ -324,7 +383,7 @@ def test_hermes_parser_streaming_just_forward_text(
 
 def test_hermes_parser_streaming_failure_case_bug_19056(
     qwen_tokenizer: TokenizerLike,
-    hermes_parser: Hermes2ProToolParser,
+    hermes_parser: ToolParser,
     any_chat_request: ChatCompletionRequest,
 ) -> None:
     text = """<tool_call>
@@ -358,7 +417,7 @@ def test_hermes_parser_streaming_failure_case_bug_19056(
 
 def test_hermes_parser_streaming(
     qwen_tokenizer: TokenizerLike,
-    hermes_parser: Hermes2ProToolParser,
+    hermes_parser: ToolParser,
     any_chat_request: ChatCompletionRequest,
 ) -> None:
     text = '<tool_call>\
@@ -387,16 +446,20 @@ def test_hermes_parser_streaming(
             delta_messages.append(delta)
     print(delta_messages)
     assert delta_messages[0].tool_calls[0].function.name == "get_current_temperature"
-    tool_call_args = "".join(
-        delta.tool_calls[0].function.arguments or "" for delta in delta_messages
-    )
-    assert tool_call_args == (
-        '{"location":"San Francisco, California, United States", "unit": "celsius"}'
+    # load to normalize whitespace
+    tool_call_args = json.loads(
+        "".join(
+            delta.tool_calls[0].function.arguments or "" for delta in delta_messages
+        )
     )
+    assert tool_call_args == {
+        "location": "San Francisco, California, United States",
+        "unit": "celsius",
+    }
 
 
 def test_hermes_parser_non_streaming_no_tool_call(
-    hermes_parser: Hermes2ProToolParser,
+    hermes_parser: ToolParser,
     any_chat_request: ChatCompletionRequest,
 ) -> None:
     text = """This is not a tool call."""
@@ -410,7 +473,7 @@ def test_hermes_parser_non_streaming_no_tool_call(
 
 
 def test_hermes_parser_non_streaming_tool_call_between_tags(
-    hermes_parser: Hermes2ProToolParser,
+    hermes_parser: ToolParser,
     any_chat_request: ChatCompletionRequest,
 ) -> None:
     text = """<tool_call>
@@ -428,9 +491,12 @@ def test_hermes_parser_non_streaming_tool_call_between_tags(
 
 
 def test_hermes_parser_non_streaming_tool_call_until_eos(
-    hermes_parser: Hermes2ProToolParser,
+    hermes_parser: ToolParser,
     any_chat_request: ChatCompletionRequest,
 ) -> None:
+    if isinstance(hermes_parser, Granite4ToolParser):
+        pytest.skip(reason="The Granite4 tool parser enforces a complete response")
+
     text = """<tool_call>
 {"name": "final_answer", "arguments": {"trigger": true}}"""
     tool_call = hermes_parser.extract_tool_calls(
@@ -445,7 +511,7 @@ def test_hermes_parser_non_streaming_tool_call_until_eos(
 
 
 def test_hermes_parser_non_streaming_tool_call_invalid_json(
-    hermes_parser: Hermes2ProToolParser,
+    hermes_parser: ToolParser,
     any_chat_request: ChatCompletionRequest,
 ) -> None:
     # Missing closing brace to trigger exception

tests/kernels/moe/modular_kernel_tools/mk_objects.py

@@ -33,7 +33,10 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
 )
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import is_deep_gemm_supported
-from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
+from vllm.utils.flashinfer import (
+    has_flashinfer_cutlass_fused_moe,
+    has_flashinfer_nvlink_one_sided,
+)
 from vllm.utils.import_utils import (
     has_aiter,
     has_deep_ep,
@@ -234,15 +237,15 @@ if has_mori():
     )
 
 if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
-    from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import (  # noqa: E501
-        FlashInferA2APrepareAndFinalize,
-    )
     from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
         FlashInferExperts,
     )
+    from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize import (  # noqa: E501
+        FlashInferNVLinkTwoSidedPrepareAndFinalize,
+    )
 
     register_prepare_and_finalize(
-        FlashInferA2APrepareAndFinalize,
+        FlashInferNVLinkTwoSidedPrepareAndFinalize,
         standard_format,
         nvfp4_types + fp8_types,
         blocked_quantization_support=True,
@@ -263,6 +266,36 @@ else:
     FlashInferCutlassMoEPrepareAndFinalize = None
     FlashInferExperts = None
 
+if (
+    has_flashinfer_nvlink_one_sided()
+    and has_flashinfer_cutlass_fused_moe()
+    and current_platform.has_device_capability(100)
+):
+    from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize import (  # noqa: E501
+        FlashInferNVLinkOneSidedPrepareAndFinalize,
+    )
+
+    register_prepare_and_finalize(
+        FlashInferNVLinkOneSidedPrepareAndFinalize,
+        standard_format,
+        nvfp4_types,
+        blocked_quantization_support=False,
+        backend="flashinfer_nvlink_one_sided",
+        supports_apply_weight_on_input=False,
+    )
+
+if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
+    from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
+        TrtLlmNvFp4ExpertsModular,
+    )
+
+    register_experts(
+        TrtLlmNvFp4ExpertsModular,
+        standard_format,
+        nvfp4_types,
+        blocked_quantization_support=False,
+        supports_expert_map=True,
+    )
 
 if has_aiter():
     from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

tests/model_executor/model_loader/runai_streamer_loader/test_runai_utils.py

@@ -19,6 +19,7 @@ from vllm.transformers_utils.runai_utils import (
 def test_is_runai_obj_uri():
     assert is_runai_obj_uri("gs://some-gcs-bucket/path")
     assert is_runai_obj_uri("s3://some-s3-bucket/path")
+    assert is_runai_obj_uri("az://some-azure-container/path")
     assert not is_runai_obj_uri("nfs://some-nfs-path")
 
 

tests/model_executor/model_loader/test_ep_weight_filter.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for EP weight filtering during model loading."""
+
+import glob
+import tempfile
+
+import huggingface_hub.constants
+import pytest
+import torch
+
+from vllm.model_executor.model_loader.ep_weight_filter import (
+    compute_local_expert_ids,
+    parse_expert_id,
+    should_skip_weight,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    safetensors_weights_iterator,
+)
+
+# ---------------------------------------------------------------------------
+# Unit tests for parse_expert_id
+# ---------------------------------------------------------------------------
+
+
+class TestParseExpertId:
+    def test_routed_expert(self):
+        name = "model.layers.0.mlp.experts.42.gate_proj.weight"
+        assert parse_expert_id(name) == 42
+
+    def test_large_expert_id(self):
+        name = "model.layers.60.mlp.experts.383.down_proj.weight"
+        assert parse_expert_id(name) == 383
+
+    def test_shared_expert(self):
+        # Shared experts use a different naming convention in most models
+        name = "model.layers.0.mlp.shared_experts.gate_proj.weight"
+        assert parse_expert_id(name) is None
+
+    def test_attention_weight(self):
+        name = "model.layers.0.self_attn.q_proj.weight"
+        assert parse_expert_id(name) is None
+
+    def test_embedding(self):
+        name = "model.embed_tokens.weight"
+        assert parse_expert_id(name) is None
+
+    def test_layernorm(self):
+        name = "model.layers.0.input_layernorm.weight"
+        assert parse_expert_id(name) is None
+
+    def test_fused_3d_expert(self):
+        # 3D fused-expert tensors (e.g. gpt-oss) have no numeric expert id.
+        # They must NOT be filtered — slicing happens later in weight_loader.
+        name = "model.layers.0.mlp.experts.gate_proj.weight"
+        assert parse_expert_id(name) is None
+
+    def test_fused_3d_expert_down_proj(self):
+        name = "model.layers.10.mlp.experts.down_proj.weight"
+        assert parse_expert_id(name) is None
+
+    def test_expert_scale(self):
+        # NVFP4 quantized models have scale tensors for experts
+        name = "model.layers.5.mlp.experts.100.gate_proj.weight_scale"
+        assert parse_expert_id(name) == 100
+
+    def test_expert_zero_id(self):
+        name = "model.layers.0.mlp.experts.0.up_proj.weight"
+        assert parse_expert_id(name) == 0
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for compute_local_expert_ids
+# ---------------------------------------------------------------------------
+
+
+class TestComputeLocalExpertIds:
+    def test_ep_disabled(self):
+        assert compute_local_expert_ids(64, ep_size=1, ep_rank=0) is None
+
+    def test_even_split(self):
+        # 64 experts, EP=8 → 8 per rank
+        ids = compute_local_expert_ids(64, ep_size=8, ep_rank=0)
+        assert ids == set(range(0, 8))
+
+        ids = compute_local_expert_ids(64, ep_size=8, ep_rank=7)
+        assert ids == set(range(56, 64))
+
+    def test_uneven_split(self):
+        # 10 experts, EP=3 → ranks get 4, 3, 3
+        ids_0 = compute_local_expert_ids(10, ep_size=3, ep_rank=0)
+        ids_1 = compute_local_expert_ids(10, ep_size=3, ep_rank=1)
+        ids_2 = compute_local_expert_ids(10, ep_size=3, ep_rank=2)
+
+        assert len(ids_0) == 4
+        assert len(ids_1) == 3
+        assert len(ids_2) == 3
+        # All experts covered, no overlap
+        assert ids_0 | ids_1 | ids_2 == set(range(10))
+        assert ids_0.isdisjoint(ids_1)
+        assert ids_1.isdisjoint(ids_2)
+
+    def test_384_experts_ep8(self):
+        # Kimi-K2.5 config: 384 experts, EP=8
+        for rank in range(8):
+            ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
+            assert len(ids) == 48
+
+        # All experts covered
+        all_ids = set()
+        for rank in range(8):
+            ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
+            all_ids |= ids
+        assert all_ids == set(range(384))
+
+    def test_384_experts_ep16(self):
+        for rank in range(16):
+            ids = compute_local_expert_ids(384, ep_size=16, ep_rank=rank)
+            assert len(ids) == 24
+
+    def test_384_experts_ep24(self):
+        # 384 / 24 = 16 exactly
+        for rank in range(24):
+            ids = compute_local_expert_ids(384, ep_size=24, ep_rank=rank)
+            assert len(ids) == 16
+
+    # round_robin placement tests
+
+    def test_round_robin_basic(self):
+        # 8 experts, EP=2: rank 0 → {0,2,4,6}, rank 1 → {1,3,5,7}
+        rr = "round_robin"
+        ids_0 = compute_local_expert_ids(8, 2, 0, placement=rr)
+        ids_1 = compute_local_expert_ids(8, 2, 1, placement=rr)
+        assert ids_0 == {0, 2, 4, 6}
+        assert ids_1 == {1, 3, 5, 7}
+
+    def test_round_robin_full_coverage(self):
+        # 384 experts, EP=8: all experts covered, no overlap
+        rr = "round_robin"
+        all_ids: set[int] = set()
+        for rank in range(8):
+            ids = compute_local_expert_ids(384, 8, rank, placement=rr)
+            assert ids is not None and len(ids) == 48
+            assert all_ids.isdisjoint(ids)
+            all_ids |= ids
+        assert all_ids == set(range(384))
+
+    def test_round_robin_uneven(self):
+        # 10 experts, EP=3: rank 0→{0,3,6,9}, rank 1→{1,4,7}, rank 2→{2,5,8}
+        rr = "round_robin"
+        ids_0 = compute_local_expert_ids(10, 3, 0, placement=rr)
+        ids_1 = compute_local_expert_ids(10, 3, 1, placement=rr)
+        ids_2 = compute_local_expert_ids(10, 3, 2, placement=rr)
+        assert ids_0 == {0, 3, 6, 9}
+        assert ids_1 == {1, 4, 7}
+        assert ids_2 == {2, 5, 8}
+        assert ids_0 | ids_1 | ids_2 == set(range(10))
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for should_skip_weight
+# ---------------------------------------------------------------------------
+
+
+class TestShouldSkipWeight:
+    def setup_method(self):
+        # Simulate EP=8, rank=0 → experts 0-47
+        self.local_ids = compute_local_expert_ids(384, ep_size=8, ep_rank=0)
+
+    def test_no_filter(self):
+        assert not should_skip_weight("anything", None)
+
+    def test_dense_not_skipped(self):
+        assert not should_skip_weight(
+            "model.layers.0.self_attn.q_proj.weight", self.local_ids
+        )
+
+    def test_local_expert_not_skipped(self):
+        assert not should_skip_weight(
+            "model.layers.0.mlp.experts.10.gate_proj.weight", self.local_ids
+        )
+
+    def test_remote_expert_skipped(self):
+        assert should_skip_weight(
+            "model.layers.0.mlp.experts.200.gate_proj.weight", self.local_ids
+        )
+
+    def test_boundary_expert(self):
+        # Expert 47 is local (last one), 48 is not
+        assert not should_skip_weight(
+            "model.layers.0.mlp.experts.47.gate_proj.weight", self.local_ids
+        )
+        assert should_skip_weight(
+            "model.layers.0.mlp.experts.48.gate_proj.weight", self.local_ids
+        )
+
+    def test_shared_expert_not_skipped(self):
+        assert not should_skip_weight(
+            "model.layers.0.mlp.shared_experts.gate_proj.weight", self.local_ids
+        )
+
+    def test_embedding_not_skipped(self):
+        assert not should_skip_weight("model.embed_tokens.weight", self.local_ids)
+
+    def test_fused_3d_expert_not_skipped(self):
+        # 3D fused-expert tensors (gpt-oss style) have no numeric id.
+        # Must not be skipped — weight_loader handles slicing later.
+        assert not should_skip_weight(
+            "model.layers.0.mlp.experts.gate_proj.weight", self.local_ids
+        )
+
+
+# ---------------------------------------------------------------------------
+# Integration test: safetensors_weights_iterator with EP filtering
+# ---------------------------------------------------------------------------
+
+
+class TestSafetensorsWeightsIteratorWithEpFilter:
+    """Verify that EP filtering produces a strict subset of unfiltered loading
+    and that all expected dense + local expert weights are present."""
+
+    @pytest.fixture(scope="class")
+    def gpt2_files(self):
+        """Download GPT-2 safetensors to a temp dir (shared across class)."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            huggingface_hub.constants.HF_HUB_OFFLINE = False
+            from vllm.model_executor.model_loader.weight_utils import (
+                download_weights_from_hf,
+            )
+
+            download_weights_from_hf(
+                "openai-community/gpt2",
+                allow_patterns=["*.safetensors"],
+                cache_dir=tmpdir,
+            )
+            files = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
+            assert len(files) > 0
+            yield files
+
+    def test_no_filter_returns_all(self, gpt2_files):
+        """With local_expert_ids=None, all weights are returned (no MoE)."""
+        all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
+        filtered_weights = dict(
+            safetensors_weights_iterator(gpt2_files, False, local_expert_ids=None)
+        )
+        assert set(all_weights.keys()) == set(filtered_weights.keys())
+
+    def test_empty_filter_skips_experts_only(self, gpt2_files):
+        """GPT-2 has no expert weights, so even an empty local_expert_ids
+        set should return all weights (all are dense)."""
+        all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
+        filtered_weights = dict(
+            safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set())
+        )
+        # GPT-2 has no experts, so nothing should be filtered
+        assert set(all_weights.keys()) == set(filtered_weights.keys())
+
+
+class TestEpFilterOnSyntheticMoeWeights:
+    """Create synthetic safetensors files with expert-like naming and verify
+    that the filter correctly skips non-local experts."""
+
+    @pytest.fixture
+    def synthetic_moe_files(self, tmp_path):
+        """Create synthetic safetensors with expert-patterned tensor names."""
+        from safetensors.torch import save_file
+
+        tensors = {}
+        # Dense weights
+        tensors["model.embed_tokens.weight"] = torch.randn(100, 64)
+        tensors["model.layers.0.self_attn.q_proj.weight"] = torch.randn(64, 64)
+        tensors["model.layers.0.input_layernorm.weight"] = torch.randn(64)
+        # Expert weights: 8 experts
+        for expert_id in range(8):
+            tensors[f"model.layers.0.mlp.experts.{expert_id}.gate_proj.weight"] = (
+                torch.randn(128, 64)
+            )
+            tensors[f"model.layers.0.mlp.experts.{expert_id}.up_proj.weight"] = (
+                torch.randn(128, 64)
+            )
+            tensors[f"model.layers.0.mlp.experts.{expert_id}.down_proj.weight"] = (
+                torch.randn(64, 128)
+            )
+        # Shared expert (should never be filtered)
+        tensors["model.layers.0.mlp.shared_experts.gate_proj.weight"] = torch.randn(
+            128, 64
+        )
+
+        filepath = str(tmp_path / "model-00001-of-00001.safetensors")
+        save_file(tensors, filepath)
+        return [filepath], tensors
+
+    def test_no_filter_returns_all(self, synthetic_moe_files):
+        files, expected = synthetic_moe_files
+        loaded = dict(safetensors_weights_iterator(files, False))
+        assert set(loaded.keys()) == set(expected.keys())
+
+    def test_ep2_rank0_gets_half_experts(self, synthetic_moe_files):
+        files, expected = synthetic_moe_files
+        # EP=2, rank=0 → experts 0-3
+        local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
+        loaded = dict(
+            safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
+        )
+
+        # Should have all dense + shared + experts 0-3 only
+        for name in loaded:
+            eid = parse_expert_id(name)
+            if eid is not None:
+                assert eid in local_ids, f"Non-local expert {eid} was loaded"
+
+        # Check expert count: 4 experts × 3 weights = 12
+        expert_names = [n for n in loaded if parse_expert_id(n) is not None]
+        assert len(expert_names) == 4 * 3
+
+        # Check all dense weights present
+        assert "model.embed_tokens.weight" in loaded
+        assert "model.layers.0.self_attn.q_proj.weight" in loaded
+        assert "model.layers.0.input_layernorm.weight" in loaded
+        assert "model.layers.0.mlp.shared_experts.gate_proj.weight" in loaded
+
+    def test_ep2_rank1_gets_other_half(self, synthetic_moe_files):
+        files, expected = synthetic_moe_files
+        local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=1)
+        loaded = dict(
+            safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
+        )
+
+        expert_names = [n for n in loaded if parse_expert_id(n) is not None]
+        assert len(expert_names) == 4 * 3
+        for name in expert_names:
+            assert parse_expert_id(name) in local_ids
+
+    def test_ep8_each_rank_gets_one_expert(self, synthetic_moe_files):
+        files, _ = synthetic_moe_files
+        all_expert_names = set()
+        for rank in range(8):
+            local_ids = compute_local_expert_ids(8, ep_size=8, ep_rank=rank)
+            loaded = dict(
+                safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
+            )
+            expert_names = {n for n in loaded if parse_expert_id(n) is not None}
+            # 1 expert × 3 weights
+            assert len(expert_names) == 3
+            all_expert_names |= expert_names
+
+        # All 8 experts × 3 weights covered across ranks
+        assert len(all_expert_names) == 24
+
+    def test_tensor_values_match(self, synthetic_moe_files):
+        """Filtered tensors have identical values to unfiltered ones."""
+        files, _ = synthetic_moe_files
+        all_weights = dict(safetensors_weights_iterator(files, False))
+
+        local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
+        filtered = dict(
+            safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
+        )
+
+        for name, tensor in filtered.items():
+            assert torch.equal(tensor, all_weights[name]), f"Tensor mismatch for {name}"

tests/model_executor/test_cpu_unquantized_gemm_dispatch.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for CPU unquantized GEMM dispatch behavior."""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers import utils
+from vllm.platforms import current_platform
+
+
+@pytest.fixture(scope="module")
+def _mock_zentorch_linear_unary():
+    """Register a mock zentorch_linear_unary op when zentorch is not installed.
+
+    Allows the dispatch tests to run in CI without a real zentorch build.
+    Skips registration when zentorch is already available.
+    """
+    if hasattr(torch.ops.zentorch, "zentorch_linear_unary"):
+        yield
+        return
+
+    lib_def = torch.library.Library("zentorch", "DEF")
+    lib_def.define(
+        "zentorch_linear_unary("
+        "Tensor input, "
+        "Tensor weight, "
+        "Tensor? bias, "
+        "bool is_weight_prepacked=False"
+        ") -> Tensor"
+    )
+
+    lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
+    lib_impl.impl(
+        "zentorch_linear_unary",
+        lambda input, weight, bias, is_weight_prepacked=False: (
+            torch.nn.functional.linear(input, weight, bias)
+        ),
+    )
+
+    yield
+
+    lib_impl._destroy()
+    lib_def._destroy()
+
+
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
+def test_dispatch_cpu_unquantized_gemm_uses_zentorch_on_zen(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
+
+    layer = torch.nn.Linear(16, 8, bias=True)
+    x = torch.randn(4, 16)
+    expected = torch.nn.functional.linear(x, layer.weight, layer.bias)
+
+    utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
+    output = layer.cpu_linear(x, layer.weight, layer.bias)
+
+    torch.testing.assert_close(output, expected)
+
+
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
+def test_dispatch_cpu_unquantized_gemm_zen_remove_weight(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
+
+    layer = torch.nn.Linear(16, 8, bias=True)
+    utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
+
+    assert layer.weight.numel() == 0

tests/models/language/generation/test_common.py

@@ -103,6 +103,10 @@ AITER_MODEL_LIST = [
             marks=[pytest.mark.core_model, pytest.mark.cpu_model],
         ),
         pytest.param("swiss-ai/Apertus-8B-Instruct-2509"),  # apertus
+        pytest.param(
+            "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",  # hyperclovax
+            marks=[large_gpu_mark(min_gb=32)],
+        ),
     ],
 )
 @pytest.mark.parametrize("max_tokens", [32])

tests/models/multimodal/generation/test_voxtral_realtime.py

@@ -1,8 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import contextlib
 from dataclasses import asdict
 
 import pytest
+import pytest_asyncio
 from mistral_common.audio import Audio
 from mistral_common.protocol.instruct.chunk import RawAudio
 from mistral_common.protocol.transcription.request import (
@@ -17,18 +19,21 @@ from vllm.assets.audio import AudioAsset
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.v1.engine.async_llm import AsyncLLM
 
+from ....utils import ROCM_ENGINE_KWARGS
+
 MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
-ENGINE_CONFIG = dict(
-    model=MODEL_NAME,
-    max_model_len=8192,
-    max_num_seqs=4,
-    limit_mm_per_prompt={"audio": 1},
-    config_format="mistral",
-    load_format="mistral",
-    tokenizer_mode="mistral",
-    enforce_eager=True,
-    gpu_memory_utilization=0.9,
-)
+ENGINE_CONFIG = {
+    "model": MODEL_NAME,
+    "max_model_len": 8192,
+    "max_num_seqs": 4,
+    "limit_mm_per_prompt": {"audio": 1},
+    "config_format": "mistral",
+    "load_format": "mistral",
+    "tokenizer_mode": "mistral",
+    "enforce_eager": True,
+    "gpu_memory_utilization": 0.9,
+    **ROCM_ENGINE_KWARGS,
+}
 
 
 EXPECTED_TEXT = [
@@ -49,6 +54,14 @@ EXPECTED_TEXT = [
 ]
 
 
+def _normalize(texts: list[str]) -> list[str]:
+    # The model occasionally transcribes "OBS" as "a base hit" and
+    # "oh, my" as "oh my", but both are acoustically valid. Normalise so
+    # the assertion is stable across runs and hardware.
+    texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
+    return texts
+
+
 @pytest.fixture
 def audio_assets() -> list[AudioAsset]:
     return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
@@ -60,15 +73,27 @@ def tokenizer() -> MistralTokenizer:
 
 
 @pytest.fixture
-def engine() -> LLM:
+def engine():
     engine_args = EngineArgs(**ENGINE_CONFIG)
-    return LLM(**asdict(engine_args))
+    llm = LLM(**asdict(engine_args))
+    try:
+        yield llm
+    finally:
+        with contextlib.suppress(Exception):
+            llm.llm_engine.engine_core.shutdown()
+        import torch
+
+        torch.accelerator.empty_cache()
 
 
-@pytest.fixture
-def async_engine() -> AsyncLLM:
+@pytest_asyncio.fixture
+async def async_engine():
     engine_args = AsyncEngineArgs(**ENGINE_CONFIG)
-    return AsyncLLM.from_engine_args(engine_args)
+    llm = AsyncLLM.from_engine_args(engine_args)
+    try:
+        yield llm
+    finally:
+        llm.shutdown()
 
 
 def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
@@ -108,8 +133,13 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
         sampling_params=sampling_params,
     )
 
-    texts = [out.outputs[0].text for out in outputs]
-    assert texts == EXPECTED_TEXT
+    texts = _normalize([out.outputs[0].text for out in outputs])
+    for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
+        assert got == expected, (
+            f"Output mismatch at index {i}:\n"
+            f"  got:      {got!r}\n"
+            f"  expected: {expected!r}"
+        )
 
 
 @pytest.mark.asyncio
@@ -149,9 +179,17 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine)
 
         output_tokens_list.append(output_tokens)
 
-    texts = [
-        tokenizer.decode(output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE)
-        for output_tokens in output_tokens_list
-    ]
-    texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
-    assert texts == EXPECTED_TEXT
+    texts = _normalize(
+        [
+            tokenizer.decode(
+                output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE
+            )
+            for output_tokens in output_tokens_list
+        ]
+    )
+    for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
+        assert got == expected, (
+            f"Output mismatch at index {i}:\n"
+            f"  got:      {got!r}\n"
+            f"  expected: {expected!r}"
+        )

tests/models/multimodal/pooling/test_clip.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
+import torch
 from transformers import CLIPModel
 
 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
@@ -50,13 +51,16 @@ def _run_test(
             if "pixel_values" in inputs:
                 pooled_output = hf_model.model.get_image_features(
                     pixel_values=inputs.pixel_values,
-                ).squeeze(0)
+                )
             else:
                 pooled_output = hf_model.model.get_text_features(
                     input_ids=inputs.input_ids,
                     attention_mask=inputs.attention_mask,
-                ).squeeze(0)
+                )
 
+            if not isinstance(pooled_output, torch.Tensor):
+                pooled_output = pooled_output.pooler_output
+            pooled_output = pooled_output.squeeze(0)
             all_outputs.append(pooled_output.tolist())
 
         hf_outputs = all_outputs

tests/models/multimodal/pooling/test_siglip.py

@@ -4,6 +4,7 @@
 from typing import Any
 
 import pytest
+import torch
 from transformers import SiglipModel
 
 from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
@@ -68,12 +69,15 @@ def _run_test(
             if "pixel_values" in inputs:
                 pooled_output = hf_model.model.get_image_features(
                     pixel_values=inputs.pixel_values,
-                ).squeeze(0)
+                )
             else:
                 pooled_output = hf_model.model.get_text_features(
                     input_ids=inputs.input_ids,
-                ).squeeze(0)
+                )
 
+            if not isinstance(pooled_output, torch.Tensor):
+                pooled_output = pooled_output.pooler_output
+            pooled_output = pooled_output.squeeze(0)
             all_outputs.append(pooled_output.tolist())
 
         hf_outputs = all_outputs

tests/models/multimodal/processing/test_audio_in_video.py

@@ -34,8 +34,22 @@ MODELS = [
 ]
 
 
+def create_mm_data(num_videos: int) -> dict[str, list]:
+    # Small video (8 frames, 64×64) and ~0.5 s of audio at 16 kHz so the test
+    # stays fast even without a GPU.
+    mm_data = dict[str, list](video=[], audio=[])
+    for i in range(num_videos):
+        rng = np.random.RandomState(i)
+        video = random_video(rng, min_frames=8, max_frames=9, min_wh=64, max_wh=65)
+        audio, sr = random_audio(rng, min_len=8000, max_len=8001, sr=16000)
+        mm_data["video"].append(video)
+        mm_data["audio"].append((audio, sr))
+    return mm_data
+
+
 @pytest.mark.parametrize("model_id", MODELS)
-def test_audio_in_video_cache_correctness(model_id: str) -> None:
+@pytest.mark.parametrize("num_videos", [1, 2])
+def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> None:
     """
     Regression test for https://github.com/vllm-project/vllm/pull/36800
 
@@ -47,7 +61,7 @@ def test_audio_in_video_cache_correctness(model_id: str) -> None:
     """
     ctx = build_model_context(
         model_id,
-        limit_mm_per_prompt={"audio": 1, "image": 0, "video": 1},
+        limit_mm_per_prompt={"audio": num_videos, "image": 0, "video": num_videos},
         mm_processor_cache_gb=1,
     )
 
@@ -65,17 +79,12 @@ def test_audio_in_video_cache_correctness(model_id: str) -> None:
 
     video_token_id = baseline_processor.info.get_hf_config().video_token_id
 
-    rng = np.random.RandomState(0)
-    # Small video (8 frames, 64×64) and ~0.5 s of audio at 16 kHz so the test
-    # stays fast even without a GPU.
-    video = random_video(rng, min_frames=8, max_frames=9, min_wh=64, max_wh=65)
-    audio, sr = random_audio(rng, min_len=8000, max_len=8001, sr=16000)
-    mm_data = {"video": [video], "audio": [(audio, sr)]}
+    mm_data = create_mm_data(num_videos)
     hf_processor_mm_kwargs = {"use_audio_in_video": True}
 
     def run(processor):
         return processor(
-            [video_token_id],
+            [video_token_id] * num_videos,
             mm_items=baseline_processor.info.parse_mm_data(mm_data),
             hf_processor_mm_kwargs=hf_processor_mm_kwargs,
         )["prompt_token_ids"]

tests/models/registry.py

@@ -320,7 +320,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "tencent/Hunyuan-A13B-Instruct", trust_remote_code=True
     ),
     "HyperCLOVAXForCausalLM": _HfExamplesInfo(
-        "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
+        "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
         trust_remote_code=True,
     ),
     "InternLMForCausalLM": _HfExamplesInfo(

tests/multimodal/media/test_audio.py

@@ -4,6 +4,7 @@ import base64
 from pathlib import Path
 from unittest.mock import patch
 
+import librosa
 import numpy as np
 import pytest
 
@@ -71,3 +72,13 @@ def test_audio_media_io_encode_base64(dummy_audio):
         decoded = base64.b64decode(out)
         assert decoded == b"dummy_wav_data"
         mock_write.assert_called_once()
+
+
+def test_audio_media_io_from_video(video_assets):
+    audio_io = AudioMediaIO()
+    video_path = video_assets[0].video_path
+    with open(video_path, "rb") as f:
+        audio, sr = audio_io.load_bytes(f.read())
+    audio_ref, sr_ref = librosa.load(video_path, sr=None)
+    assert sr == sr_ref
+    np.testing.assert_allclose(audio_ref, audio, atol=1e-4)

tests/reasoning/test_gptoss_reasoning_parser.py

@@ -1,11 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
+from unittest.mock import Mock
+
 import pytest
 from transformers import AutoTokenizer
 
+from vllm.entrypoints.mcp.tool_server import ToolServer
 from vllm.reasoning import ReasoningParser
-from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser
+from vllm.reasoning.gptoss_reasoning_parser import (
+    GptOssReasoningParser,
+    from_builtin_tool_to_tag,
+    no_func_reasoning_tag,
+)
 
 REASONING_MODEL_NAME = "openai/gpt-oss-120b"
 
@@ -142,3 +150,133 @@ def test_gptoss_is_reasoning_end(
     output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(output)
     actual_is_reasoning_end = parser.is_reasoning_end(output_ids)
     assert is_reasoning_end == actual_is_reasoning_end
+
+
+class TestGptOssStructuralTags:
+    """Test cases for GptOssReasoningParser structural tag functionality."""
+
+    @pytest.fixture
+    def mock_tokenizer(self):
+        """Create a mock tokenizer for testing."""
+        tokenizer = Mock()
+        tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
+        tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
+        return tokenizer
+
+    @pytest.fixture
+    def reasoning_parser(self, mock_tokenizer):
+        """Create a GptOssReasoningParser instance."""
+        return GptOssReasoningParser(mock_tokenizer)
+
+    def test_prepare_structured_tag_no_tool_server(self, reasoning_parser):
+        """Test prepare_structured_tag with no tool server."""
+        result = reasoning_parser.prepare_structured_tag(None, None)
+        expected = json.dumps(no_func_reasoning_tag)
+
+        assert result == expected
+
+        # Verify the structure is correct
+        parsed = json.loads(result)
+        assert parsed["type"] == "structural_tag"
+        assert parsed["format"]["type"] == "triggered_tags"
+        assert len(parsed["format"]["tags"]) == 1
+        assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>"
+        assert parsed["format"]["triggers"] == ["<|channel|>analysis"]
+
+    def test_prepare_structured_tag_with_original_tag(self, reasoning_parser):
+        """Test prepare_structured_tag when original_tag is provided."""
+        original_tag = '{"custom": "tag"}'
+        result = reasoning_parser.prepare_structured_tag(original_tag, None)
+
+        # Should return the original tag unchanged
+        assert result == original_tag
+
+    def test_from_builtin_tool_to_tag(self):
+        """Test from_builtin_tool_to_tag function."""
+        tags = from_builtin_tool_to_tag("python")
+
+        assert len(tags) == 2
+        assert tags[0]["begin"] == "<|channel|>commentary to=python"
+        assert tags[0]["content"]["type"] == "any_text"
+        assert tags[0]["end"] == "<|end|>"
+
+        assert tags[1]["begin"] == "<|channel|>analysis to=python"
+        assert tags[1]["content"]["type"] == "any_text"
+        assert tags[1]["end"] == "<|end|>"
+
+    @pytest.mark.parametrize(
+        "tools",
+        [
+            [],
+            ["browser"],
+            ["python"],
+            ["container"],
+            ["browser", "python"],
+            ["browser", "container"],
+            ["python", "container"],
+            ["browser", "python", "container"],
+        ],
+    )
+    def test_json_validity_comprehensive(self, reasoning_parser, tools):
+        """Test JSON validity across all possible tool combinations."""
+        tool_server = Mock(spec=ToolServer)
+        tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools)
+
+        result = reasoning_parser.prepare_structured_tag(None, tool_server)
+        parsed_result = json.loads(result)
+
+        assert parsed_result["type"] == "structural_tag"
+        assert "format" in parsed_result
+        assert "tags" in parsed_result["format"]
+        assert "triggers" in parsed_result["format"]
+
+        # Tag count should be: 1 (analysis) + 2 * len(tools)
+        expected_tag_count = 1 + (2 * len(tools))
+        assert len(parsed_result["format"]["tags"]) == expected_tag_count
+
+        # Verify triggers are correctly configured
+        expected_triggers = ["<|channel|>analysis"]
+        if tools:
+            expected_triggers.append("<|channel|>commentary to=")
+        assert set(parsed_result["format"]["triggers"]) == set(expected_triggers)
+
+    def test_no_cross_request_state_pollution(self, reasoning_parser):
+        """Test that sequential calls with different tool servers produce
+        independent results, guarding against shared mutable state
+        (e.g. missing deepcopy in tag_with_builtin_funcs)."""
+        tool_server_1 = Mock(spec=ToolServer)
+        tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python")
+
+        tool_server_2 = Mock(spec=ToolServer)
+        tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser")
+
+        result_1 = reasoning_parser.prepare_structured_tag(None, tool_server_1)
+        result_2 = reasoning_parser.prepare_structured_tag(None, tool_server_2)
+
+        tags_1 = [tag["begin"] for tag in json.loads(result_1)["format"]["tags"]]
+        tags_2 = [tag["begin"] for tag in json.loads(result_2)["format"]["tags"]]
+
+        assert "<|channel|>commentary to=python" in tags_1
+        assert "<|channel|>commentary to=browser" not in tags_1
+
+        assert "<|channel|>commentary to=browser" in tags_2
+        assert "<|channel|>commentary to=python" not in tags_2
+
+    def test_tag_format_consistency(self, reasoning_parser):
+        """Test that all generated tags follow consistent format,
+        catching malformed tags from from_builtin_tool_to_tag."""
+        tool_server = Mock(spec=ToolServer)
+        tool_server.has_tool = Mock(
+            side_effect=lambda tool: tool in ["python", "browser"]
+        )
+
+        result = reasoning_parser.prepare_structured_tag(None, tool_server)
+        parsed_result = json.loads(result)
+
+        for tag in parsed_result["format"]["tags"]:
+            assert "begin" in tag
+            assert "content" in tag
+            assert "end" in tag
+            assert tag["content"]["type"] == "any_text"
+            assert tag["end"] == "<|end|>"
+            assert tag["begin"].startswith("<|channel|>")

tests/test_regression.py

@@ -55,7 +55,7 @@ def test_gc():
     # The memory allocated for model and KV cache should be released.
     # The memory allocated for PyTorch and others should be less than 50MB.
     # Usually, it's around 10MB.
-    allocated = torch.cuda.memory_allocated()
+    allocated = torch.accelerator.memory_allocated()
     assert allocated < 50 * 1024 * 1024
 
 

tests/test_zen_cpu_platform_detection.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from unittest.mock import mock_open, patch
+
+from vllm.platforms import _is_amd_zen_cpu
+
+
+def test_is_amd_zen_cpu_detects_amd_with_avx512():
+    cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2 avx512f avx512bw"
+    with (
+        patch("os.path.exists", return_value=True),
+        patch("builtins.open", mock_open(read_data=cpuinfo)),
+    ):
+        assert _is_amd_zen_cpu()
+
+
+def test_is_amd_zen_cpu_returns_false_for_amd_without_avx512():
+    cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2"
+    with (
+        patch("os.path.exists", return_value=True),
+        patch("builtins.open", mock_open(read_data=cpuinfo)),
+    ):
+        assert not _is_amd_zen_cpu()
+
+
+def test_is_amd_zen_cpu_returns_false_for_intel_with_avx512():
+    cpuinfo = "vendor_id: GenuineIntel\nflags: avx avx2 avx512f"
+    with (
+        patch("os.path.exists", return_value=True),
+        patch("builtins.open", mock_open(read_data=cpuinfo)),
+    ):
+        assert not _is_amd_zen_cpu()
+
+
+def test_is_amd_zen_cpu_returns_false_when_cpuinfo_missing():
+    with patch("os.path.exists", return_value=False):
+        assert not _is_amd_zen_cpu()

tests/tool_parsers/test_glm4_moe_tool_parser.py

@@ -560,19 +560,23 @@ def test_streaming_empty_tool_call(glm4_moe_tool_parser, mock_request):
     assert glm4_moe_tool_parser.current_tool_id == -1
 
 
-def test_streaming_prev_tool_call_arr_finalization(glm4_moe_tool_parser, mock_request):
+def test_streaming_prev_tool_call_arr_updates(glm4_moe_tool_parser, mock_request):
     """Test that prev_tool_call_arr contains parsed dict after tool call."""
     _reset_streaming_state(glm4_moe_tool_parser)
 
     # Stream a complete tool call
+    name_only = {"name": "get_weather", "arguments": {}}
+    name_and_args = {"name": "get_weather", "arguments": {"city": "Beijing"}}
     chunks = [
-        "<tool_call>get_weather\n",
-        "<arg_key>city</arg_key>",
-        "<arg_value>Beijing</arg_value>",
-        "</tool_call>",
+        # Delta, expected streamed_args_for_tool, expected prev_tool_call_arr
+        ("<tool_call>get_weather\n", "", name_only),
+        ("<arg_key>city</arg_key>", "", name_only),
+        ("<arg_value>Beijing</arg_value>", '{"city": "Beijing"', name_only),
+        # Note: arguments are only updated when the tool call is complete.
+        ("</tool_call>", '{"city": "Beijing"}', name_and_args),
     ]
 
-    for chunk in chunks:
+    for chunk, exp_streamed, exp_prev_tc in chunks:
         glm4_moe_tool_parser.extract_tool_calls_streaming(
             previous_text="",
             current_text="",
@@ -582,6 +586,8 @@ def test_streaming_prev_tool_call_arr_finalization(glm4_moe_tool_parser, mock_re
             delta_token_ids=[],
             request=mock_request,
         )
+        assert glm4_moe_tool_parser.streamed_args_for_tool[0] == exp_streamed
+        assert glm4_moe_tool_parser.prev_tool_call_arr[0] == exp_prev_tc
 
     # After the tool call completes, prev_tool_call_arr should have parsed dict
     assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1
@@ -592,6 +598,12 @@ def test_streaming_prev_tool_call_arr_finalization(glm4_moe_tool_parser, mock_re
     assert isinstance(args, dict), f"Expected dict, got {type(args)}"
     assert args.get("city") == "Beijing"
 
+    # Test equivalence of prev_tool_call_arr and streamed_args_for_tool
+    # Simulates logic in chat_completion/serving.py:chat_completion_stream_generator
+    tool_call_json = json.dumps(tool_entry.get("arguments", {}))
+    streamed_content = glm4_moe_tool_parser.streamed_args_for_tool[0]
+    assert tool_call_json.startswith(streamed_content)
+
 
 def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_request):
     """Test streaming multiple sequential tool calls."""

tests/transformers_utils/test_utils.py

@@ -11,6 +11,7 @@ from vllm.transformers_utils.gguf_utils import (
     split_remote_gguf,
 )
 from vllm.transformers_utils.utils import (
+    is_azure,
     is_cloud_storage,
     is_gcs,
     is_s3,
@@ -31,9 +32,17 @@ def test_is_s3():
     assert not is_s3("nfs://nfs-fqdn.local")
 
 
+def test_is_azure():
+    assert is_azure("az://model-container/path")
+    assert not is_azure("s3://model-path/path-to-model")
+    assert not is_azure("/unix/local/path")
+    assert not is_azure("nfs://nfs-fqdn.local")
+
+
 def test_is_cloud_storage():
     assert is_cloud_storage("gs://model-path")
     assert is_cloud_storage("s3://model-path/path-to-model")
+    assert is_cloud_storage("az://model-container/path")
     assert not is_cloud_storage("/unix/local/path")
     assert not is_cloud_storage("nfs://nfs-fqdn.local")
 

tests/utils.py

@@ -122,6 +122,12 @@ ROCM_EXTRA_ARGS = (
     if current_platform.is_rocm()
     else []
 )
+# Python-API equivalent of ROCM_EXTRA_ARGS for use with EngineArgs kwargs.
+ROCM_ENGINE_KWARGS: dict = (
+    {"enable_prefix_caching": False, "max_num_seqs": 1}
+    if current_platform.is_rocm()
+    else {}
+)
 
 
 class RemoteVLLMServer:

tests/utils_/test_mem_utils.py

@@ -29,7 +29,7 @@ def test_memory_profiling():
     def measure_current_non_torch():
         free, total = torch.cuda.mem_get_info()
         current_used = total - free
-        current_torch = torch.cuda.memory_reserved()
+        current_torch = torch.accelerator.memory_reserved()
         current_non_torch = current_used - current_torch
         return current_non_torch
 

tests/v1/attention/test_gdn_metadata_builder.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for GDNAttentionMetadataBuilder.build() — specifically the
+reclassification of non-spec decodes as prefills when spec decodes exist.
+Covers the fix for https://github.com/vllm-project/vllm/issues/34845.
+"""
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from tests.v1.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+    create_vllm_config,
+)
+from vllm.config import SpeculativeConfig
+from vllm.v1.attention.backends.gdn_attn import (
+    GDNAttentionMetadata,
+    GDNAttentionMetadataBuilder,
+)
+from vllm.v1.kv_cache_interface import MambaSpec
+
+BLOCK_SIZE = 16
+DEVICE = torch.device("cpu")
+
+
+@dataclass
+class GDNBuildTestCase:
+    """Specification for a GDN metadata builder classification test."""
+
+    seq_lens: list[int]
+    query_lens: list[int]
+    num_decode_draft_tokens: list[int] | None  # None = no spec config
+    num_speculative_tokens: int
+    expected_num_decodes: int
+    expected_num_prefills: int
+    expected_num_prefill_tokens: int
+    expected_num_spec_decodes: int
+
+
+GDN_BUILD_TEST_CASES = {
+    # The original #34845 crash: non-spec query_len=1 + spec decode
+    "mixed_decode_and_spec_decode": GDNBuildTestCase(
+        seq_lens=[65, 20],
+        query_lens=[1, 3],
+        num_decode_draft_tokens=[-1, 2],
+        num_speculative_tokens=2,
+        expected_num_decodes=0,
+        expected_num_prefills=1,
+        expected_num_prefill_tokens=1,
+        expected_num_spec_decodes=1,
+    ),
+    # All requests are spec decodes — no reclassification needed
+    "pure_spec_decode": GDNBuildTestCase(
+        seq_lens=[50, 30],
+        query_lens=[3, 3],
+        num_decode_draft_tokens=[2, 2],
+        num_speculative_tokens=2,
+        expected_num_decodes=0,
+        expected_num_prefills=0,
+        expected_num_prefill_tokens=0,
+        expected_num_spec_decodes=2,
+    ),
+    # No speculative config at all — standard decode path
+    "pure_regular_decode": GDNBuildTestCase(
+        seq_lens=[40, 30, 20],
+        query_lens=[1, 1, 1],
+        num_decode_draft_tokens=None,
+        num_speculative_tokens=0,
+        expected_num_decodes=3,
+        expected_num_prefills=0,
+        expected_num_prefill_tokens=0,
+        expected_num_spec_decodes=0,
+    ),
+    # Multi-token prefill alongside spec decode — no decode to reclassify
+    "spec_decode_with_real_prefill": GDNBuildTestCase(
+        seq_lens=[100, 20],
+        query_lens=[50, 3],
+        num_decode_draft_tokens=[-1, 2],
+        num_speculative_tokens=2,
+        expected_num_decodes=0,
+        expected_num_prefills=1,
+        expected_num_prefill_tokens=50,
+        expected_num_spec_decodes=1,
+    ),
+    # All three types in one batch — decode gets reclassified
+    "prefill_decode_and_spec_decode": GDNBuildTestCase(
+        seq_lens=[100, 65, 20],
+        query_lens=[50, 1, 3],
+        num_decode_draft_tokens=[-1, -1, 2],
+        num_speculative_tokens=2,
+        expected_num_decodes=0,
+        expected_num_prefills=2,
+        expected_num_prefill_tokens=51,
+        expected_num_spec_decodes=1,
+    ),
+    # Multiple non-spec query_len=1 requests all reclassified
+    "multiple_decodes_reclassified": GDNBuildTestCase(
+        seq_lens=[40, 50, 60, 20],
+        query_lens=[1, 1, 1, 3],
+        num_decode_draft_tokens=[-1, -1, -1, 2],
+        num_speculative_tokens=2,
+        expected_num_decodes=0,
+        expected_num_prefills=3,
+        expected_num_prefill_tokens=3,
+        expected_num_spec_decodes=1,
+    ),
+    # Zero-length padded sequence excluded from counts
+    "zero_length_padding_with_spec": GDNBuildTestCase(
+        seq_lens=[16, 65, 20],
+        query_lens=[0, 1, 3],
+        num_decode_draft_tokens=[-1, -1, 2],
+        num_speculative_tokens=2,
+        expected_num_decodes=0,
+        expected_num_prefills=1,
+        expected_num_prefill_tokens=1,
+        expected_num_spec_decodes=1,
+    ),
+}
+
+
+def _create_gdn_builder(
+    num_speculative_tokens: int = 0,
+) -> GDNAttentionMetadataBuilder:
+    """Create a GDNAttentionMetadataBuilder with minimal config."""
+    vllm_config = create_vllm_config(block_size=BLOCK_SIZE)
+    if num_speculative_tokens > 0:
+        vllm_config.speculative_config = SpeculativeConfig(
+            method="ngram",
+            num_speculative_tokens=num_speculative_tokens,
+        )
+    mamba_spec = MambaSpec(
+        block_size=BLOCK_SIZE,
+        shapes=((16, 64),),
+        dtypes=(torch.float16,),
+    )
+    return GDNAttentionMetadataBuilder(
+        kv_cache_spec=mamba_spec,
+        layer_names=["layer.0"],
+        vllm_config=vllm_config,
+        device=DEVICE,
+    )
+
+
+def _build(
+    builder: GDNAttentionMetadataBuilder,
+    batch_spec: BatchSpec,
+    num_decode_draft_tokens: list[int] | None = None,
+) -> GDNAttentionMetadata:
+    """Build GDN attention metadata, optionally with spec-decode kwargs."""
+    common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE)
+    kwargs: dict = {}
+    if num_decode_draft_tokens is not None:
+        kwargs["num_decode_draft_tokens_cpu"] = torch.tensor(
+            num_decode_draft_tokens, dtype=torch.int32
+        )
+        kwargs["num_accepted_tokens"] = torch.ones(
+            batch_spec.batch_size, dtype=torch.int32, device=DEVICE
+        )
+    return builder.build(common_prefix_len=0, common_attn_metadata=common, **kwargs)
+
+
+@pytest.mark.parametrize(
+    "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys()
+)
+def test_gdn_build_classification(test_case: GDNBuildTestCase):
+    """Test that GDN metadata builder classifies requests correctly."""
+    builder = _create_gdn_builder(test_case.num_speculative_tokens)
+    batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens)
+    meta = _build(builder, batch, test_case.num_decode_draft_tokens)
+
+    assert meta.num_decodes == test_case.expected_num_decodes
+    assert meta.num_prefills == test_case.expected_num_prefills
+    assert meta.num_prefill_tokens == test_case.expected_num_prefill_tokens
+    assert meta.num_spec_decodes == test_case.expected_num_spec_decodes
+
+
+def test_has_initial_state_after_reclassification():
+    """After reclassification, num_prefills > 0 so the prefill kernel path
+    should compute has_initial_state. For the reclassified request with
+    context_lens > 0, the corresponding entry must be True."""
+    builder = _create_gdn_builder(num_speculative_tokens=2)
+    batch = BatchSpec(seq_lens=[65, 20], query_lens=[1, 3])
+    meta = _build(builder, batch, num_decode_draft_tokens=[-1, 2])
+
+    assert meta.num_prefills > 0, "reclassification should produce prefills"
+    assert meta.has_initial_state is not None
+    # req0 has context_lens = 65 - 1 = 64 > 0, so has_initial_state[0] = True
+    assert meta.has_initial_state[0].item() is True

tests/v1/e2e/general/test_async_scheduling.py

@@ -7,7 +7,10 @@ from typing import Any
 import pytest
 import torch._dynamo.config as dynamo_config
 
-from tests.utils import large_gpu_mark, single_gpu_only
+from tests.utils import (
+    large_gpu_mark,
+    single_gpu_only,
+)
 from vllm import SamplingParams
 from vllm.logprobs import Logprob
 from vllm.platforms import current_platform
@@ -150,6 +153,7 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke
     run_tests(monkeypatch, MTP_MODEL, test_configs, test_sampling_params)
 
 
+@pytest.mark.flaky(reruns=2, only_on=current_platform.is_rocm())
 def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
     """Test ngram_gpu speculative decoding with different configurations.
 
@@ -202,7 +206,6 @@ def run_tests(
     with monkeypatch.context() as m:
         # lock matmul precision to full FP32 (IEEE)
         m.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "highest")
-        # m.setenv("VLLM_BATCH_INVARIANT", "1")
         outputs: list[tuple[str, list, list]] = []
         for n, (
             test_preemption,
@@ -351,6 +354,7 @@ def run_test(
         speculative_config=spec_config,
         disable_log_stats=False,
         attention_config=attention_config,
+        enable_prefix_caching=False if current_platform.is_rocm() else None,
         **cache_arg,
     ) as vllm_model:
         results = []

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

@@ -18,11 +18,19 @@ dp_ep_configs=(
 "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
 "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
 )
+hybrid_ssm_configs=(
+  "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
+  # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
+  "ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
+)
 
 # Select config array based on DP_EP env var
 if [[ -n "${DP_EP:-}" ]]; then
   configs=("${dp_ep_configs[@]}")
   echo "DP_EP is set, using dp_ep_configs"
+elif [[ -n "${HYBRID_SSM:-}" ]]; then
+  configs=("${hybrid_ssm_configs[@]}")
+  echo "HYBRID_SSM is set, using hybrid_ssm_configs."
 else
   configs=("${tp_configs[@]}")
 fi

tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

@@ -21,6 +21,11 @@
 #   MODEL_NAME          - target model (default: meta-llama/Llama-3.1-8B-Instruct)
 #   NUM_SPEC_TOKENS     - number of speculative tokens (default: 3)
 #   GPU_MEMORY_UTILIZATION - (default: 0.7)
+#   ATTENTION_BACKEND   - attention backend to use
+#                         Default: TRITON_ATTN on ROCm, FLASH_ATTN on NVIDIA
+#                         ROCm options: TRITON_ATTN, ROCM_ATTN, ROCM_AITER_FA,
+#                                       ROCM_AITER_UNIFIED_ATTN
+#                         NVIDIA options: FLASH_ATTN, FLASHINFER
 set -x
 
 # ── Model & spec decode config ──────────────────────────────────────────
@@ -51,6 +56,28 @@ GIT_ROOT=$(git rev-parse --show-toplevel)
 
 SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
 
+# ── Detect platform (NVIDIA vs ROCm) ────────────────────────────────────
+
+if [[ "$SMI_BIN" == *"rocm"* ]]; then
+  GPU_PLATFORM="rocm"
+  GPU_DEVICE_VAR="HIP_VISIBLE_DEVICES"
+else
+  GPU_PLATFORM="nvidia"
+  GPU_DEVICE_VAR="CUDA_VISIBLE_DEVICES"
+fi
+echo "Detected GPU platform: ${GPU_PLATFORM} (using ${GPU_DEVICE_VAR})"
+
+# ── Attention backend config ─────────────────────────────────────────────
+
+if [[ -z "${ATTENTION_BACKEND:-}" ]]; then
+  if [[ "$GPU_PLATFORM" == "rocm" ]]; then
+    ATTENTION_BACKEND="TRITON_ATTN"
+  else
+    ATTENTION_BACKEND="FLASH_ATTN"
+  fi
+fi
+echo "Using attention backend: ${ATTENTION_BACKEND}"
+
 cleanup_instances() {
   echo ""
   echo "Cleaning up..."
@@ -84,13 +111,16 @@ wait_for_server() {
 
 # ── Resolve GPU list ─────────────────────────────────────────────────────
 
-if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
-  IFS=',' read -ra ALL_GPUS <<< "$CUDA_VISIBLE_DEVICES"
+# Accept either CUDA_VISIBLE_DEVICES or HIP_VISIBLE_DEVICES
+VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES:-}}"
+
+if [[ -n "${VISIBLE_DEVICES}" ]]; then
+  IFS=',' read -ra ALL_GPUS <<< "$VISIBLE_DEVICES"
 else
   ALL_GPUS=()
-  if [[ "$SMI_BIN" == *"nvidia"* ]]; then
+  if [[ "$GPU_PLATFORM" == "nvidia" ]]; then
     num=$($SMI_BIN --query-gpu=name --format=csv,noheader | wc -l)
-  elif [[ "$SMI_BIN" == *"rocm"* ]]; then
+  elif [[ "$GPU_PLATFORM" == "rocm" ]]; then
     num=$($SMI_BIN -l | grep -c GPU)
   else
     num=1
@@ -100,7 +130,7 @@ fi
 
 TOTAL_GPUS_NEEDED=$(( (NUM_PREFILL_INSTANCES * PREFILLER_TP_SIZE) + (NUM_DECODE_INSTANCES * DECODER_TP_SIZE) ))
 if [[ ${#ALL_GPUS[@]} -lt $TOTAL_GPUS_NEEDED ]]; then
-  echo "FAIL: Need $TOTAL_GPUS_NEEDED GPUs but only have ${#ALL_GPUS[@]} (CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set})"
+  echo "FAIL: Need $TOTAL_GPUS_NEEDED GPUs but only have ${#ALL_GPUS[@]} (visible devices=${VISIBLE_DEVICES:-not set})"
   exit 1
 fi
 
@@ -119,12 +149,14 @@ run_test_for_device() {
   echo "================================================================"
   echo "NixlConnector PD + Spec Decode Acceptance Test (kv_buffer_device=${kv_device})"
   echo "================================================================"
-  echo "Model:            ${MODEL_NAME}"
-  echo "SD method:        ${SD_METHOD}"
-  echo "SD model:         ${SD_MODEL}"
-  echo "Spec tokens:      ${NUM_SPEC_TOKENS}"
-  echo "KV buffer device: ${kv_device}"
-  echo "GPUs available:   ${ALL_GPUS[*]}"
+  echo "Model:              ${MODEL_NAME}"
+  echo "SD method:          ${SD_METHOD}"
+  echo "SD model:           ${SD_MODEL}"
+  echo "Spec tokens:        ${NUM_SPEC_TOKENS}"
+  echo "KV buffer device:   ${kv_device}"
+  echo "Attention backend:  ${ATTENTION_BACKEND}"
+  echo "GPU platform:       ${GPU_PLATFORM}"
+  echo "GPUs available:     ${ALL_GPUS[*]}"
   echo "================================================================"
 
   local PREFILL_HOSTS=()
@@ -146,7 +178,8 @@ run_test_for_device() {
     local SIDE_CHANNEL_PORT=$((5559 + i))
 
     echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
-    CUDA_VISIBLE_DEVICES=$GPU_ID \
+    env \
+    ${GPU_DEVICE_VAR}=$GPU_ID \
     VLLM_KV_CACHE_LAYOUT='HND' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
@@ -159,7 +192,7 @@ run_test_for_device() {
       --tensor-parallel-size $PREFILLER_TP_SIZE \
       --kv-transfer-config "$kv_config" \
       --speculative-config "$PREFILL_SPEC_CONFIG" \
-      --attention-backend FLASH_ATTN &
+      --attention-backend $ATTENTION_BACKEND &
 
     PREFILL_HOSTS+=("localhost")
     PREFILL_PORTS+=("$PORT")
@@ -178,7 +211,8 @@ run_test_for_device() {
     local SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
 
     echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
-    CUDA_VISIBLE_DEVICES=$GPU_ID \
+    env \
+    ${GPU_DEVICE_VAR}=$GPU_ID \
     VLLM_KV_CACHE_LAYOUT='HND' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
@@ -191,7 +225,7 @@ run_test_for_device() {
       --tensor-parallel-size $DECODER_TP_SIZE \
       --kv-transfer-config "$kv_config" \
       --speculative-config "$DECODE_SPEC_CONFIG" \
-      --attention-backend FLASH_ATTN &
+      --attention-backend $ATTENTION_BACKEND &
 
     DECODE_HOSTS+=("localhost")
     DECODE_PORTS+=("$PORT")
@@ -218,7 +252,7 @@ run_test_for_device() {
   sleep 5
 
   # Run test
-  echo "Running spec decode acceptance test (kv_buffer_device=${kv_device})..."
+  echo "Running spec decode acceptance test (kv_buffer_device=${kv_device}, backend=${ATTENTION_BACKEND})..."
   DECODE_PORT=${DECODE_PORTS[0]} \
   TEST_MODEL=$MODEL_NAME \
   python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py"
@@ -234,4 +268,4 @@ for device in $KV_BUFFER_DEVICES; do
   run_test_for_device "$device"
 done
 
-echo "=== All spec decode acceptance tests passed ==="
+echo "=== All spec decode acceptance tests passed (backend=${ATTENTION_BACKEND}) ==="

tests/v1/kv_connector/nixl_integration/test_accuracy.py

@@ -18,6 +18,7 @@ EXPECTED_VALUES = {
     "deepseek-ai/deepseek-vl2-tiny": 0.19,
     "deepseek-ai/DeepSeek-V2-Lite-Chat": 0.65,
     "google/gemma-3-4b-it": 0.74,
+    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8": 0.84,
 }
 
 SIMPLE_PROMPT = (

tests/v1/kv_connector/unit/test_kv_cache_layout.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def test_mla_backend_rejects_cross_layer_kv_cache():
+    """MLA backends return identity permutation (layers dim first)
+    to signal cross-layer KV cache is unsupported."""
+    from vllm.model_executor.layers.attention.mla_attention import (
+        MLACommonBackend,
+    )
+
+    stride_order = MLACommonBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=True
+    )
+    assert stride_order == (0, 1, 2, 3)
+    assert stride_order[0] == 0  # layers dim first => no cross-layer
+    assert MLACommonBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=False
+    ) == (0, 1, 2)
+
+
+def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
+    """DeepseekV32Indexer returns identity permutation (layers dim first)
+    to signal cross-layer KV cache is unsupported."""
+    from vllm.v1.attention.backends.mla.indexer import (
+        DeepseekV32IndexerBackend,
+    )
+
+    stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=True
+    )
+    assert stride_order == (0, 1, 2, 3)
+    assert stride_order[0] == 0  # layers dim first => no cross-layer
+    assert DeepseekV32IndexerBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=False
+    ) == (0, 1, 2)

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -84,10 +84,13 @@ def mock_parallel_groups():
         yield mock_group
 
 
-def _setup_kv_transfer_request(request, remote_host="127.0.0.1", fake_port=4789):
+def _setup_kv_transfer_request(
+    request, remote_host="127.0.0.1", fake_port=4789, fake_transfer_id="0"
+):
     """Setup KV transfer parameters for a request."""
     request.kv_transfer_params.update(
         {
+            "transfer_id": fake_transfer_id,
             "remote_notify_port": fake_port,
             "remote_block_ids": None,
             "remote_host": remote_host,

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -53,7 +53,13 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
 from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.output_processor import OutputProcessor
-from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+)
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 from vllm.v1.request import RequestStatus
 from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
@@ -332,8 +338,20 @@ def test_kv_transfer_handshake(dist_init):
 
         # Prefill connector will register KV cache to populate proper handshake
         # metadata.
-        # TODO this must match with values used in kv cache config
-        kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
+        kv_cache_groups = [
+            KVCacheGroupSpec(
+                ["layer0", "layer1", "layer2"],
+                FullAttentionSpec(
+                    block_size=16,
+                    num_kv_heads=4,
+                    head_size=16,
+                    dtype=torch.float16,
+                ),
+            )
+        ]
+        kv_cache_config = KVCacheConfig(
+            num_blocks=2, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
+        )
         prefill_connector = NixlConnector(
             vllm_config, KVConnectorRole.WORKER, kv_cache_config
         )
@@ -437,7 +455,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         self.kv_cache_layout = kv_cache_layout
         # Mock register_kv_caches attribute needed for tests that do not call it.
         self.src_xfer_handles_by_block_size = {self.block_size: 1}
-        test_shape = self.attn_backend.get_kv_cache_shape(
+        test_shape = self.attn_backends[0].get_kv_cache_shape(
             num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
         )
         self.kv_topo = TpKVTopology(
@@ -447,7 +465,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
             remote_block_size=self._block_size,  # shared state
             is_mla=self.use_mla,
             total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
-            attn_backend=self.attn_backend,
+            attn_backends=self.attn_backends,
             tensor_shape=test_shape,
         )
 
@@ -501,6 +519,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
                     # is started. We mock HND here.
                     kv_cache_layout="HND",
                     block_size=self.block_size,
+                    ssm_sizes=(0, 0),
                 ),
                 remote_tp_rank=remote_tp_rank,
                 remote_tp_size=remote_tp_size,
@@ -951,6 +970,7 @@ class TestNixlHandshake:
                 block_lens=worker.block_len_per_layer,
                 kv_cache_layout=mismatched_layout,
                 block_size=worker.block_size,
+                ssm_sizes=(0, 0),
             )
 
             with pytest.raises(RuntimeError):
@@ -1006,6 +1026,7 @@ class TestNixlHandshake:
                 block_lens=[i * 2 for i in worker.block_len_per_layer],
                 kv_cache_layout="HND",
                 block_size=worker.block_size,
+                ssm_sizes=(0, 0),
             )
 
             # We don't check layout for homogeneous TP and MLA for now, as the
@@ -1496,9 +1517,47 @@ def test_register_kv_caches(
         # test run if not mocking.
         mock_get_attn_backend.return_value = backend_cls
         mock_get_attn_backends.return_value = [backend_cls]
+        num_layers = 32
+        block_size = 16
+        num_blocks = 8
+        num_heads = 4
+        head_size = 16
 
+        # TODO (NickLucche) the fact that connector depends on kv_cache_config for init
+        # but cross-layer preference cant be inferred prior to creating kv_cache_config
+        # is a bit awkward.
+        dummy_connector = NixlConnector(
+            vllm_config,
+            KVConnectorRole.WORKER,
+            make_kv_cache_config(block_size=block_size),
+        )
+        kv_cache_spec = FullAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=num_heads,
+            head_size=head_size,
+            dtype=torch.float16,
+        )
+        if dummy_connector.prefer_cross_layer_blocks:
+            kv_cache_config = KVCacheConfig(
+                num_blocks=num_blocks,
+                kv_cache_tensors=[
+                    KVCacheTensor(
+                        size=kv_cache_spec.page_size_bytes * num_blocks,
+                        shared_by=["all-layers"],
+                    )
+                    for _ in range(num_layers)
+                ],
+                kv_cache_groups=[KVCacheGroupSpec(["all-layers"], kv_cache_spec)],
+            )
+        else:
+            kv_cache_config = KVCacheConfig(
+                num_blocks=num_blocks,
+                kv_cache_tensors=[],
+                kv_cache_groups=[
+                    KVCacheGroupSpec(["layer0", "layer1", "layer2"], kv_cache_spec)
+                ],
+            )
         # Create connector
-        kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
         connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
         connector.connector_worker = FakeNixlConnectorWorker(
             vllm_config,
@@ -1526,35 +1585,6 @@ def test_register_kv_caches(
             or connector.prefer_cross_layer_blocks
         )
         if connector.prefer_cross_layer_blocks:
-            num_layers = 32
-            block_size = 16
-            num_blocks = 8
-            # Keep the fake worker's expected num_blocks in sync with the
-            # cross-layer tensor we are about to register.
-            worker_kv_cache_config = make_kv_cache_config(
-                block_size=block_size, num_blocks=num_blocks
-            )
-            connector.connector_worker.kv_cache_config = worker_kv_cache_config
-            connector.connector_worker.num_blocks = worker_kv_cache_config.num_blocks
-            kv_cache_spec = AttentionSpec(
-                block_size=block_size,
-                num_kv_heads=4,
-                head_size=64,
-                dtype=torch.bfloat16,
-            )
-            kv_cache_config = KVCacheConfig(
-                num_blocks=num_blocks,
-                kv_cache_tensors=[
-                    KVCacheTensor(
-                        size=kv_cache_spec.page_size_bytes * num_blocks,
-                        shared_by=["dummy-layer"],
-                    )
-                    for i in range(num_layers)
-                ],
-                # allocate_uniform_kv_caches does not use this
-                kv_cache_groups=[],
-            )
-
             with set_current_vllm_config(vllm_config):
                 _, cross_layers_kv_cache, _ = (
                     KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
@@ -1586,12 +1616,8 @@ def test_register_kv_caches(
             expected_blocks_count = 8
 
             kv_caches = {"all-layers": cross_layers_kv_cache}
-
         else:
             # Create test kv cache tensors using proper backend shape
-            kv_cache_spec = cast(
-                AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
-            )
             kv_cache_shape = backend_cls.get_kv_cache_shape(
                 num_blocks=kv_cache_config.num_blocks,
                 block_size=kv_cache_spec.block_size,
@@ -2261,7 +2287,7 @@ def test_compatibility_hash_validation(
     kv_cache_spec = cast(
         AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
     )
-    kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
+    kv_cache_shape = decode_worker.attn_backends[0].get_kv_cache_shape(
         num_blocks=kv_cache_config.num_blocks,
         block_size=kv_cache_spec.block_size,
         num_kv_heads=kv_cache_spec.num_kv_heads,
@@ -2269,10 +2295,14 @@ def test_compatibility_hash_validation(
     )
     shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
     unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
+    # Build kv_caches from the actual layer names in kv_cache_config so that
+    # _layer_specs lookups in register_kv_caches always find a matching key.
+    layer_names = [
+        name for group in kv_cache_config.kv_cache_groups for name in group.layer_names
+    ]
     kv_caches = {
-        "layer0": shared_tensor,
-        "layer1": unique_tensor,
-        "layer2": shared_tensor,
+        name: shared_tensor if i % 2 == 0 else unique_tensor
+        for i, name in enumerate(layer_names)
     }
     decode_connector.register_kv_caches(kv_caches)
 
@@ -2312,6 +2342,7 @@ def test_compatibility_hash_validation(
         block_lens=[4096 * prefill_block_size],  # slot_size * block_size
         kv_cache_layout="HND",
         block_size=prefill_block_size,
+        ssm_sizes=(0, 0),
     )
     handshake_payload = NixlHandshakePayload(
         compatibility_hash=remote_hash,
@@ -2391,7 +2422,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
         remote_block_size=decode_worker._block_size,  # shared state
         is_mla=decode_worker.use_mla,
         total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
-        attn_backend=backend,
+        attn_backends=[backend],
         tensor_shape=test_shape,
     )
 

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -74,6 +74,8 @@ def test_logical_to_kernel_block_ids_with_hma():
     # Simulate HMA scenario: logical block size = 32, kernel block size = 16
     # So each logical block maps to 2 kernel blocks eg [0]->[0,1]
     worker._physical_blocks_per_logical_kv_block = 2
+    # FA + SW groups (neither is MambaSpec, so both get expanded)
+    worker.kv_cache_config = make_kv_cache_config(block_size=16, hma_enabled=True)
 
     # Test conversion: FA + SW group
     logical_block_ids = [[0, 1, 2], [3, 4]]
@@ -201,3 +203,113 @@ def test_nixl_metadata_hma_block_ids_structure():
     assert len(req_meta.remote.block_ids) == 2
     assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17]
     assert list(req_meta.remote.block_ids[1]) == [18, 19, 20, 21]
+
+
+@pytest.mark.cpu_test
+def test_get_block_descs_ids_hybrid_ssm():
+    """Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM
+    when ratio=1 (no kernel block size mismatch)."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorWorker,
+    )
+
+    worker = object.__new__(NixlConnectorWorker)
+
+    num_blocks = 100
+    engine_id = "test-engine"
+    worker.num_regions = 2
+    worker.dst_num_blocks = {engine_id: num_blocks}
+    worker._has_mamba = True
+    worker._is_mamba_group = [False, True]
+    worker._physical_blocks_per_logical_kv_block = 1
+    # num_descs = num_regions * num_blocks (no blocks_first doubling)
+    worker.num_descs = 2 * num_blocks
+
+    fa_blocks = [3, 5]
+    ssm_blocks = [1, 2]
+    result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))
+
+    # FA group: stride=num_blocks=100, offset=0
+    #   region0: [3, 5],  region1: [103, 105]
+    # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1),
+    #   offset=num_descs=200
+    #   region0: [201, 202],  region1: [301, 302]
+    expected = [3, 5, 103, 105, 201, 202, 301, 302]
+    assert list(result) == expected, f"Expected {expected}, got {list(result)}"
+
+
+@pytest.mark.cpu_test
+def test_get_block_descs_ids_kernel_block_mismatch():
+    """Test _get_block_descs_ids uses different strides for FA (kernel blocks)
+    vs SSM (logical blocks) when ratio > 1."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorWorker,
+    )
+
+    worker = object.__new__(NixlConnectorWorker)
+
+    ratio = 4
+    logical_blocks = 100
+    num_blocks = logical_blocks * ratio  # 400 kernel blocks
+    engine_id = "test-engine"
+    worker.num_regions = 2
+    worker.dst_num_blocks = {engine_id: num_blocks}
+    worker._has_mamba = True
+    worker._is_mamba_group = [False, True]
+    worker._physical_blocks_per_logical_kv_block = ratio
+    worker.num_descs = 2 * num_blocks  # 800
+
+    fa_blocks = [3, 7]  # kernel-level block IDs
+    ssm_blocks = [1, 2]  # logical block IDs
+    result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))
+
+    # FA group: stride=num_blocks=400, offset=0
+    #   region0: [3, 7],  region1: [403, 407]
+    # SSM group: stride=logical_blocks=400//4=100, offset=num_descs=800
+    #   region0: [801, 802],  region1: [901, 902]
+    expected = [3, 7, 403, 407, 801, 802, 901, 902]
+    assert list(result) == expected, f"Expected {expected}, got {list(result)}"
+
+
+@pytest.mark.cpu_test
+def test_nixl_metadata_hybrid_ssm_block_ids():
+    """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM
+    groups with different block counts (kernel mismatch active)."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
+        NixlConnectorMetadata,
+    )
+
+    metadata = NixlConnectorMetadata()
+
+    # FA: 8 kernel blocks (2 logical * ratio=4), SSM: 2 logical blocks
+    fa_blocks = [0, 1, 2, 3, 4, 5, 6, 7]
+    ssm_blocks = [0, 1]
+
+    metadata.add_new_req_to_recv(
+        request_id="test-req-hybrid",
+        local_block_ids=(fa_blocks, ssm_blocks),
+        kv_transfer_params={
+            "remote_block_ids": ([10, 11, 12, 13, 14, 15, 16, 17], [20, 21]),
+            "remote_engine_id": "remote-engine",
+            "remote_request_id": "prefill-test-req-hybrid",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "tp_size": 1,
+        },
+    )
+
+    assert "test-req-hybrid" in metadata.reqs_to_recv
+    req_meta = metadata.reqs_to_recv["test-req-hybrid"]
+
+    # Verify local block IDs: different lengths per group
+    assert len(req_meta.local_block_ids) == 2
+    assert list(req_meta.local_block_ids[0]) == fa_blocks
+    assert list(req_meta.local_block_ids[1]) == ssm_blocks
+    assert len(req_meta.local_block_ids[0]) != len(req_meta.local_block_ids[1])
+
+    # Verify remote block IDs: same asymmetry preserved
+    assert req_meta.remote is not None
+    assert len(req_meta.remote.block_ids) == 2
+    assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17]
+    assert list(req_meta.remote.block_ids[1]) == [20, 21]
+    assert len(req_meta.remote.block_ids[0]) != len(req_meta.remote.block_ids[1])

tests/v1/spec_decode/test_extract_hidden_states.py

@@ -252,29 +252,22 @@ def test_propose():
     ]
 
     # Sampled token IDs from target model
-    sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
-
-    # Mock scheduler output
-    mock_scheduler_output = mock.MagicMock()
+    sampled_token_ids = torch.tensor(
+        [42, 60], dtype=torch.int32, device=device
+    ).unsqueeze(-1)
 
     # Call propose
-    with mock.patch(
-        "vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
-    ) as mock_has_kv:
-        mock_has_kv.return_value = False
-
-        draft_tokens, kv_connector_output = proposer.propose(
-            sampled_token_ids=sampled_token_ids,
-            target_hidden_states=target_hidden_states,
-            common_attn_metadata=common_attn_metadata,
-            scheduler_output=mock_scheduler_output,
-            slot_mappings=None,
-        )
+    draft_tokens = proposer.propose(
+        sampled_token_ids=sampled_token_ids,
+        target_hidden_states=target_hidden_states,
+        common_attn_metadata=common_attn_metadata,
+        slot_mappings=None,
+    )
 
     # Verify draft tokens match sampled tokens
     # Shape should be [batch_size, 1] for num_speculative_tokens=1
     assert draft_tokens.shape == (batch_size, 1)
-    assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
+    assert torch.equal(draft_tokens, sampled_token_ids)
 
     # Verify the model was called
     model_mock.assert_called_once()
@@ -326,21 +319,16 @@ def test_propose_different_layer_counts(num_hidden_layers):
         for _ in range(num_hidden_layers)
     ]
 
-    sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
-    mock_scheduler_output = mock.MagicMock()
+    sampled_token_ids = torch.tensor(
+        [42, 60], dtype=torch.int32, device=device
+    ).unsqueeze(-1)
 
-    with mock.patch(
-        "vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
-    ) as mock_has_kv:
-        mock_has_kv.return_value = False
-
-        draft_tokens, _ = proposer.propose(
-            sampled_token_ids=sampled_token_ids,
-            target_hidden_states=target_hidden_states,
-            common_attn_metadata=common_attn_metadata,
-            scheduler_output=mock_scheduler_output,
-            slot_mappings=None,
-        )
+    draft_tokens = proposer.propose(
+        sampled_token_ids=sampled_token_ids,
+        target_hidden_states=target_hidden_states,
+        common_attn_metadata=common_attn_metadata,
+        slot_mappings=None,
+    )
 
     assert draft_tokens.shape == (batch_size, 1)
-    assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
+    assert torch.equal(draft_tokens, sampled_token_ids)

tests/v1/structured_output/test_gptoss_structural_tags.py

@@ -1,173 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Unit tests for GPT-OSS structural tag support in reasoning (PR #25515)."""
-
-import json
-from unittest.mock import Mock
-
-import pytest
-
-from vllm.entrypoints.mcp.tool_server import ToolServer
-from vllm.reasoning.gptoss_reasoning_parser import (
-    GptOssReasoningParser,
-    from_builtin_tool_to_tag,
-    no_func_reaonsing_tag,
-    tag_with_builtin_funcs,
-)
-
-
-class TestGptOssReasoningParser:
-    """Test cases for GptOssReasoningParser structural tag functionality."""
-
-    @pytest.fixture
-    def mock_tokenizer(self):
-        """Create a mock tokenizer for testing."""
-        tokenizer = Mock()
-        tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
-        tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
-        return tokenizer
-
-    @pytest.fixture
-    def reasoning_parser(self, mock_tokenizer):
-        """Create a GptOssReasoningParser instance."""
-        return GptOssReasoningParser(mock_tokenizer)
-
-    @pytest.fixture
-    def mock_tool_server_empty(self):
-        """Create a mock ToolServer with no tools."""
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(return_value=False)
-        return tool_server
-
-    @pytest.fixture
-    def mock_tool_server_with_browser(self):
-        """Create a mock ToolServer with browser tool."""
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(side_effect=lambda tool: tool == "browser")
-        return tool_server
-
-    @pytest.fixture
-    def mock_tool_server_with_all_tools(self):
-        """Create a mock ToolServer with all builtin tools."""
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(
-            side_effect=lambda tool: tool in ["browser", "python", "container"]
-        )
-        return tool_server
-
-    def test_prepare_structured_tag_no_tool_server(self, reasoning_parser):
-        """Test prepare_structured_tag with no tool server."""
-        result = reasoning_parser.prepare_structured_tag(None, None)
-        expected = json.dumps(no_func_reaonsing_tag)
-
-        assert result == expected
-
-        # Verify the structure is correct
-        parsed = json.loads(result)
-        assert parsed["type"] == "structural_tag"
-        assert parsed["format"]["type"] == "triggered_tags"
-        assert len(parsed["format"]["tags"]) == 1
-        assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>"
-        assert parsed["format"]["triggers"] == ["<|channel|>analysis"]
-
-    def test_prepare_structured_tag_with_all_tools(
-        self, reasoning_parser, mock_tool_server_with_all_tools
-    ):
-        """Test prepare_structured_tag with all builtin tools."""
-        result = reasoning_parser.prepare_structured_tag(
-            None, mock_tool_server_with_all_tools
-        )
-        parsed = json.loads(result)
-
-        # Should have analysis tag + tags for all 3 tools (2 tags each)
-        assert len(parsed["format"]["tags"]) == 7  # 1 analysis + 6 tool tags
-
-        # Check all tool tags are present
-        tag_begins = [tag["begin"] for tag in parsed["format"]["tags"]]
-        for tool in ["browser", "python", "container"]:
-            assert f"<|channel|>commentary to={tool}" in tag_begins
-            assert f"<|channel|>analysis to={tool}" in tag_begins
-
-    def test_prepare_structured_tag_with_original_tag(self, reasoning_parser):
-        """Test prepare_structured_tag when original_tag is provided."""
-        original_tag = '{"custom": "tag"}'
-        result = reasoning_parser.prepare_structured_tag(original_tag, None)
-
-        # Should return the original tag unchanged
-        assert result == original_tag
-
-    def test_from_builtin_tool_to_tag(self):
-        """Test from_builtin_tool_to_tag function."""
-        tags = from_builtin_tool_to_tag("python")
-
-        assert len(tags) == 2
-        assert tags[0]["begin"] == "<|channel|>commentary to=python"
-        assert tags[0]["content"]["type"] == "any_text"
-        assert tags[0]["end"] == "<|end|>"
-
-        assert tags[1]["begin"] == "<|channel|>analysis to=python"
-        assert tags[1]["content"]["type"] == "any_text"
-        assert tags[1]["end"] == "<|end|>"
-
-    def test_tag_with_builtin_funcs(self):
-        """Test tag_with_builtin_funcs function."""
-        builtin_tools = ["browser", "python"]
-        result = tag_with_builtin_funcs(no_func_reaonsing_tag, builtin_tools)
-
-        assert result["type"] == "structural_tag"
-        # Should have original analysis tag + 2 tags per tool
-        assert len(result["format"]["tags"]) == 5  # 1 + 2*2
-
-        # Should have added commentary trigger
-        assert "<|channel|>commentary to=" in result["format"]["triggers"]
-        assert "<|channel|>analysis" in result["format"]["triggers"]
-
-    def test_tag_structure_invariants(self):
-        """Test that the basic tag structure follows expected format."""
-        # Test the base no_func_reaonsing_tag structure
-        assert no_func_reaonsing_tag["type"] == "structural_tag"
-        assert no_func_reaonsing_tag["format"]["type"] == "triggered_tags"
-        assert no_func_reaonsing_tag["format"]["stop_after_first"] is False
-
-        # Verify analysis tag structure
-        analysis_tag = no_func_reaonsing_tag["format"]["tags"][0]
-        assert analysis_tag["begin"] == "<|channel|>analysis<|message|>"
-        assert analysis_tag["content"]["type"] == "any_text"
-        assert analysis_tag["end"] == "<|end|>"
-
-    def test_json_serialization_valid(
-        self, reasoning_parser, mock_tool_server_with_all_tools
-    ):
-        """Test that all generated tags produce valid JSON."""
-        # Test with no tool server
-        result1 = reasoning_parser.prepare_structured_tag(None, None)
-        json.loads(result1)  # Should not raise
-
-        # Test with empty tool server
-        empty_server = Mock(spec=ToolServer)
-        empty_server.has_tool = Mock(return_value=False)
-        result2 = reasoning_parser.prepare_structured_tag(None, empty_server)
-        json.loads(result2)  # Should not raise
-
-        # Test with tools
-        result3 = reasoning_parser.prepare_structured_tag(
-            None, mock_tool_server_with_all_tools
-        )
-        json.loads(result3)  # Should not raise
-
-    @pytest.mark.parametrize("tool_name", ["browser", "python", "container"])
-    def test_single_tool_integration(self, reasoning_parser, tool_name):
-        """Test integration with individual tools."""
-        tool_server = Mock(spec=ToolServer)
-        tool_server.has_tool = Mock(side_effect=lambda tool: tool == tool_name)
-
-        result = reasoning_parser.prepare_structured_tag(None, tool_server)
-        parsed = json.loads(result)
-
-        # Should have 1 analysis + 2 tool-specific tags
-        assert len(parsed["format"]["tags"]) == 3
-
-        tag_begins = [tag["begin"] for tag in parsed["format"]["tags"]]
-        assert f"<|channel|>commentary to={tool_name}" in tag_begins
-        assert f"<|channel|>analysis to={tool_name}" in tag_begins

tools/pre_commit/check_torch_cuda.py

@@ -8,7 +8,7 @@ import regex as re
 # Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
 # --------------------------------------------------------------------------- #
 _TORCH_CUDA_PATTERNS = [
-    r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|set_device|device\()\b",
+    r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|set_device|device\()\b",
     r"\bwith\storch\.cuda\.device\b",
 ]
 

vllm/_xpu_ops.py

@@ -8,6 +8,7 @@ from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func
 
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
 
 logger = init_logger(__name__)
 
@@ -54,6 +55,37 @@ if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
         return torch.empty((M, N), dtype=input.dtype, device=input.device)
 
 
+def _xpu_ops_deepseek_scaling_rope_impl(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor | None,
+    offsets: torch.Tensor | None,
+    cos_sin_cache: torch.Tensor | None,
+    rotary_dim: int,
+    is_neox_style: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert key is not None
+    return torch.ops._xpu_C.deepseek_scaling_rope(
+        positions, query, key, offsets, cos_sin_cache, rotary_dim, is_neox_style
+    )
+
+
+def _xpu_ops_deepseek_scaling_rope_fake(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor | None,
+    offsets: torch.Tensor | None,
+    cos_sin_cache: torch.Tensor | None,
+    rotary_dim: int,
+    is_neox_style: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return query, key
+
+
+# Global flag to ensure ops are registered only once
+_OPS_REGISTERED = False
+
+
 class xpu_ops:
     @staticmethod
     def flash_attn_varlen_func(
@@ -402,3 +434,21 @@ class xpu_ops:
         raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
             topk_indices
         )
+
+    @staticmethod
+    def register_ops_once() -> None:
+        global _OPS_REGISTERED
+        if not _OPS_REGISTERED:
+            # register all the custom ops here
+            direct_register_custom_op(
+                op_name="xpu_ops_deepseek_scaling_rope",
+                op_func=_xpu_ops_deepseek_scaling_rope_impl,
+                mutates_args=[],
+                fake_impl=_xpu_ops_deepseek_scaling_rope_fake,
+                dispatch_key=current_platform.dispatch_key,
+            )
+
+            _OPS_REGISTERED = True
+
+
+xpu_ops.register_ops_once()

vllm/compilation/cuda_graph.py

@@ -16,7 +16,11 @@ from vllm.compilation.counter import compilation_counter
 from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
 from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
-from vllm.forward_context import BatchDescriptor, get_forward_context
+from vllm.forward_context import (
+    BatchDescriptor,
+    get_forward_context,
+    is_forward_context_available,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.offloader.base import get_offloader
 from vllm.platforms import current_platform
@@ -224,6 +228,12 @@ class CUDAGraphWrapper:
         self.concrete_cudagraph_entries.clear()
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
+        if not is_forward_context_available():
+            # No forward context means we are outside the normal
+            # inference path (e.g. a vision encoder forward pass).
+            # Just run the underlying function without cudagraphs.
+            return self.runnable(*args, **kwargs)
+
         forward_context = get_forward_context()
         batch_descriptor = forward_context.batch_descriptor
         cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode

vllm/compilation/wrapper.py

@@ -10,7 +10,6 @@ from types import CodeType
 from typing import Any, ParamSpec, TypeVar
 
 import torch
-import torch._C._dynamo.guards
 
 import vllm.envs as envs
 from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config
@@ -24,65 +23,23 @@ R = TypeVar("R")
 P = ParamSpec("P")
 
 
-def _noop_add_global_state_guard(
-    self: torch._C._dynamo.guards.GuardManager, *args: Any, **kwargs: Any
-) -> None:
-    """No-op to skip the GLOBAL_STATE guard entirely"""
-    pass
-
-
-def _noop_add_torch_function_mode_stack_guard(
-    self: torch._C._dynamo.guards.GuardManager, *args: Any, **kwargs: Any
-) -> None:
-    """No-op to skip the TORCH_FUNCTION_MODE_STACK guard entirely"""
-    pass
-
-
 @contextmanager
 def _compilation_context() -> Generator[None, None, None]:
-    """Context manager for compilation settings and patches.
+    """Context manager for compilation settings.
 
-    This manager:
-    1. Sets higher dynamo cache limits for compilation. (Needed for
-        qwen2_5_vl see test_qwen2_5_vl_evs_functionality).
-        Generally a recompilation can happen whenever we use a new
-        backend instance in torch.compile.
-    2. Patches out add_global_state_guard to skip GLOBAL_STATE guards
-    3. Patches out add_torch_function_mode_stack_guard to skip
-        TORCH_FUNCTION_MODE_STACK guards.
-    4. Restores everything when compilation completes
+    This manager sets higher dynamo cache limits for compilation.
+    (Needed for qwen2_5_vl see test_qwen2_5_vl_evs_functionality).
+    Generally a recompilation can happen whenever we use a new
+    backend instance in torch.compile.
     """
-    # Save original values
-    original_global_state_guard = (
-        torch._C._dynamo.guards.GuardManager.add_global_state_guard
-    )
-    original_torch_function_mode_stack_guard = (
-        torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard
-    )
     original_cache_size = torch._dynamo.config.cache_size_limit
     original_accumulated_cache = torch._dynamo.config.accumulated_cache_size_limit
 
     try:
-        # Set higher cache limits for compilation
         torch._dynamo.config.cache_size_limit = 2048
         torch._dynamo.config.accumulated_cache_size_limit = 8192
-
-        # Patch guard manager
-        torch._C._dynamo.guards.GuardManager.add_global_state_guard = (
-            _noop_add_global_state_guard
-        )
-        torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard = (
-            _noop_add_torch_function_mode_stack_guard
-        )
         yield
     finally:
-        # Restore original values
-        torch._C._dynamo.guards.GuardManager.add_global_state_guard = (
-            original_global_state_guard
-        )
-        torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard = (
-            original_torch_function_mode_stack_guard
-        )
         torch._dynamo.config.cache_size_limit = original_cache_size
         torch._dynamo.config.accumulated_cache_size_limit = original_accumulated_cache
 
@@ -155,7 +112,7 @@ class TorchCompileWithNoGuardsWrapper:
                     entry.guard_type == "SHAPE_ENV" for entry in x
                 ]
             else:
-                options["guard_filter_fn"] = lambda x: [False for _ in x]
+                options["guard_filter_fn"] = torch.compiler.skip_all_guards_unsafe
 
         compiled_ptr: Any = self.forward
         # Validate that unbacked dynamic shapes require VLLM_USE_BYTECODE_HOOK=False

vllm/config/cache.py

@@ -13,6 +13,7 @@ logger = init_logger(__name__)
 
 CacheDType = Literal[
     "auto",
+    "float16",
     "bfloat16",
     "fp8",
     "fp8_e4m3",

vllm/config/load.py

@@ -62,6 +62,9 @@ class LoadConfig:
       This is recommended for models on network filesystems (e.g., Lustre, NFS)
       as it avoids inefficient random reads, significantly speeding up model
       initialization. However, it uses more CPU RAM.
+    - "prefetch": Checkpoint files are read into the OS page cache before
+      workers load them, speeding up the model loading phase. Useful on
+      network or high-latency storage.
     - "torchao": Weights are loaded in upfront and then reconstructed
       into torchao tensor subclasses. This is used when the checkpoint
       was quantized using torchao and saved using safetensors.

vllm/config/model.py

@@ -2021,6 +2021,15 @@ def _get_and_verify_max_len(
 
                 if rope_type == "yarn":
                     derived_max_model_len = rp["original_max_position_embeddings"]
+        if scaling_factor is None:
+            # Fallback the factor to 1.0 if a user assigned `null`
+            logger.warning_once(
+                "The model's RoPE configuration has a null scaling "
+                "factor which is unexpected. This likely indicates a bug "
+                "in the model's HuggingFace config.json. Please notify the "
+                "model vendor. Falling back the value to 1.0. "
+            )
+            scaling_factor = 1.0
         # Do this outside loop since all layer types should have the same scaling
         derived_max_model_len *= scaling_factor
 

vllm/config/parallel.py

@@ -45,7 +45,9 @@ All2AllBackend = Literal[
     "mori",
     "nixl_ep",
     "allgather_reducescatter",
-    "flashinfer_all2allv",
+    "flashinfer_all2allv",  # temporary alias for flashinfer_nvlink_two_sided
+    "flashinfer_nvlink_two_sided",
+    "flashinfer_nvlink_one_sided",
 ]
 
 
@@ -158,7 +160,8 @@ class ParallelConfig:
     - "deepep_low_latency": Use deepep low-latency kernels\n
     - "mori": Use mori kernels\n
     - "nixl_ep": Use nixl-ep kernels\n
-    - "flashinfer_all2allv": Use flashinfer alltoallv kernels for mnnvl"""
+    - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl
+    - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels"""
 
     max_parallel_loading_workers: int | None = None
     """Maximum number of parallel loading workers when loading model

vllm/config/vllm.py

@@ -1574,8 +1574,9 @@ class VllmConfig:
                 "runai_streamer_sharded",
             ):
                 raise ValueError(
-                    f"To load a model from S3, 'load_format' "
-                    f"must be 'runai_streamer' or 'runai_streamer_sharded', "
+                    f"To load a model from object storage (S3/GCS/Azure), "
+                    f"'load_format' must be 'runai_streamer' or "
+                    f"'runai_streamer_sharded', "
                     f"but got '{self.load_config.load_format}'. "
                     f"Model: {self.model_config.model}"
                 )

vllm/distributed/device_communicators/all2all.py

@@ -4,23 +4,36 @@ import threading
 from typing import Any
 
 import torch
+import torch.distributed as dist
 
 import vllm.envs as envs
 from vllm.distributed import get_dp_group, get_ep_group
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
-from vllm.utils.flashinfer import has_flashinfer_all2all
+from vllm.utils.flashinfer import (
+    has_flashinfer_nvlink_one_sided,
+    has_flashinfer_nvlink_two_sided,
+)
 from vllm.utils.import_utils import has_deep_ep, has_mori
 
 from .base_device_communicator import All2AllManagerBase, Cache
 
-if has_flashinfer_all2all():
+if has_flashinfer_nvlink_two_sided():
     from flashinfer.comm import Mapping  # type: ignore[import-not-found]
     from flashinfer.comm.mnnvl import MnnvlConfig  # type: ignore[import-not-found]
     from flashinfer.comm.trtllm_alltoall import (
         MnnvlMoe,  # type: ignore[import-not-found]
     )
 
+if has_flashinfer_nvlink_one_sided():
+    from flashinfer.comm import Mapping  # type: ignore[import-not-found]
+    from flashinfer.comm.mnnvl import MnnvlConfig  # type: ignore[import-not-found]
+    from flashinfer.comm.trtllm_moe_alltoall import (
+        MoeAlltoAll,  # type: ignore[import-not-found]
+        moe_a2a_get_workspace_size_per_rank,
+    )
+
+
 logger = init_logger(__name__)
 
 
@@ -529,9 +542,9 @@ class NixlEPAll2AllManager(All2AllManagerBase):
         return 0
 
 
-class FlashInferAllToAllManager(All2AllManagerBase):
+class FlashInferNVLinkTwoSidedManager(All2AllManagerBase):
     """
-    All2All communication based on flashinfer kernels.
+    All2All communication based on flashinfer all2allv/two-sided NVLink kernels.
     """
 
     # This type lint could be removed after all of the work in
@@ -540,7 +553,7 @@ class FlashInferAllToAllManager(All2AllManagerBase):
     world_size: int
 
     def __init__(self, cpu_group, tcp_store_group=None):
-        assert has_flashinfer_all2all(), (
+        assert has_flashinfer_nvlink_two_sided(), (
             "flashinfer all2all module not found. Please install/check flashinfer"
         )  # noqa
         super().__init__(cpu_group, tcp_store_group)
@@ -597,7 +610,7 @@ class FlashInferAllToAllManager(All2AllManagerBase):
 
     def ensure_alltoall_workspace_initialized(self):
         """Ensure workspace is initialized"""
-        if not has_flashinfer_all2all():
+        if not has_flashinfer_nvlink_two_sided():
             return False
 
         if self.world_size <= 1:
@@ -633,6 +646,119 @@ class FlashInferAllToAllManager(All2AllManagerBase):
                 self.initialized = False
 
 
+class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
+    """
+    All2All communication based on FlashInfer's MoeAlltoAll/One-sided NVLink kernel.
+    This is a newer kernel from trtllm that should perform better than the kernel
+    used by flashinfer_nvlink_two_sided.
+    """
+
+    rank: int
+    world_size: int
+
+    def __init__(self, cpu_group):
+        assert has_flashinfer_nvlink_one_sided(), (
+            "flashinfer trtllm_moe_alltoall module not found. "
+            "Please install/check flashinfer"
+        )
+        super().__init__(cpu_group)
+        logger.debug(
+            "Initialize FlashInfer One-sided NVLink rank=%d, world size=%d",
+            self.rank,
+            self.world_size,
+        )
+        self.initialized = False
+        self.moe_alltoall: MoeAlltoAll | None = None
+        self.mapping = None
+
+    def initialize(
+        self,
+        max_num_tokens: int,
+        top_k: int,
+        num_experts: int,
+        hidden_size: int,
+    ):
+        """Initialize the MoeAlltoAll workspace."""
+        if self.initialized:
+            return
+
+        self.cleanup()
+        gpus_per_node = torch.accelerator.device_count()
+        logger.debug(
+            "Making One-sided NVLink mapping: rank=%d, world size=%d",
+            self.rank,
+            self.world_size,
+        )
+        self.mapping = Mapping(
+            self.world_size,
+            self.rank,
+            gpus_per_node,
+            tp_size=self.world_size,
+            moe_ep_size=self.world_size,
+        )
+
+        from vllm.distributed.device_communicators.mnnvl_compat import (
+            CustomCommunicator,
+        )
+
+        dp_config = MnnvlConfig(
+            comm_backend=CustomCommunicator(get_dp_group().cpu_group),
+        )
+        total_dispatch_payload_size_per_token = (
+            hidden_size // 2  # nvfp4 hidden states
+            + hidden_size // 16  # fp8 scaling factors
+            + top_k * 4  # int32 topks ids
+            + top_k * 4  # float32 topk weights
+        )
+        combine_payload_size_per_token = hidden_size * 2  # bf16 hidden states
+        self.workspace_size = moe_a2a_get_workspace_size_per_rank(
+            ep_size=self.world_size,
+            max_num_tokens=max_num_tokens,
+            total_dispatch_payload_size_per_token=total_dispatch_payload_size_per_token,
+            combine_payload_size_per_token=combine_payload_size_per_token,
+        )
+
+        self.moe_alltoall = MoeAlltoAll(
+            mapping=self.mapping,
+            max_num_tokens=max_num_tokens,
+            top_k=top_k,
+            num_experts=num_experts,
+            workspace_size_per_rank=self.workspace_size,
+            mnnvl_config=dp_config,
+        )
+
+        self.gpus_per_node = gpus_per_node
+        self.max_num_tokens = max_num_tokens
+        self.top_k = top_k
+        self.num_experts = num_experts
+        self.hidden_size = hidden_size
+        self.initialized = True
+
+        logger.info(
+            "FlashInfer One-sided NVLink initialized for rank %s, size %s",
+            self.rank,
+            self.world_size,
+        )
+        dist.barrier()
+
+    def get_handle(self, kwargs):
+        return self
+
+    def cleanup(self):
+        """Clean up resources."""
+        if self.initialized and self.moe_alltoall is not None:
+            try:
+                del self.moe_alltoall
+            except Exception as e:
+                logger.warning(
+                    "Failed to cleanup FlashInfer One-sided NVLink workspace: %s", e
+                )
+            finally:
+                self.moe_alltoall = None
+                self.mapping = None
+                self.initialized = False
+
+
 class MoriAll2AllManager(All2AllManagerBase):
     def __init__(self, cpu_group):
         assert has_mori(), (

vllm/distributed/device_communicators/cuda_communicator.py

@@ -149,12 +149,25 @@ class CudaCommunicator(DeviceCommunicatorBase):
                 self.all2all_manager = NixlEPAll2AllManager(
                     self.cpu_group, tcp_store_group
                 )
-            elif self.all2all_backend == "flashinfer_all2allv":
-                from .all2all import FlashInferAllToAllManager
+            elif (
+                self.all2all_backend == "flashinfer_all2allv"
+                or self.all2all_backend == "flashinfer_nvlink_two_sided"
+            ):
+                if self.all2all_backend == "flashinfer_all2allv":
+                    logger.warning_once(
+                        "'flashinfer_all2allv' is deprecated and has been renamed to"
+                        "'flashinfer_nvlink_two_sided'. It will be removed in a future"
+                        "release."
+                    )
+                from .all2all import FlashInferNVLinkTwoSidedManager
 
-                self.all2all_manager = FlashInferAllToAllManager(
+                self.all2all_manager = FlashInferNVLinkTwoSidedManager(
                     self.cpu_group, tcp_store_group
                 )
+            elif self.all2all_backend == "flashinfer_nvlink_one_sided":
+                from .all2all import FlashInferNVLinkOneSidedManager
+
+                self.all2all_manager = FlashInferNVLinkOneSidedManager(self.cpu_group)
             else:
                 raise ValueError(f"Unknown all2all backend: {self.all2all_backend}")
 

vllm/distributed/device_communicators/mnnvl_compat.py

@@ -5,9 +5,9 @@ from typing import Any
 import torch.distributed as dist
 from flashinfer.comm.mnnvl import CommBackend as CommBackend
 
-from vllm.utils.flashinfer import has_flashinfer_all2all
+from vllm.utils.flashinfer import has_flashinfer_nvlink_two_sided
 
-assert has_flashinfer_all2all(), "Flashinfer alltoallv module cannot be found"
+assert has_flashinfer_nvlink_two_sided(), "Flashinfer alltoallv module cannot be found"
 
 
 class CustomCommunicator(CommBackend):
@@ -25,14 +25,14 @@ class CustomCommunicator(CommBackend):
         dist.all_gather_object(gathered, data, group=self._group)
         return gathered
 
-    # NOTE(rob): CommBackend is an abstract class, and bcast/barrier
-    # are unimplemented on vLLM side. If we need to utilize these
-    # methods in the future, can create a concrete implementation.
     def bcast(self, data: Any, root: int) -> Any:
-        raise NotImplementedError
+        obj_list = [data]
+        # broadcast_object_list mutates obj_list in-place
+        dist.broadcast_object_list(obj_list, src=root, group=self._group)
+        return obj_list[0]
 
     def barrier(self) -> None:
-        raise NotImplementedError
+        dist.barrier(group=self._group)
 
     def Split(self, color: int, key: int) -> "CustomCommunicator":
         return self

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -16,10 +16,12 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.platforms import current_platform
 from vllm.v1.attention.backend import AttentionBackend
+from vllm.v1.kv_cache_interface import MambaSpec
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 
 if TYPE_CHECKING:
     from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
+    from vllm.v1.kv_cache_interface import KVCacheSpec
 
 logger = init_logger(__name__)
 
@@ -328,22 +330,26 @@ class TpKVTopology:
     remote_tp_size: dict[EngineId, int]
     is_mla: bool
     total_num_kv_heads: int
-    attn_backend: type[AttentionBackend]
+    attn_backends: list[type[AttentionBackend]]
     engine_id: EngineId
     remote_block_size: dict[EngineId, int]
     tensor_shape: torch.Size | None = None
+    is_mamba: bool = False
 
     def __post_init__(self):
         # Figure out whether the first dimension of the cache is K/V
         # or num_blocks. This is used to register the memory regions correctly.
-        _MOCK_BLOCK_SIZE = 16
-        kv_cache_shape = self.attn_backend.get_kv_cache_shape(
-            num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1
-        )
-        logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
+        attn_backend = self.attn_backends[0]
+        if not self.is_mamba:
+            _MOCK_BLOCK_SIZE = 16
+            kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape(
+                num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1
+            )
+            logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
         # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
         # we just mock num_blocks to 1 for the dimension check below.
-        self._is_kv_layout_blocks_first = (
+        # Hybrid SSM models assume a single blocks_first layout
+        self._is_kv_layout_blocks_first = self.is_mamba or (
             len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
         )
 
@@ -360,7 +366,7 @@ class TpKVTopology:
             _MOCK_NUM_LAYERS = 80
             kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape
             try:
-                kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order(
+                kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
                     include_num_layers_dimension=self._cross_layers_blocks
                 )
             except (AttributeError, NotImplementedError):
@@ -483,6 +489,30 @@ class TpKVTopology:
         remote_tp_size = self.remote_tp_size[remote_engine_id]
         return self.get_target_remote_ranks(remote_tp_size)
 
+    def get_transfer_cache_regions(
+        self, cache: torch.Tensor, layer_spec: "KVCacheSpec"
+    ) -> list[torch.Tensor] | torch.Tensor:
+        """Return the cache tensor(s) to register as NIXL memory regions,
+        also accounting for hybrid SSM models specificities.
+        """
+        if isinstance(layer_spec, MambaSpec):
+            # Register the whole kv cache shared tensor, including SSM/Conv. This is
+            # similar to FI with the difference that SSM/Conv have different sizes
+            conv, ssm = cache
+            return [conv]
+
+        # Check may be hacky but it's matching `_update_hybrid_attention_mamba_layout`.
+        if self.is_mamba and cache.shape[0] == 2:
+            # When MAMBA is present, all backends are blocks first, so that blocks
+            # can be shared between attention layers and mamba layers. Runner
+            # `_update_hybrid_attention_mamba_layout` already adjusted strides
+            # for FlashAttn-like backends so its num_blocks first.
+            # Swap [2<>num_blocks] dims to get required layout for hybrid SSM.
+            cache = cache.transpose(0, 1)
+
+        # Regular case: backends like FA register K/V in separate regions
+        return cache if self.split_k_and_v else [cache]
+
 
 def get_current_attn_backends(
     vllm_config: VllmConfig, layer_names: list[str] | None = None

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

@@ -286,7 +286,9 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1):
             cached_req = self._active_requests[req_id]
             req_block_ids = self._req_blocks[req_id]
 
-            assert new_block_ids is not None
+            if new_block_ids is None:
+                continue
+
             block_ids = new_block_ids[0]
 
             req_block_ids.extend(block_ids)

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -564,7 +564,7 @@ class MooncakeConnectorWorker:
             remote_block_size=self._block_size,  # shared state
             is_mla=self.use_mla,
             total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
-            attn_backend=backend,
+            attn_backends=[backend],
         )
 
         self.async_zmq_ctx = zmq.asyncio.Context()

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py

@@ -39,11 +39,13 @@ logger = init_logger(__name__)
 Transfer = tuple[int, float]
 EngineId = str
 ReqId = str
+TransferId = str
 
 
 @dataclass
 class WriteTask:
-    request_id: str
+    request_id: ReqId
+    transfer_id: TransferId
     dst_engine_id: str
     local_block_ids: list[int]
     remote_block_ids_hint: list[int] | None
@@ -59,7 +61,8 @@ class WriteTask:
 class LayerTransferPlan:
     """Plan for transferring a single layer."""
 
-    request_id: str
+    request_id: ReqId
+    transfer_id: TransferId
     layer_name: str
     sess_idx: int
     transfer_local_offsets: list[int]
@@ -234,6 +237,7 @@ class MoRIIOConstants:
     POP_DONE_RECV = b"pop_done_recv"
     OVER = b"OVER"
     COMPLETION_PREFIX = "cmpl"
+    TRANSFER_PREFIX = "tx"
 
     PING_INTERVAL = 5
     MAX_PING_RETRIES = 100
@@ -247,6 +251,7 @@ class MoRIIOConstants:
 class ReqMeta:
     """Metadata for a single request."""
 
+    transfer_id: TransferId
     local_block_ids: list[int]
     remote_block_ids: list[int]
     remote_host: str
@@ -263,21 +268,15 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
         self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
         self.reqs_to_save: dict[ReqId, ReqMeta] = {}
         self.reqs_to_send: dict[ReqId, float] = {}
+        self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}
 
     def __repr__(self):
-        return_str = ""
-        for req_id, req_meta in self.reqs_to_recv.items():
-            return_str += (
-                f"{req_id = },{req_meta.local_block_ids = },"
-                f"{req_meta.remote_host = },{req_meta.remote_port = }"
-                f"{req_meta.remote_engine_id = },{req_meta.tp_size = }"
-            )
-        return_str = f"MoRIIOConnectorMetadata:reqs_to_recv:{return_str},"
-
-        for req_id, expiry in self.reqs_to_send.items():
-            return_str += f"{req_id = },{expiry = }"
-        return_str = f"MoRIIOConnectorMetadata:reqs_to_send:{return_str},"
-        return return_str
+        return (
+            f"MoRIIOConnectorMetadata: reqs_to_recv={self.reqs_to_recv}, "
+            f"reqs_to_save={self.reqs_to_save}, "
+            f"reqs_to_send={self.reqs_to_send}, "
+            f"transfer_id_to_request_id={self.transfer_id_to_request_id}"
+        )
 
     def add_new_req(
         self,
@@ -286,7 +285,9 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata):
         kv_transfer_params: dict[str, Any],
         write_mode=False,
     ):
+        transfer_id = kv_transfer_params["transfer_id"]
         _req = ReqMeta(
+            transfer_id=transfer_id,
             local_block_ids=local_block_ids,
             remote_block_ids=kv_transfer_params["remote_block_ids"],
             remote_engine_id=kv_transfer_params["remote_engine_id"],

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

@@ -32,6 +32,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
     MoRIIOMode,
     ReqId,
     ReqMeta,
+    TransferId,
     WriteTask,
     get_moriio_mode,
     get_port_offset,
@@ -277,6 +278,30 @@ class MoRIIOConnectorScheduler:
         # Reqs to send and their expiration time
         self._reqs_need_send: dict[ReqId, float] = {}
         self.paths: dict[str, zmq.Socket] = {}
+        self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}
+        self.request_id_to_transfer_id: dict[ReqId, TransferId] = {}
+
+    def map_request_id(self, request_id: ReqId, transfer_id: TransferId):
+        self.transfer_id_to_request_id[transfer_id] = request_id
+        self.request_id_to_transfer_id[request_id] = transfer_id
+
+    def unmap_request_id(self, request_id: ReqId):
+        if request_id in self.request_id_to_transfer_id:
+            transfer_id = self.request_id_to_transfer_id[request_id]
+            del self.request_id_to_transfer_id[request_id]
+            if transfer_id in self.transfer_id_to_request_id:
+                del self.transfer_id_to_request_id[transfer_id]
+            else:
+                logger.warning(
+                    "transfer id not in transfer_id_to_request_id lookup"
+                    "table. there is likely a bug!"
+                )
+        else:
+            logger.warning(
+                "Could not find %s  in transfer_id_to_request_id"
+                "lookup table.  This could lead to a possible hang.",
+                request_id,
+            )
 
     def get_num_new_matched_tokens(
         self,
@@ -309,7 +334,12 @@ class MoRIIOConnectorScheduler:
         return len(token_ids) - 1 - num_computed_tokens, False
 
     def send_notify_block(
-        self, req_id: str, block_notify_list: list[int], host=None, port=None
+        self,
+        req_id: ReqId,
+        transfer_id: TransferId,
+        block_notify_list: list[int],
+        host=None,
+        port=None,
     ):
         path = make_zmq_path("tcp", host, port)
         if path not in self.paths:
@@ -321,6 +351,7 @@ class MoRIIOConnectorScheduler:
 
         data = {
             "req_id": req_id,
+            "transfer_id": transfer_id,
             "block_notify_list": block_notify_list or [],
             "decode_rank": self.dp_rank,
             "type": "remote_blocks",
@@ -338,6 +369,9 @@ class MoRIIOConnectorScheduler:
         params = request.kv_transfer_params
         if not params:
             return
+        transfer_id = params["transfer_id"]
+        request_id = request.request_id
+        self.map_request_id(request_id, transfer_id)
         if params.get("do_remote_decode"):
             local_block_ids = blocks.get_block_ids()[0]
             self._reqs_need_save[request.request_id] = (request, local_block_ids)
@@ -386,6 +420,7 @@ class MoRIIOConnectorScheduler:
 
                     self.send_notify_block(
                         req_id=request.request_id,
+                        transfer_id=request.kv_transfer_params["transfer_id"],
                         block_notify_list=blocks.get_block_ids()[0],
                         host=params.get("remote_host"),
                         port=target_port,
@@ -400,6 +435,7 @@ class MoRIIOConnectorScheduler:
         scheduler_output: SchedulerOutput,
     ) -> KVConnectorMetadata:
         meta = MoRIIOConnectorMetadata()
+        meta.transfer_id_to_request_id = self.transfer_id_to_request_id
 
         if self.mode == MoRIIOMode.WRITE:
             # when async_load_kv finished,
@@ -506,6 +542,9 @@ class MoRIIOConnectorScheduler:
         should be freed now or will be sent asynchronously and freed later.
         """
 
+        request_id = request.request_id
+        self.unmap_request_id(request_id)
+
         params = request.kv_transfer_params
         logger.debug(
             "MoriioConnector request_finished, request_status=%s, "
@@ -728,6 +767,7 @@ class MoRIIOConnectorWorker:
             self.cache_config.cache_dtype,
             use_mla=self.use_mla,
         )
+        self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}
 
         # TODO: consider the integration of flashinfer or other backends.
         self.backend_name = backend.get_name()
@@ -735,7 +775,8 @@ class MoRIIOConnectorWorker:
 
     def schedule_write_blocks(
         self,
-        request_id: str,
+        request_id: ReqId,
+        transfer_id: TransferId,
         dst_engine_id: str,
         local_block_ids: list[int],
         remote_block_ids: list[int] | None,
@@ -748,6 +789,7 @@ class MoRIIOConnectorWorker:
 
         Args:
             request_id: Unique identifier for the request
+            transfer_id: Unique identifier for the transfer
             dst_engine_id: Destination engine ID
             local_block_ids: Local block IDs to transfer
             remote_block_ids: Hint for remote block IDs
@@ -768,6 +810,7 @@ class MoRIIOConnectorWorker:
 
         task = WriteTask(
             request_id=request_id,
+            transfer_id=transfer_id,
             dst_engine_id=dst_engine_id,
             local_block_ids=local_block_ids,
             remote_block_ids_hint=remote_block_ids,
@@ -1010,7 +1053,7 @@ class MoRIIOConnectorWorker:
         return {remote_agent_name}
 
     def _background_moriio_handshake(
-        self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta
+        self, req_id: ReqId, remote_engine_id: EngineId, meta: ReqMeta
     ):
         # Do MoRIIO handshake in background and add to _ready_requests when done.
         fut = None
@@ -1189,6 +1232,13 @@ class MoRIIOConnectorWorker:
             else:
                 done_recving = self._pop_done_transfers()
 
+        done_recving = {
+            self.transfer_id_to_request_id[id]
+            for id in filter(
+                lambda id: id in self.transfer_id_to_request_id, done_recving
+            )
+        }
+
         return done_sending, done_recving
 
     def _pop_done_transfers(self) -> set[str]:
@@ -1269,6 +1319,7 @@ class MoRIIOConnectorWorker:
         Start loading by triggering non-blocking moriio_xfer.
         We check for these trnxs to complete in each step().
         """
+        self.transfer_id_to_request_id = metadata.transfer_id_to_request_id
         if self.is_producer:
             self.moriio_wrapper.async_wait_reqid()
             return
@@ -1332,9 +1383,10 @@ class MoRIIOConnectorWorker:
             remote_notify_port=meta.remote_notify_port,
         )
 
-    def _write_blocks_for_req(self, req_id: str, meta: ReqMeta, layer_name, kv_layer):
+    def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer):
         self.schedule_write_blocks(
             request_id=req_id,
+            transfer_id=meta.transfer_id,
             dst_engine_id=meta.remote_engine_id,
             local_block_ids=meta.local_block_ids,
             remote_block_ids=meta.remote_block_ids,

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py

@@ -29,6 +29,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
     MoRIIOError,
     RemoteAllocInfo,
     TransferError,
+    TransferId,
     WriteTask,
     get_port_offset,
     get_role,
@@ -162,14 +163,14 @@ class MoRIIOWriter:
             True if remote blocks are ready
         """
         return (
-            task.request_id in self.worker.moriio_wrapper.done_remote_allocate_req_dict
+            task.transfer_id in self.worker.moriio_wrapper.done_remote_allocate_req_dict
         )
 
-    def _get_remote_alloc_info(self, request_id: str) -> RemoteAllocInfo:
+    def _get_remote_alloc_info(self, transfer_id: str) -> RemoteAllocInfo:
         """Get remote allocation info for a request.
 
         Args:
-            request_id: The request ID
+            transfer_id:TransferId The request ID
 
         Returns:
             Remote allocation information
@@ -178,10 +179,10 @@ class MoRIIOWriter:
             KeyError: If allocation info is missing
         """
         try:
-            return self.worker.moriio_wrapper.done_remote_allocate_req_dict[request_id]
+            return self.worker.moriio_wrapper.done_remote_allocate_req_dict[transfer_id]
         except KeyError as e:
             raise KeyError(
-                f"Remote allocation info missing for request {request_id}"
+                f"Remote allocation info missing for transfer {transfer_id}"
             ) from e
 
     def _execute_write_task(self, task: WriteTask) -> None:
@@ -192,10 +193,14 @@ class MoRIIOWriter:
 
         """
         # Get remote allocation info
-        request_info = self._get_remote_alloc_info(task.request_id)
+        request_info = self._get_remote_alloc_info(task.transfer_id)
 
         if request_info.block_ids is None:
-            logger.debug("Request %s remote block IDs not ready", task.request_id)
+            logger.debug(
+                "Request remote block IDs not ready:request_id = %s, transfer_id = %s",
+                task.request_id,
+                task.transfer_id,
+            )
             return
 
         # Wait for CUDA event
@@ -257,6 +262,7 @@ class MoRIIOWriter:
 
         return LayerTransferPlan(
             request_id=task.request_id,
+            transfer_id=task.transfer_id,
             layer_name=task.layer_name,
             sess_idx=sess_idx,
             transfer_local_offsets=local_off,
@@ -312,17 +318,18 @@ class MoRIIOWriter:
 
             # Send completion notification
             self.worker.moriio_wrapper.send_notify(
-                task.request_id, task.remote_ip, remote_port
+                task.transfer_id, task.remote_ip, remote_port
             )
             # mark request as done, then we can free the blocks
             with self.worker.moriio_wrapper.lock:
                 self.worker.moriio_wrapper.done_req_ids.append(task.request_id)
             del self.worker.moriio_wrapper.done_remote_allocate_req_dict[
-                task.request_id
+                task.transfer_id
             ]
             logger.debug(
-                "Completed transfer for request %s, notified port %d",
+                "Completed transfer for (request, transfer) %s, %s, notified port %d",
                 task.request_id,
+                task.transfer_id,
                 remote_port,
             )
 
@@ -355,7 +362,7 @@ class MoRIIOWrapper:
         self.notify_port: int | None = None
         self.lock = threading.Lock()
         self.done_req_ids: list[str] = []
-        self.done_remote_allocate_req_dict: dict[str, RemoteAllocInfo] = {}
+        self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {}
         self.done_write_cache_req_ids: list[str] = []
         self.notify_thread: threading.Thread | None = None
         self.sessions: list[IOEngine.Session] = []
@@ -525,7 +532,7 @@ class MoRIIOWrapper:
 
         try:
             msg_str = msg.decode("UTF-8")
-            if msg_str.startswith(MoRIIOConstants.COMPLETION_PREFIX):
+            if msg_str.startswith(MoRIIOConstants.TRANSFER_PREFIX):
                 self._handle_completion_message(msg_str)
                 handled = True
         except UnicodeDecodeError:
@@ -535,7 +542,7 @@ class MoRIIOWrapper:
 
     def _handle_structured_message(self, data: dict):
         assert get_role() == ROLE.PRODUCER, "Only prefill can get block messages"
-        req_id = data["req_id"]
+        transfer_id = data["transfer_id"]
         block_notify_list = data.get("block_notify_list", [])
         decode_dp_rank = data.get("decode_rank", 0)
         assert len(block_notify_list) > 0, (
@@ -543,7 +550,7 @@ class MoRIIOWrapper:
         )
 
         with self.lock:
-            self.done_remote_allocate_req_dict[req_id] = RemoteAllocInfo(
+            self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(
                 block_ids=block_notify_list, decode_dp_rank=decode_dp_rank
             )
 

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

@@ -59,7 +59,12 @@ from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
 from vllm.v1.attention.backends.utils import get_kv_cache_layout
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, SlidingWindowSpec
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    MambaSpec,
+    SlidingWindowSpec,
+    UniformTypeKVCacheSpecs,
+)
 from vllm.v1.worker.block_table import BlockTable
 from vllm.v1.worker.utils import select_common_block_size
 
@@ -159,6 +164,7 @@ class NixlAgentMetadata:
     block_lens: list[int]
     kv_cache_layout: str
     block_size: int
+    ssm_sizes: tuple[int, int]
 
 
 @dataclass
@@ -310,6 +316,15 @@ class NixlConnectorMetadata(KVConnectorMetadata):
 class NixlConnector(KVConnectorBase_V1, SupportsHMA):
     @property
     def prefer_cross_layer_blocks(self) -> bool:
+        if any(
+            [
+                isinstance(group.kv_cache_spec, MambaSpec)
+                for group in self.kv_cache_config.kv_cache_groups
+            ]
+        ):
+            # Hybrid SSM models do not yet support cross-layer layout
+            return False
+
         backend = get_current_attn_backend(self._vllm_config)
         if backend.get_name() not in (
             "FLASH_ATTN",
@@ -335,12 +350,9 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
         kv_cache_config: "KVCacheConfig",
     ):
         super().__init__(vllm_config, role, kv_cache_config)
-
         assert vllm_config.kv_transfer_config is not None
         assert vllm_config.kv_transfer_config.engine_id is not None
-        for group in kv_cache_config.kv_cache_groups:
-            if isinstance(group.kv_cache_spec, MambaSpec):
-                raise ValueError("NixlConnector does not support Mamba models.")
+        self.kv_cache_config = kv_cache_config
         self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
         self.kv_transfer_config = vllm_config.kv_transfer_config
         if role == KVConnectorRole.SCHEDULER:
@@ -403,6 +415,14 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
         assert self.connector_scheduler is not None
         return self.connector_scheduler.build_connector_meta(scheduler_output)
 
+    def request_finished(
+        self,
+        request: "Request",
+        block_ids: list[int],
+    ) -> tuple[bool, dict[str, Any] | None]:
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.request_finished(request, (block_ids,))
+
     def request_finished_all_groups(
         self,
         request: "Request",
@@ -434,11 +454,7 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):
         self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
     ):
         assert self.connector_worker is not None
-
-        cross_layer_name = "ALL_LAYERS"
-        kv_caches = {cross_layer_name: kv_cache}
-
-        self.connector_worker.register_kv_caches(kv_caches)
+        self.connector_worker.register_cross_layers_kv_caches(kv_cache)
 
     def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
         assert self.connector_worker is not None
@@ -962,6 +978,40 @@ class NixlConnectorWorker:
             )
         )
         self.kv_cache_config = kv_cache_config
+        self._layer_specs = {
+            layer: group.kv_cache_spec
+            for group in kv_cache_config.kv_cache_groups
+            for layer in group.layer_names
+        }
+        self.hma_group_size = len(kv_cache_config.kv_cache_tensors)
+
+        # Mamba metadata
+        self._is_mamba_group = [
+            isinstance(group.kv_cache_spec, MambaSpec)
+            for group in kv_cache_config.kv_cache_groups
+        ]
+        mamba_ssm_size = (0, 0)
+        self._has_mamba = any(self._is_mamba_group)
+        if self._has_mamba:
+            assert self._is_hma_required
+            mamba_spec = next(
+                spec
+                for spec in self._layer_specs.values()
+                if isinstance(spec, MambaSpec)
+            )
+            conv_nbytes, ssm_nbytes = (
+                torch.tensor([], dtype=mamba_spec.dtypes[0]).element_size(),  # type: ignore[misc]
+                torch.tensor([], dtype=mamba_spec.dtypes[1]).element_size(),  # type: ignore[misc]
+            )
+            conv_shape, ssm_shape = (
+                torch.Size(mamba_spec.shapes[0]),
+                torch.Size(mamba_spec.shapes[1]),
+            )
+            mamba_ssm_size = (
+                conv_shape.numel() * conv_nbytes,
+                ssm_shape.numel() * ssm_nbytes,
+            )
+        self._mamba_ssm_size = mamba_ssm_size
 
         # Agent.
         non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
@@ -1106,9 +1156,9 @@ class NixlConnectorWorker:
 
         # Get the attention backend from the first layer
         # NOTE (NickLucche) models with multiple backends are not supported yet
-        self.attn_backend = get_current_attn_backend(vllm_config)
+        self.attn_backends = get_current_attn_backends(vllm_config)
+        self.backend_name = self.attn_backends[0].get_name()
 
-        self.backend_name = self.attn_backend.get_name()
         self.kv_cache_layout = get_kv_cache_layout()
         self.host_buffer_kv_cache_layout = self.kv_cache_layout
         logger.info("Detected attention backend %s", self.backend_name)
@@ -1135,6 +1185,8 @@ class NixlConnectorWorker:
     def _sync_block_size_with_kernel(self) -> None:
         backends = get_current_attn_backends(self.vllm_config)
         kernel_block_size = select_common_block_size(self.block_size, backends)
+        # Number of blocks not accounting for kernel block mismatches
+        self._logical_num_blocks = self.num_blocks
         if self.block_size != kernel_block_size:
             logger.info_once(
                 "User-specified logical block size (%s) does not match"
@@ -1428,9 +1480,19 @@ class NixlConnectorWorker:
 
         fut.add_done_callback(request_ready)
 
+    def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
+        """Register a cross-layers KV cache tensor with NIXL.
+
+        `use_uniform_kv_cache()` guarantees a single KV cache group whose
+        layers all share the same `AttentionSpec`, so any layer name from
+        `_layer_specs` yields the correct per-layer spec for `page_size_bytes`.
+        """
+        first_layer = next(iter(self._layer_specs))
+        # Forwarding a real layer name rather than a synthetic key
+        self.register_kv_caches({first_layer: kv_cache})
+
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in nixl."""
-
         self.kv_topo = TpKVTopology(
             tp_rank=self.tp_rank,
             engine_id=self.engine_id,
@@ -1438,8 +1500,12 @@ class NixlConnectorWorker:
             remote_block_size=self._block_size,  # shared state
             is_mla=self.use_mla,
             total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
-            attn_backend=self.attn_backend,
-            tensor_shape=next(iter(kv_caches.values())).shape,
+            attn_backends=self.attn_backends,
+            # SSM States come in tuples (ssm, conv)
+            tensor_shape=next(iter(kv_caches.values())).shape
+            if not self._has_mamba
+            else None,
+            is_mamba=self._has_mamba,
         )
         self.compat_hash = compute_nixl_compatibility_hash(
             self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
@@ -1481,12 +1547,50 @@ class NixlConnectorWorker:
         # to better exploit the memory layout (ie num_blocks is the first dim).
         tensor_size_bytes = None
 
-        # Enable different block lengths for different layers when MLA is used.
+        # Enable different block lengths for different layers *only* when MLA is used.
+        # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.
         self.block_len_per_layer = list[int]()
         for layer_name, cache_or_caches in xfer_buffers.items():
-            cache_list = (
-                cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches]
+            # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
+            # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
+            # However, physical page_size may differ when kernel requires a specific
+            # block size. This leads to SSM and FA layers having different num_blocks.
+            # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.
+            layer_spec = self._layer_specs[layer_name]
+            if isinstance(layer_spec, UniformTypeKVCacheSpecs):
+                # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
+                layer_spec = layer_spec.kv_cache_specs[layer_name]
+            cache_list = self.kv_topo.get_transfer_cache_regions(
+                cache_or_caches, layer_spec
             )
+            # `layer_spec.page_size_bytes` only accounts for logical page_size, that is
+            # the page_size assuming constant `self._logical_num_blocks`.
+            physical_page_size = (
+                layer_spec.page_size_bytes
+                if isinstance(layer_spec, MambaSpec)
+                else layer_spec.page_size_bytes
+                // self._physical_blocks_per_logical_kv_block
+            )
+            # For when registering multiple tensors eg K/V in separate regions.
+            physical_page_size = physical_page_size // len(cache_list)
+            if self.kv_topo._cross_layers_blocks:
+                # When cross-layers blocks are used, multiply by number of layers
+                physical_page_size = physical_page_size * len(
+                    self.kv_cache_config.kv_cache_tensors
+                )
+            num_blocks = (
+                self._logical_num_blocks
+                if isinstance(layer_spec, MambaSpec)
+                else self.num_blocks
+            )
+            # `page_size` accounts for physical blocks, st KVCache is always
+            # [`num_blocks` * `page_size`]
+            curr_tensor_size_bytes = num_blocks * physical_page_size
+            if tensor_size_bytes is None:
+                tensor_size_bytes = curr_tensor_size_bytes
+
+            # TODO (NickLucche) we could eventually unify how we handle FA/FI regions,
+            # registering a single tensor for both K/V and splitting logically like FI.
             for cache in cache_list:
                 base_addr = cache.data_ptr()
                 if base_addr in seen_base_addresses:
@@ -1494,27 +1598,27 @@ class NixlConnectorWorker:
                     # across groups. This results in skipping all tensors but the ones
                     # pointed to by group0. Also, generally we will have more blocks
                     # per tensor but fewer regions.
+                    logger.debug("Skipping %s because it's already seen", layer_name)
                     continue
-
                 logger.debug(
                     "Registering layer %s with cache shape: %s", layer_name, cache.shape
                 )
                 seen_base_addresses.append(base_addr)
-                curr_tensor_size_bytes = cache.numel() * cache.element_size()
+                # Only record non-Mamba page sizes.
+                if isinstance(layer_spec, MambaSpec):
+                    self.block_len_per_layer.append(
+                        physical_page_size // self._physical_blocks_per_logical_kv_block
+                    )
+                else:
+                    self.block_len_per_layer.append(physical_page_size)
 
-                if tensor_size_bytes is None:
-                    tensor_size_bytes = curr_tensor_size_bytes
-
-                assert cache.shape[0] == self.num_blocks, (
+                assert cache.shape[0] == num_blocks, (
                     "All kv cache tensors must have the same number of blocks"
                 )
 
-                self.block_len_per_layer.append(
-                    curr_tensor_size_bytes // self.num_blocks
-                )
-
                 if not self.use_mla:
-                    # Different kv cache shape is not supported by HeteroTP
+                    # Different kv cache shape is not supported by HeteroTP.
+                    # This must also hold true for Mamba-like models.
                     assert tensor_size_bytes == curr_tensor_size_bytes, (
                         "All kv cache tensors must have the same size"
                     )
@@ -1533,6 +1637,21 @@ class NixlConnectorWorker:
         self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
         self.num_regions = len(caches_data)
 
+        if self.kv_topo.is_kv_layout_blocks_first:
+            # NOTE (NickLucche) When FlashInfer is used, memory is registered
+            # with joint KV for each block. This minimizes the overhead in
+            # registerMem allowing faster descs queries. In order to be able to
+            # split on kv_heads dim as required by heterogeneous TP, one must
+            # be able to index K/V separately. Hence we double the number
+            # of 'virtual' regions here and halve `block_len` below.
+            # Similarly for Mamba layers, we register SSM+Conv as a single region and
+            # then duplicate it logically to be able to index SSM/Conv separately.
+            self.num_regions *= 2
+
+        # TODO (NickLucche) Adapt to different descs views (engine_id->tp_rank) to
+        # support heterogeneous TP.
+        self.num_descs = self.num_regions * self.num_blocks
+
         descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
         logger.debug("Registering descs: %s", caches_data)
         self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
@@ -1542,17 +1661,21 @@ class NixlConnectorWorker:
         self.device_kv_caches = kv_caches
         self.dst_num_blocks[self.engine_id] = self.num_blocks
 
-        if self.kv_topo.is_kv_layout_blocks_first:
-            # NOTE (NickLucche) When FlashInfer is used, memory is registered
-            # with joint KV for each block. This minimizes the overhead in
-            # registerMem allowing faster descs queries. In order to be able to
-            # split on kv_heads dim as required by heterogeneous TP, one must
-            # be able to index K/V separately. Hence we double the number
-            # of 'virtual' regions here and halve `block_len` below.
-            self.num_regions *= 2
+        if self._has_mamba:
+            logger.info(
+                "Hybrid SSM registration: num_blocks=%s, "
+                "logical_num_blocks=%s, ratio=%s, num_regions=%s, "
+                "num_descs=%s, mamba_ssm_size=%s, block_len_per_layer=%s",
+                self.num_blocks,
+                self._logical_num_blocks,
+                self._physical_blocks_per_logical_kv_block,
+                self.num_regions,
+                self.num_descs,
+                self._mamba_ssm_size,
+                set(self.block_len_per_layer),
+            )
 
         # Register local/src descr for NIXL xfer.
-        self.seen_base_addresses = seen_base_addresses
         self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = (
             self.register_local_xfer_handler(self.block_size)
         )
@@ -1569,6 +1692,7 @@ class NixlConnectorWorker:
             if not self.use_host_buffer
             else self.host_buffer_kv_cache_layout,
             block_size=self.block_size,
+            ssm_sizes=self._mamba_ssm_size,
         )
         # Wrap metadata in payload with hash for defensive decoding
         assert self.compat_hash is not None
@@ -1594,40 +1718,65 @@ class NixlConnectorWorker:
         data copy correctness.
         """
         assert self.kv_topo is not None
+        kv_topo = self.kv_topo
 
         block_size_ratio = self.block_size // block_size
-        blocks_data = []
-        for i, base_addr in enumerate(self.seen_base_addresses):
-            # The new block_len is using prefill block_len;
-            # and num_blocks is multiple with N
-            kv_block_len = (
-                self.get_backend_aware_kv_block_len(layer_idx=i) // block_size_ratio
-            )
-            block_len_per_layer = self.block_len_per_layer[i] // block_size_ratio
-            num_blocks = self.num_blocks * block_size_ratio
-            for block_id in range(num_blocks):
-                block_offset = block_id * block_len_per_layer
-                addr = base_addr + block_offset
-                # (addr, len, device id)
-                blocks_data.append((addr, kv_block_len, self.device_id))
+        blocks_data: list[tuple[int, int, int]] = []
+        local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
 
-            if self.kv_topo.is_kv_layout_blocks_first:
-                # Separate and interleave K/V regions to maintain the same
-                # descs ordering. This is needed for selecting contiguous heads
-                # when split across TP ranks.
+        def register_blocks(blocks_data: list[tuple[int, int, int]], mamba: bool):
+            for i, base_addr in enumerate(local_base_addresses):
+                # The new block_len is using prefill block_len;
+                # and num_blocks is multiple with N
+                kv_block_len = (
+                    self.get_backend_aware_kv_block_len(
+                        layer_idx=i, first_split=True, mamba_view=mamba
+                    )
+                    // block_size_ratio
+                )
+                # Jump one page_size, but ssm page_size may be bigger when kernel
+                # locks block size to a specific value.
+                block_len_per_layer = (
+                    self.block_len_per_layer[i]
+                    // block_size_ratio
+                    * (1 if not mamba else self._physical_blocks_per_logical_kv_block)
+                )
+                num_blocks = self._logical_num_blocks if mamba else self.num_blocks
+                num_blocks = num_blocks * block_size_ratio
                 for block_id in range(num_blocks):
                     block_offset = block_id * block_len_per_layer
                     addr = base_addr + block_offset
-                    # Register addresses for V cache (K registered first).
-                    v_addr = addr + kv_block_len
-                    blocks_data.append((v_addr, kv_block_len, self.device_id))
-        logger.debug(
-            "Created %s blocks for src engine %s and rank %s on device id %s",
-            len(blocks_data),
-            self.engine_id,
-            self.tp_rank,
-            self.device_id,
-        )
+                    # (addr, len, device id)
+                    blocks_data.append((addr, kv_block_len, self.device_id))
+
+                if kv_topo.is_kv_layout_blocks_first:
+                    second_split = self.get_backend_aware_kv_block_len(
+                        layer_idx=i, first_split=False, mamba_view=mamba
+                    )
+                    # Separate and interleave K/V regions to maintain the same
+                    # descs ordering. This is needed for selecting contiguous heads
+                    # when split across TP ranks.
+                    for block_id in range(num_blocks):
+                        block_offset = block_id * block_len_per_layer
+                        addr = base_addr + block_offset
+                        # Register addresses for V cache (K registered first).
+                        v_addr = addr + kv_block_len
+                        blocks_data.append((v_addr, second_split, self.device_id))
+            logger.debug(
+                "Created %s blocks for src engine %s and rank %s on device id %s",
+                len(blocks_data),
+                self.engine_id,
+                self.tp_rank,
+                self.device_id,
+            )
+
+        register_blocks(blocks_data, mamba=False)
+        if self._has_mamba:
+            assert self.num_descs == len(blocks_data)
+            logger.debug(
+                "Registering additional %s local Mamba blocks", len(blocks_data)
+            )
+            register_blocks(blocks_data, mamba=True)
 
         descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
         # NIXL_INIT_AGENT to be used for preparations of local descs.
@@ -1708,7 +1857,8 @@ class NixlConnectorWorker:
         # local origin:|          0|          1|          8|         12|
         # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
         assert self.kv_topo is not None
-        block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id)
+        kv_topo = self.kv_topo
+        block_size_ratio = kv_topo.block_size_ratio_from_engine_id(engine_id)
 
         if engine_id not in self.dst_num_blocks:
             self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
@@ -1768,48 +1918,86 @@ class NixlConnectorWorker:
         # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].
 
         # Register all remote blocks, but only the corresponding kv heads.
-        for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
-            # Read our whole local region size from remote.
-            local_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)
-            remote_kv_block_len = local_block_len // block_size_ratio
-            if block_size_ratio > 1:
-                # using remote kv_block_len as transfer unit
-                local_block_len = remote_kv_block_len
+        def register_remote_blocks(
+            blocks_data: list[tuple[int, int, int]], mamba: bool
+        ):
+            for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
+                # Read our whole local region size from remote.
+                local_block_len = self.get_backend_aware_kv_block_len(
+                    layer_idx=i, first_split=True, mamba_view=mamba
+                )
+                remote_kv_block_len = local_block_len // block_size_ratio
+                if block_size_ratio > 1:
+                    # using remote kv_block_len as transfer unit
+                    local_block_len = remote_kv_block_len
 
-            if tp_ratio < 0 and not self.use_mla:
-                # Remote tp is bigger: read a chunk of local region from remote
-                local_block_len = local_block_len // (-tp_ratio)
-            rank_offset = (
-                self.tp_rank % tp_ratio * remote_kv_block_len
-                if indexes_into_remote
-                else 0
-            )
-            for block_id in range(nixl_agent_meta.num_blocks):
-                block_offset = block_id * nixl_agent_meta.block_lens[i]
-                # For each block, grab the heads chunk belonging to rank_i
-                # of size remote_nheads // tp_ratio, which correspond to
-                # self.block_len == remote_block_len//tp_ratio bytes.
-                addr = base_addr + block_offset + rank_offset
-                # (addr, len, device id)
-                blocks_data.append((addr, local_block_len, nixl_agent_meta.device_id))
+                if tp_ratio < 0 and not self.use_mla:
+                    # Remote tp is bigger: read a chunk of local region from remote
+                    local_block_len = local_block_len // (-tp_ratio)
+                rank_offset = (
+                    self.tp_rank % tp_ratio * remote_kv_block_len
+                    if indexes_into_remote
+                    else 0
+                )
 
-            if self.kv_topo.is_kv_layout_blocks_first:
-                # With FlashInfer index V separately to allow head splitting.
-                for block_id in range(nixl_agent_meta.num_blocks):
-                    block_offset = block_id * nixl_agent_meta.block_lens[i]
+                # Assume same num_blocks for mamba and fa
+                num_blocks = (
+                    nixl_agent_meta.num_blocks
+                    if not mamba
+                    else nixl_agent_meta.num_blocks
+                    // self._physical_blocks_per_logical_kv_block
+                )
+                page_size = nixl_agent_meta.block_lens[i] * (
+                    1 if not mamba else self._physical_blocks_per_logical_kv_block
+                )
+                for block_id in range(num_blocks):
+                    block_offset = block_id * page_size
+                    # For each block, grab the heads chunk belonging to rank_i
+                    # of size remote_nheads // tp_ratio, which correspond to
+                    # self.block_len == remote_block_len//tp_ratio bytes.
                     addr = base_addr + block_offset + rank_offset
-                    v_addr = addr + nixl_agent_meta.block_lens[i] // 2
+                    # (addr, len, device id)
                     blocks_data.append(
-                        (v_addr, local_block_len, nixl_agent_meta.device_id)
+                        (addr, local_block_len, nixl_agent_meta.device_id)
                     )
 
-        logger.debug(
-            "Created %s blocks for dst engine %s with remote rank %s and local rank %s",
-            len(blocks_data),
-            engine_id,
-            remote_tp_rank,
-            self.tp_rank,
-        )
+                if kv_topo.is_kv_layout_blocks_first:
+                    # With FlashInfer index V separately to allow head splitting.
+                    second_split = self.get_backend_aware_kv_block_len(
+                        layer_idx=i, first_split=False, mamba_view=mamba
+                    )
+                    # Apply the same scaling as local_block_len above for when we read
+                    # a chunk of local V from `tp_ratio` separate remote workers.
+                    if tp_ratio < 0 and not self.use_mla:
+                        second_split = second_split // (-tp_ratio)
+                    for block_id in range(num_blocks):
+                        block_offset = block_id * page_size
+                        addr = base_addr + block_offset + rank_offset
+                        # Hop over the first split of remote page: either K or Conv.
+                        if mamba:
+                            v_addr = addr + nixl_agent_meta.ssm_sizes[0]
+                        else:
+                            v_addr = addr + nixl_agent_meta.block_lens[i] // 2
+                        blocks_data.append(
+                            (v_addr, second_split, nixl_agent_meta.device_id)
+                        )
+
+            logger.debug(
+                "Created %s blocks for dst engine %s"
+                " with remote rank %s and local rank %s",
+                len(blocks_data),
+                engine_id,
+                remote_tp_rank,
+                self.tp_rank,
+            )
+
+        register_remote_blocks(blocks_data, mamba=False)
+        if self._has_mamba:
+            # Create extra descs for the Mamba "view" of the same KV cache tensors.
+            logger.debug(
+                "Registering additional %s remote Mamba blocks", len(blocks_data)
+            )
+            register_remote_blocks(blocks_data, mamba=True)
 
         # Register with NIXL.
         descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
@@ -1849,6 +2037,9 @@ class NixlConnectorWorker:
             assert block_size_ratio == 1, (
                 "HMA does not support different remote block size yet"
             )
+        # Mamba additional constraints
+        if self._has_mamba:
+            assert tp_ratio == 1, "Mamba does not support heterogeneous TP yet"
 
         kv_cache_layout = (
             self.kv_cache_layout
@@ -2495,6 +2686,7 @@ class NixlConnectorWorker:
         A single flattened array is returned for all groups anyway.
         """
         region_ids = np.arange(self.num_regions)
+
         # NOTE (NickLucche) With HMA, every kv group has the same number of layers and
         # layers from different groups share the same kv tensor.
         # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be read across all regions,
@@ -2505,11 +2697,33 @@ class NixlConnectorWorker:
         if block_size_ratio is not None:
             num_blocks = int(num_blocks * block_size_ratio)
 
-        # Compute the desc ids for each block.
+        # Compute desc ids per group using the right stride: FA descs have
+        # num_blocks entries per region (kernel granularity), SSM descs have
+        # logical_blocks entries per region (no kernel splitting).
         region_ids = region_ids[:, None]
-        block_ids = np.concatenate(block_ids)[None, :]
-        descs_ids = region_ids * num_blocks + block_ids
-        return descs_ids.flatten()
+        if not self._has_mamba:
+            block_ids = np.concatenate(block_ids)[None, :]
+            descs_ids = region_ids * num_blocks + block_ids
+            return descs_ids.flatten()
+        else:
+            # NOTE (NickLucche) SSM and Attention blocks regions can be exchanged
+            # arbitrarily by manager. Therefore, descs are duplicated for SSM and
+            # Attention like so:
+            # desc_handle->[descs_fa (all regions) | descs_ssm (all regions)].
+            # This is like having two "low-level views" of the same storage.
+            # `num_fa_descs` offset must be computed per-engine since P and D can
+            # have different num_blocks (and thus different FA descs counts).
+            ratio = self._physical_blocks_per_logical_kv_block
+            # SSM may register fewer num_blocks than FA
+            logical_blocks = num_blocks // ratio
+            num_fa_descs = self.num_regions * num_blocks
+            all_descs = []
+            for i, group in enumerate(block_ids):
+                stride = logical_blocks if self._is_mamba_group[i] else num_blocks
+                group_arr = np.asarray(group)[None, :]
+                offset = num_fa_descs if self._is_mamba_group[i] else 0
+                all_descs.append((region_ids * stride + group_arr + offset).flatten())
+            return np.concatenate(all_descs)
 
     def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:
         """
@@ -2523,16 +2737,22 @@ class NixlConnectorWorker:
         block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape(
             1, -1
         )
+        # Mamba blocks have no logical<>physical discrepancy
+        group_specs = self.kv_cache_config.kv_cache_groups
         return [
             BlockTable.map_to_kernel_blocks(
                 np.array(group),
                 self._physical_blocks_per_logical_kv_block,
                 block_arange,
             ).tolist()
-            for group in block_ids
+            if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)
+            else group
+            for i, group in enumerate(block_ids)
         ]
 
-    def get_backend_aware_kv_block_len(self, layer_idx: int) -> int:
+    def get_backend_aware_kv_block_len(
+        self, layer_idx: int, first_split: bool = True, mamba_view: bool = False
+    ) -> int:
         """
         Get the block length for one K/V element (K and V have the same size).
 
@@ -2540,11 +2760,38 @@ class NixlConnectorWorker:
         block, as K and V are in separate regions.
         For FlashInfer, this is half the length of the whole block, as K and V
         share the same region.
+        Similarly, for SSM-based models, state and conv are interleaved, but crucially
+        the their size differs.
+        Reference diagram:
+                            KVCacheTensor (Shared)
+                               /       \
+                              /         \
+                             /           \
+        Attention (FlashInfer) View      Mamba View
+                  |                          |
+                  |                          |
+           +-------------------+         +-------------------+
+           | KVCacheTensor     |         | KVCacheTensor      |
+           |                   |         |                    |
+           |<----- page ------>|         |<----- page ------->|
+           |       size        |         |       size         |
+           |  Key 0  |  Val 0  |         |Conv 0  |   SSM 0   |
+           |  Key 1  |  Val 1  |         |Conv 1  |   SSM 1   |
+           |   ...   |   ...   |         |  ...   |    ...    |
+           | Key N-2 | Val N-2 |         |Conv N-2|   SSM N-2 |
+           | Key N-1 | Val N-1 |         |Conv N-1|   SSM N-1 |
+           +-------------------+         +--------------------+
+           |1st_split-2nd_split|         |1st_split-2nd_split |
         """
         assert self.kv_topo is not None
         if self.kv_topo.is_kv_layout_blocks_first:
             # For indexing only half (either just the K or V part).
-            block_len = self.block_len_per_layer[layer_idx] // 2
+            if mamba_view:
+                # NOTE (NickLucche) Mamba Opt: this is already skipping the padding so
+                # we're only transferring the minimum required bytes.
+                block_len = self._mamba_ssm_size[not first_split]
+            else:
+                block_len = self.block_len_per_layer[layer_idx] // 2
         else:
             block_len = self.block_len_per_layer[layer_idx]
         return block_len

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -24,7 +24,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
 )
 from vllm.forward_context import ForwardContext
 from vllm.logger import init_logger
-from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.kv_cache_utils import BlockHash
@@ -601,7 +601,9 @@ class OffloadingConnectorWorker:
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         layer_names = list(kv_caches.keys())
         layers = get_layers_from_vllm_config(
-            self.spec.vllm_config, Attention, layer_names
+            self.spec.vllm_config,
+            AttentionLayerBase,  # type: ignore[type-abstract]
+            layer_names,
         )
         attn_backends = {
             layer_name: layers[layer_name].get_attn_backend()

vllm/engine/arg_utils.py

@@ -614,6 +614,7 @@ class EngineArgs:
     )
 
     fail_on_environ_validation: bool = False
+    gdn_prefill_backend: Literal["flashinfer", "triton"] | None = None
 
     def __post_init__(self):
         # support `EngineArgs(compilation_config={...})`
@@ -1318,6 +1319,13 @@ class EngineArgs:
             help="Shutdown timeout in seconds. 0 = abort, >0 = wait.",
         )
 
+        parser.add_argument(
+            "--gdn-prefill-backend",
+            dest="gdn_prefill_backend",
+            choices=["flashinfer", "triton"],
+            default=None,
+            help="Select GDN prefill backend.",
+        )
         return parser
 
     @classmethod
@@ -1903,6 +1911,9 @@ class EngineArgs:
             ),
         )
 
+        if self.gdn_prefill_backend is not None:
+            self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend
+
         config = VllmConfig(
             model_config=model_config,
             cache_config=cache_config,

vllm/entrypoints/anthropic/protocol.py

@@ -34,7 +34,14 @@ class AnthropicUsage(BaseModel):
 class AnthropicContentBlock(BaseModel):
     """Content block in message"""
 
-    type: Literal["text", "image", "tool_use", "tool_result", "thinking"]
+    type: Literal[
+        "text",
+        "image",
+        "tool_use",
+        "tool_result",
+        "thinking",
+        "redacted_thinking",
+    ]
     text: str | None = None
     # For image content
     source: dict[str, Any] | None = None
@@ -48,6 +55,8 @@ class AnthropicContentBlock(BaseModel):
     # For thinking content
     thinking: str | None = None
     signature: str | None = None
+    # For redacted thinking content (safety-filtered by the API)
+    data: str | None = None
 
 
 class AnthropicMessage(BaseModel):

vllm/entrypoints/anthropic/serving.py

@@ -224,6 +224,12 @@ class AnthropicServingMessages(OpenAIServingChat):
             content_parts.append({"type": "image_url", "image_url": {"url": image_url}})
         elif block.type == "thinking" and block.thinking is not None:
             reasoning_parts.append(block.thinking)
+        elif block.type == "redacted_thinking":
+            # Redacted thinking blocks contain safety-filtered reasoning.
+            # We skip them as the content is opaque (base64 'data' field),
+            # but accepting the block prevents a validation error when the
+            # client echoes back the full assistant message.
+            pass
         elif block.type == "tool_use":
             cls._convert_tool_use_block(block, tool_calls)
         elif block.type == "tool_result":

vllm/entrypoints/cli/launch.py

@@ -116,6 +116,11 @@ async def run_launch_fastapi(args: argparse.Namespace) -> None:
     # 2. Build and serve the API server
     engine_args = AsyncEngineArgs.from_cli_args(args)
     model_config = engine_args.create_model_config()
+
+    # Render servers preprocess data only — no inference, no quantized kernels.
+    # Clear quantization so VllmConfig skips quant dtype/capability validation.
+    model_config.quantization = None
+
     vllm_config = VllmConfig(model_config=model_config)
     shutdown_task = await build_and_serve_renderer(
         vllm_config, listen_address, sock, args

vllm/entrypoints/openai/api_server.py

@@ -29,11 +29,13 @@ from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.launcher import serve_http
 from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
+from vllm.entrypoints.openai.engine.protocol import GenerationError
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.server_utils import (
     engine_error_handler,
     exception_handler,
+    generation_error_handler,
     get_uvicorn_log_config,
     http_exception_handler,
     lifespan,
@@ -263,6 +265,7 @@ def build_app(
     app.exception_handler(RequestValidationError)(validation_exception_handler)
     app.exception_handler(EngineGenerateError)(engine_error_handler)
     app.exception_handler(EngineDeadError)(engine_error_handler)
+    app.exception_handler(GenerationError)(generation_error_handler)
     app.exception_handler(Exception)(exception_handler)
 
     # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -7,7 +7,6 @@ import json
 import time
 from typing import Annotated, Any, ClassVar, Literal
 
-import torch
 from openai.types.chat.chat_completion_audio import (
     ChatCompletionAudio as OpenAIChatCompletionAudio,
 )
@@ -48,7 +47,8 @@ from vllm.utils import random_uuid
 logger = init_logger(__name__)
 
 
-_LONG_INFO = torch.iinfo(torch.long)
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
 
 
 class ChatMessage(OpenAIBaseModel):
@@ -165,7 +165,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
     n: int | None = 1
     presence_penalty: float | None = 0.0
     response_format: AnyResponseFormat | None = None
-    seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)
+    seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
     stop: str | list[str] | None = []
     stream: bool | None = False
     stream_options: StreamOptions | None = None
@@ -198,9 +198,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
     min_tokens: int = 0
     skip_special_tokens: bool = True
     spaces_between_special_tokens: bool = True
-    truncate_prompt_tokens: Annotated[int, Field(ge=-1, le=_LONG_INFO.max)] | None = (
-        None
-    )
+    truncate_prompt_tokens: Annotated[int, Field(ge=-1, le=_INT64_MAX)] | None = None
     prompt_logprobs: int | None = None
     allowed_token_ids: list[int] | None = None
     bad_words: list[str] = Field(default_factory=list)
@@ -285,6 +283,8 @@ class ChatCompletionRequest(OpenAIBaseModel):
     )
     priority: int = Field(
         default=0,
+        ge=_INT64_MIN,
+        le=_INT64_MAX,
         description=(
             "The priority of the request (lower means earlier handling; "
             "default: 0). Any priority other than 0 will raise an error "

vllm/entrypoints/openai/chat_completion/serving.py

@@ -6,6 +6,7 @@ import json
 import time
 from collections.abc import AsyncGenerator, AsyncIterator
 from collections.abc import Sequence as GenericSequence
+from http import HTTPStatus
 from typing import TYPE_CHECKING, Any, Final
 
 import partial_json_parser
@@ -1289,7 +1290,12 @@ class OpenAIServingChat(OpenAIServing):
         except asyncio.CancelledError:
             return self.create_error_response("Client disconnected")
 
-        assert final_res is not None
+        if final_res is None:
+            return self.create_error_response(
+                "No output received from the engine.",
+                err_type="InternalServerError",
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+            )
 
         choices: list[ChatCompletionResponseChoice] = []
         if self.tool_call_id_type == "kimi_k2":

vllm/entrypoints/openai/completion/protocol.py

@@ -7,7 +7,6 @@ import json
 import time
 from typing import Annotated, Any, Literal
 
-import torch
 from pydantic import Field, model_validator
 
 from vllm.config import ModelConfig
@@ -36,7 +35,8 @@ from vllm.utils import random_uuid
 logger = init_logger(__name__)
 
 
-_LONG_INFO = torch.iinfo(torch.long)
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
 
 
 class CompletionRequest(OpenAIBaseModel):
@@ -57,7 +57,7 @@ class CompletionRequest(OpenAIBaseModel):
     max_tokens: int | None = 16
     n: int = 1
     presence_penalty: float | None = 0.0
-    seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)
+    seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
     stop: str | list[str] | None = []
     stream: bool | None = False
     stream_options: StreamOptions | None = None
@@ -78,9 +78,7 @@ class CompletionRequest(OpenAIBaseModel):
     min_tokens: int = 0
     skip_special_tokens: bool = True
     spaces_between_special_tokens: bool = True
-    truncate_prompt_tokens: Annotated[int, Field(ge=-1, le=_LONG_INFO.max)] | None = (
-        None
-    )
+    truncate_prompt_tokens: Annotated[int, Field(ge=-1, le=_INT64_MAX)] | None = None
     allowed_token_ids: list[int] | None = None
     prompt_logprobs: int | None = None
     # --8<-- [end:completion-sampling-params]
@@ -108,6 +106,8 @@ class CompletionRequest(OpenAIBaseModel):
     )
     priority: int = Field(
         default=0,
+        ge=_INT64_MIN,
+        le=_INT64_MAX,
         description=(
             "The priority of the request (lower means earlier handling; "
             "default: 0). Any priority other than 0 will raise an error "

vllm/entrypoints/openai/responses/protocol.py

@@ -6,7 +6,6 @@
 import time
 from typing import Any, Literal, TypeAlias
 
-import torch
 from openai.types.responses import (
     ResponseCodeInterpreterCallCodeDeltaEvent,
     ResponseCodeInterpreterCallCodeDoneEvent,
@@ -78,7 +77,8 @@ from vllm.utils import random_uuid
 
 logger = init_logger(__name__)
 
-_LONG_INFO = torch.iinfo(torch.long)
+_INT64_MIN = -(2**63)
+_INT64_MAX = 2**63 - 1
 
 
 class InputTokensDetails(OpenAIBaseModel):
@@ -210,6 +210,8 @@ class ResponsesRequest(OpenAIBaseModel):
     )
     priority: int = Field(
         default=0,
+        ge=_INT64_MIN,
+        le=_INT64_MAX,
         description=(
             "The priority of the request (lower means earlier handling; "
             "default: 0). Any priority other than 0 will raise an error "
@@ -246,7 +248,7 @@ class ResponsesRequest(OpenAIBaseModel):
     )
 
     repetition_penalty: float | None = None
-    seed: int | None = Field(None, ge=_LONG_INFO.min, le=_LONG_INFO.max)
+    seed: int | None = Field(None, ge=_INT64_MIN, le=_INT64_MAX)
     stop: str | list[str] | None = []
     ignore_eos: bool = False
     vllm_xargs: dict[str, str | int | float | list[str | int | float]] | None = Field(

vllm/entrypoints/openai/server_utils.py

@@ -21,7 +21,11 @@ from starlette.types import ASGIApp, Message, Receive, Scope, Send
 from vllm import envs
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.launcher import terminate_if_errored
-from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
+from vllm.entrypoints.openai.engine.protocol import (
+    ErrorInfo,
+    ErrorResponse,
+    GenerationError,
+)
 from vllm.entrypoints.utils import create_error_response, sanitize_message
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import init_logger
@@ -354,6 +358,17 @@ async def engine_error_handler(
     return JSONResponse(err.model_dump(), status_code=err.error.code)
 
 
+async def generation_error_handler(req: Request, exc: GenerationError):
+    """Handle GenerationError without logging stack traces.
+
+    GenerationError is a known, expected error (e.g. KV cache load failure)
+    that should be returned to the client as a 500 response without polluting
+    server logs with stack traces.
+    """
+    err = create_error_response(exc)
+    return JSONResponse(err.model_dump(), status_code=err.error.code)
+
+
 async def exception_handler(req: Request, exc: Exception):
     if req.app.state.args.log_error_stack:
         logger.exception(

vllm/entrypoints/openai/speech_to_text/speech_to_text.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
+import io
 import math
 import time
 import zlib
@@ -35,7 +36,6 @@ from vllm.entrypoints.openai.speech_to_text.protocol import (
     TranslationSegment,
     TranslationStreamResponse,
 )
-from vllm.entrypoints.openai.speech_to_text.utils import load_audio_bytes
 from vllm.entrypoints.utils import get_max_tokens
 from vllm.exceptions import VLLMValidationError
 from vllm.inputs import EncoderDecoderInputs, ProcessorInputs
@@ -43,6 +43,7 @@ from vllm.logger import init_logger
 from vllm.logprobs import FlatLogprobs, Logprob
 from vllm.model_executor.models import SupportsTranscription
 from vllm.multimodal.audio import split_audio
+from vllm.multimodal.media.audio import extract_audio_from_video_bytes
 from vllm.outputs import RequestOutput
 from vllm.renderers.inputs import DictPrompt, EncoderDecoderDictPrompt
 from vllm.renderers.inputs.preprocess import parse_enc_dec_prompt, parse_model_prompt
@@ -55,6 +56,19 @@ try:
 except ImportError:
     librosa = PlaceholderModule("librosa")  # type: ignore[assignment]
 
+try:
+    import soundfile as sf
+except ImportError:
+    sf = PlaceholderModule("soundfile")  # type: ignore[assignment]
+
+# Public libsndfile error codes exposed via `soundfile.LibsndfileError.code`, soundfile
+# being librosa's main backend. Used to validate if an audio loading error is due to a
+# server error vs a client error (invalid audio file).
+# 1 = unrecognised format      (file is not a supported audio container)
+# 3 = malformed file           (corrupt or structurally invalid audio)
+# 4 = unsupported encoding     (codec not supported by this libsndfile build)
+_BAD_SF_CODES = {1, 3, 4}
+
 SpeechToTextResponse: TypeAlias = TranscriptionResponse | TranslationResponse
 SpeechToTextResponseVerbose: TypeAlias = (
     TranscriptionResponseVerbose | TranslationResponseVerbose
@@ -198,7 +212,30 @@ class OpenAISpeechToText(OpenAIServing):
         # transparently falls back to ffmpeg via an in-memory fd.
         # NOTE resample to model SR here for efficiency. This is also a
         # pre-requisite for chunking, as it assumes Whisper SR.
-        y, sr = load_audio_bytes(audio_data, sr=self.asr_config.sample_rate)
+        try:
+            with io.BytesIO(audio_data) as buf:
+                y, sr = librosa.load(buf, sr=self.asr_config.sample_rate)  # type: ignore[return-value]
+        except sf.LibsndfileError as exc:
+            # Only fall back for known format-detection failures.
+            # Re-raise anything else (e.g. corrupt but recognised format).
+            if exc.code not in _BAD_SF_CODES:
+                raise
+            logger.debug(
+                "librosa/soundfile could not decode audio from BytesIO "
+                "(code=%s: %s); falling back to pyav in-process decode",
+                exc.code,
+                exc,
+            )
+            try:
+                native_y, native_sr = extract_audio_from_video_bytes(audio_data)
+                sr = self.asr_config.sample_rate
+                y = librosa.resample(native_y, orig_sr=native_sr, target_sr=sr)
+            except Exception as pyav_exc:
+                logger.debug(
+                    "pyAV fallback also failed: %s",
+                    pyav_exc,
+                )
+                raise ValueError("Invalid or unsupported audio file.") from pyav_exc
 
         duration = librosa.get_duration(y=y, sr=sr)
         do_split_audio = (

vllm/entrypoints/openai/speech_to_text/utils.py

@@ -1,106 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Audio decoding utilities for the speech-to-text endpoints."""
-
-import io
-
-import numpy as np
-import torchaudio
-
-from vllm.logger import init_logger
-from vllm.utils.import_utils import PlaceholderModule
-
-try:
-    import librosa
-except ImportError:
-    librosa = PlaceholderModule("librosa")  # type: ignore[assignment]
-
-try:
-    import soundfile as sf
-except ImportError:
-    sf = PlaceholderModule("soundfile")  # type: ignore[assignment]
-
-logger = init_logger(__name__)
-
-# Public libsndfile error codes exposed via ``soundfile.LibsndfileError.code``.
-# soundfile is librosa's primary backend.  These codes indicate that the audio
-# data itself is problematic (unrecognised container, corrupt file, or
-# unsupported encoding) rather than a transient server error.
-# 1 = unrecognised format, 3 = malformed file, 4 = unsupported encoding
-_BAD_SF_CODES = {1, 3, 4}
-
-
-def _decode_audio_bytes_torchaudio(
-    audio_data: bytes,
-    sr: int,
-) -> tuple[np.ndarray, int]:
-    """Decode audio bytes to mono float32 PCM via torchaudio, in-process.
-
-    ``torchaudio.load`` (backed by TorchCodec / FFmpeg) can decode
-    container formats (MP4, M4A, WebM) directly from a ``BytesIO``
-    buffer without spawning a subprocess.  The decoded waveform is
-    down-mixed to mono and resampled to *sr* Hz, matching the return
-    convention of ``librosa.load``.
-    """
-    buf = io.BytesIO(audio_data)
-    waveform, orig_sr = torchaudio.load(buf)
-
-    # Down-mix to mono (average across channels).
-    if waveform.shape[0] > 1:
-        waveform = waveform.mean(dim=0, keepdim=True)
-
-    # Resample to the target sample rate when necessary.
-    if orig_sr != sr:
-        waveform = torchaudio.functional.resample(
-            waveform, orig_freq=orig_sr, new_freq=sr
-        )
-
-    # Squeeze channel dim → 1-D float32 numpy array (same as librosa.load).
-    y = waveform.squeeze(0).numpy()
-    if y.size == 0:
-        raise RuntimeError(
-            "torchaudio produced no audio samples (file may be empty or corrupt)"
-        )
-    return y, sr
-
-
-def load_audio_bytes(
-    audio_data: bytes,
-    sr: int | float,
-) -> tuple[np.ndarray, int]:
-    """Load audio from raw bytes, with an in-process torchaudio fallback.
-
-    First tries ``librosa.load(BytesIO(...))`` which works for formats
-    that *soundfile* can auto-detect (WAV, FLAC, MP3, OGG, ...).  If
-    that fails with a ``LibsndfileError`` indicating an unrecognised or
-    unsupported format (typically container formats like MP4/M4A/WebM),
-    the bytes are decoded in-process via ``torchaudio`` (backed by
-    TorchCodec / FFmpeg) which handles these containers natively.
-    """
-    sr = int(sr)
-
-    # Fast path: librosa + soundfile (works for most formats).
-    try:
-        with io.BytesIO(audio_data) as buf:
-            return librosa.load(buf, sr=sr)  # type: ignore[return-value]
-    except sf.LibsndfileError as exc:
-        # Only fall back for known format-detection failures.
-        # Re-raise anything else (e.g. corrupt but recognised format).
-        if exc.code not in _BAD_SF_CODES:
-            raise
-        logger.debug(
-            "librosa/soundfile could not decode audio from BytesIO "
-            "(code=%s: %s); falling back to torchaudio in-process decode",
-            exc.code,
-            exc,
-        )
-
-    # Fallback: torchaudio in-process decode (no subprocess overhead).
-    try:
-        return _decode_audio_bytes_torchaudio(audio_data, sr)
-    except Exception as ta_exc:
-        logger.debug(
-            "torchaudio fallback also failed: %s",
-            ta_exc,
-        )
-        raise ValueError("Invalid or unsupported audio file.") from ta_exc

vllm/entrypoints/pooling/base/protocol.py

@@ -34,6 +34,8 @@ class PoolingBasicRequestMixin(OpenAIBaseModel):
     )
     priority: int = Field(
         default=0,
+        ge=-(2**63),
+        le=2**63 - 1,
         description=(
             "The priority of the request (lower means earlier handling; "
             "default: 0). Any priority other than 0 will raise an error "

vllm/entrypoints/serve/disagg/protocol.py

@@ -93,6 +93,8 @@ class GenerateRequest(BaseModel):
     )
     priority: int = Field(
         default=0,
+        ge=-(2**63),
+        le=2**63 - 1,
         description=(
             "The priority of the request (lower means earlier handling; "
             "default: 0). Any priority other than 0 will raise an error "

vllm/entrypoints/serve/render/serving.py

@@ -506,6 +506,7 @@ class OpenAIServingRender:
         (ResponsesRequest not supported here); TODO comment dropped accordingly.
         """
         renderer = self.renderer
+        mm_config = self.model_config.multimodal_config
 
         default_template_kwargs = merge_kwargs(
             default_template_kwargs,
@@ -518,7 +519,11 @@ class OpenAIServingRender:
         tok_params = request.build_tok_params(self.model_config)
         chat_params = request.build_chat_params(
             default_template, default_template_content_format
-        ).with_defaults(default_template_kwargs)
+        ).with_defaults(
+            default_template_kwargs,
+            default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
+            default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
+        )
 
         (conversation,), (engine_prompt,) = await renderer.render_chat_async(
             [messages],