Revert "[TPU][Bugfix] fix OOM issue in CI test (#21550 )"

This reverts commit 40d86ee412.
[TPU][Bugfix] fix OOM issue in CI test (#21550 )
2025-07-24 23:04:37 -07:00 · 2025-07-24 23:01:53 -07:00 · 2025-07-24 22:54:23 -07:00 · 2025-07-24 22:53:59 -07:00 · 2025-07-24 22:46:43 -07:00 · 2025-07-25 13:46:06 +08:00
183 changed files with 6053 additions and 3002 deletions
@@ -62,7 +62,8 @@ echo "Results will be stored in: $RESULTS_DIR"
 echo "--- Installing Python dependencies ---"
 python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
    && python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
-    && python3 -m pip install --progress-bar off lm_eval[api]==0.4.4
+    && python3 -m pip install --progress-bar off lm_eval[api]==0.4.4 \
+    && python3 -m pip install --progress-bar off hf-transfer
 echo "--- Python dependencies installed ---"
 export VLLM_USE_V1=1
 export VLLM_XLA_CHECK_RECOMPILATION=1
@@ -134,7 +135,7 @@ run_and_track_test 1 "test_compilation.py" \
 run_and_track_test 2 "test_basic.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_basic.py"
 run_and_track_test 3 "test_accuracy.py::test_lm_eval_accuracy_v1_engine" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine"
+    "HF_HUB_DISABLE_XET=1 python3 -m pytest -s -v /workspace/vllm/tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine"
 run_and_track_test 4 "test_quantization_accuracy.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_quantization_accuracy.py"
 run_and_track_test 5 "examples/offline_inference/tpu.py" \
@@ -150,7 +151,7 @@ run_and_track_test 9 "test_multimodal.py" \
 run_and_track_test 10 "test_pallas.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py"
 run_and_track_test 11 "test_struct_output_generate.py" \
-    "python3 -m pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\""
+    "HF_HUB_DISABLE_XET=1 python3 -m pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py -k \"not test_structured_output_with_reasoning_matrices\""
 run_and_track_test 12 "test_moe_pallas.py" \
    "python3 -m pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py"
 run_and_track_test 13 "test_lora.py" \
@@ -31,4 +31,13 @@ docker run \
    VLLM_USE_V1=1 python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
    cd tests
    pytest -v -s v1/core
+    pytest -v -s v1/engine
+    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py
+    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py
+    pytest -v -s v1/structured_output
+    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_eagle.py
+    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_nixl_connector.py
+    pytest -v -s v1/test_serial_utils.py
+    pytest -v -s v1/test_utils.py
+    pytest -v -s v1/test_metrics_reader.py
 '
@@ -166,6 +166,7 @@ steps:
  - tests/v1/test_async_llm_dp.py
  - tests/v1/test_external_lb_dp.py
  - tests/v1/test_internal_lb_dp.py
+  - tests/v1/test_hybrid_lb_dp.py
  - tests/v1/engine/test_engine_core_client.py
  commands:
  # test with tp=2 and external_dp=2
@@ -178,6 +179,7 @@ steps:
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/test_external_lb_dp.py
  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/test_internal_lb_dp.py
+  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/test_hybrid_lb_dp.py
  - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
  - pytest -v -s distributed/test_utils.py
  - pytest -v -s compile/test_basic_correctness.py
@@ -718,6 +720,7 @@ steps:
  - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
  - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s kv_transfer/test_disagg.py
  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
+  - pytest -v -s models/multimodal/generation/test_maverick.py

 - label: Plugin Tests (2 GPUs) # 40min
  mirror_hardwares: [amdexperimental]
@@ -52,3 +52,15 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 # Docs
 /docs @hmellor
 mkdocs.yaml @hmellor
+
+# CPU
+/vllm/v1/worker/^cpu @bigPYJ1151
+/csrc/cpu @bigPYJ1151
+/vllm/platforms/cpu.py @bigPYJ1151
+/cmake/cpu_extension.cmake @bigPYJ1151
+/docker/Dockerfile.cpu @bigPYJ1151
+
+# Intel GPU
+/vllm/v1/worker/^xpu @jikunshang
+/vllm/platforms/xpu.py @jikunshang
+/docker/Dockerfile.xpu @jikunshang
@@ -635,7 +635,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                     "in CUDA target architectures.")
    endif()
  endif()
-  
+
  cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a" "${CUDA_ARCHS}")
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
    set(SRCS "csrc/quantization/cutlass_w8a8/moe/blockwise_scaled_group_mm_sm100.cu")
@@ -768,6 +768,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  list(APPEND VLLM_MOE_EXT_SRC "csrc/moe/moe_wna16.cu")
 endif()

+if(VLLM_GPU_LANG STREQUAL "CUDA")
+  set(MOE_PERMUTE_SRC
+      "csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu"
+      "csrc/moe/moe_permute_unpermute_op.cu")
+
+  list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}")
+endif()
+
 set_gencode_flags_for_srcs(
  SRCS "${VLLM_MOE_EXT_SRC}"
  CUDA_ARCHS "${CUDA_ARCHS}")
@@ -836,17 +844,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  endif()
 endif()

-if(VLLM_GPU_LANG STREQUAL "CUDA")
-  set(MOE_PERMUTE_SRC
-      "csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu"
-      "csrc/moe/moe_permute_unpermute_op.cu")
-
-  set_gencode_flags_for_srcs(
-    SRCS "${MARLIN_PERMUTE_SRC}"
-    CUDA_ARCHS "${MOE_PERMUTE_ARCHS}")
-
-  list(APPEND VLLM_MOE_EXT_SRC "${MOE_PERMUTE_SRC}")
-endif()
 message(STATUS "Enabling moe extension.")
 define_gpu_extension_target(
  _moe_C
@@ -39,6 +39,7 @@ You must set the following variables at the top of the script before execution.
 | `DOWNLOAD_DIR` | **Required.** Directory to download and load model weights from. | `""` (default download path) |
 | `INPUT_LEN` | **Required.** Request input length. | `4000` |
 | `OUTPUT_LEN` | **Required.** Request output length. | `16` |
+| `MAX_MODEL_LEN` | **Required.** Max model length. | `4096` |
 | `MIN_CACHE_HIT_PCT` | Prefix cache hit rate in percentage (0-100). Set to `0` to disable. | `60` |
 | `MAX_LATENCY_ALLOWED_MS` | The maximum allowed P99 end-to-end latency in milliseconds. Set to a very large number (e.g., `100000000000`) to effectively ignore the latency constraint. | `500` |
 | `NUM_SEQS_LIST` | A space-separated string of `max-num-seqs` values to test. | `"128 256"` |
@@ -69,6 +70,7 @@ Here are a few examples of how to configure the script for different goals:
 ```bash
 INPUT_LEN=1800
 OUTPUT_LEN=20
+MAX_MODEL_LEN=2048
 MIN_CACHE_HIT_PCT=0
 MAX_LATENCY_ALLOWED_MS=100000000000 # A very large number
 ```
@@ -80,6 +82,7 @@ MAX_LATENCY_ALLOWED_MS=100000000000 # A very large number
 ```bash
 INPUT_LEN=1800
 OUTPUT_LEN=20
+MAX_MODEL_LEN=2048
 MIN_CACHE_HIT_PCT=0
 MAX_LATENCY_ALLOWED_MS=500
 ```
@@ -91,6 +94,7 @@ MAX_LATENCY_ALLOWED_MS=500
 ```bash
 INPUT_LEN=1800
 OUTPUT_LEN=20
+MAX_MODEL_LEN=2048
 MIN_CACHE_HIT_PCT=60
 MAX_LATENCY_ALLOWED_MS=500
 ```
@@ -4,13 +4,15 @@
 # See details in README (benchmarks/auto_tune/README.md).

 TAG=$(date +"%Y_%m_%d_%H_%M")
-BASE=""
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+BASE="$SCRIPT_DIR/../../.."
 MODEL="meta-llama/Llama-3.1-8B-Instruct"
 SYSTEM="TPU"
 TP=1
 DOWNLOAD_DIR=""
 INPUT_LEN=4000
 OUTPUT_LEN=16
+MAX_MODEL_LEN=4096
 MIN_CACHE_HIT_PCT=0
 MAX_LATENCY_ALLOWED_MS=100000000000
 NUM_SEQS_LIST="128 256"
@@ -36,6 +38,13 @@ current_hash=$(git rev-parse HEAD)
 echo "hash:$current_hash" >> "$RESULT"
 echo "current_hash: $current_hash"

+TOTAL_LEN=$((INPUT_LEN + OUTPUT_LEN))
+RED='\033[0;31m'
+if (( TOTAL_LEN > MAX_MODEL_LEN )); then
+    echo -e "${RED}FAILED: INPUT_LEN($INPUT_LEN) + OUTPUT_LEN($OUTPUT_LEN) = $TOTAL_LEN, which is > MAX_MODEL_LEN = $MAX_MODEL_LEN.\033[0m" >&2
+    exit 1
+fi
+
 best_throughput=0
 best_max_num_seqs=0
 best_num_batched_tokens=0
@@ -60,7 +69,7 @@ start_server() {
        --enable-prefix-caching \
        --load-format dummy \
        --download-dir "$DOWNLOAD_DIR" \
-        --max-model-len $(( INPUT_LEN+OUTPUT_LEN )) > "$vllm_log" 2>&1 &
+        --max-model-len $MAX_MODEL_LEN > "$vllm_log" 2>&1 &

    # wait for 10 minutes...
    server_started=0
@@ -245,4 +254,3 @@ done
 echo "finish permutations"
 echo "best_max_num_seqs: $best_max_num_seqs, best_num_batched_tokens: $best_num_batched_tokens, best_throughput: $best_throughput, profile saved in: $PROFILE_PATH"
 echo "best_max_num_seqs: $best_max_num_seqs, best_num_batched_tokens: $best_num_batched_tokens, best_throughput: $best_throughput, profile saved in: $PROFILE_PATH" >> "$RESULT"
-
@@ -167,7 +167,8 @@ async def run_vllm_async(
    from vllm import SamplingParams

    async with build_async_engine_client_from_engine_args(
-        engine_args, disable_frontend_multiprocessing
+        engine_args,
+        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
    ) as llm:
        model_config = await llm.get_model_config()
        assert all(
@@ -58,6 +58,22 @@ function (find_isa CPUINFO TARGET OUT)
    endif()
 endfunction()

+
+function(check_sysctl TARGET OUT)
+    execute_process(COMMAND sysctl -n "${TARGET}"
+                    RESULT_VARIABLE SYSCTL_RET
+                    OUTPUT_VARIABLE SYSCTL_INFO
+                    ERROR_QUIET
+                    OUTPUT_STRIP_TRAILING_WHITESPACE)
+    if(SYSCTL_RET EQUAL 0 AND
+      (SYSCTL_INFO STREQUAL "1" OR SYSCTL_INFO GREATER 0))
+        set(${OUT} ON PARENT_SCOPE)
+    else()
+        set(${OUT} OFF PARENT_SCOPE)
+    endif()
+endfunction()
+
+
 function (is_avx512_disabled OUT)
    set(DISABLE_AVX512 $ENV{VLLM_CPU_DISABLE_AVX512})
    if(DISABLE_AVX512 AND DISABLE_AVX512 STREQUAL "true")
@@ -70,7 +86,10 @@ endfunction()
 is_avx512_disabled(AVX512_DISABLED)

 if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
-    set(APPLE_SILICON_FOUND TRUE)
+    message(STATUS "Apple Silicon Detected")
+    set(ENABLE_NUMA OFF)
+    check_sysctl(hw.optional.neon ASIMD_FOUND)
+    check_sysctl(hw.optional.arm.FEAT_BF16 ARM_BF16_FOUND)
 else()
    find_isa(${CPUINFO} "avx2" AVX2_FOUND)
    find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
@@ -82,7 +101,6 @@ else()
    find_isa(${CPUINFO} "S390" S390_FOUND)
 endif()

-
 if (AVX512_FOUND AND NOT AVX512_DISABLED)
    list(APPEND CXX_COMPILE_FLAGS
        "-mavx512f"
@@ -149,9 +167,6 @@ elseif (ASIMD_FOUND)
        set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16")  
    endif()
    list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})     
-elseif(APPLE_SILICON_FOUND)
-    message(STATUS "Apple Silicon Detected")
-    set(ENABLE_NUMA OFF)
 elseif (S390_FOUND)
    message(STATUS "S390 detected")
    # Check for S390 VXE support
@@ -24,7 +24,7 @@

 #include "attention_dtypes.h"
 #include "attention_utils.cuh"
-#include "cuda_compat.h"
+#include "../cuda_compat.h"

 #ifdef USE_ROCM
  #include <hip/hip_bf16.h>
@@ -16,9 +16,8 @@
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
-
 #include "attention_kernels.cuh"
-#include "cuda_compat.h"
+#include "../cuda_compat.h"

 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -75,7 +74,7 @@ void paged_attention_v1_launcher(
  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());

-  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
+  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
  int padded_max_seq_len =
      DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
  int logits_size = padded_max_seq_len * sizeof(float);
@@ -16,9 +16,8 @@
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
-
 #include "attention_kernels.cuh"
-#include "cuda_compat.h"
+#include "../cuda_compat.h"

 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -79,7 +78,7 @@ void paged_attention_v2_launcher(
  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());

-  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
+  const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
  int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
  int logits_size = PARTITION_SIZE * sizeof(float);
  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
@@ -4,8 +4,35 @@
  #include <hip/hip_runtime.h>
 #endif

-#if defined(USE_ROCM) && defined(__GFX9__)
-  #define WARP_SIZE 64
+#ifdef USE_ROCM
+struct Utils {
+  static __host__ int get_warp_size() {
+    static bool is_cached = false;
+    static int result;
+
+    if (!is_cached) {
+      int device_id;
+      cudaDeviceProp deviceProp;
+      cudaGetDevice(&device_id);
+      cudaGetDeviceProperties(&deviceProp, device_id);
+
+      result = deviceProp.warpSize;
+      is_cached = true;
+    }
+
+    return result;
+  }
+
+  static __device__ constexpr int get_warp_size() {
+  #ifdef __GFX9__
+    return 64;
+  #else
+    return 32;
+  #endif
+  }
+};
+
+  #define WARP_SIZE Utils::get_warp_size()
 #else
  #define WARP_SIZE 32
 #endif
@@ -190,8 +190,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
  2) This implementation assumes k is small, but will work for any k.
 */

-template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, typename IndType>
-__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
+template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
    void topkGatingSoftmax(const float* input, const bool* finished, float* output, const int num_rows, IndType* indices,
        int* source_rows, const int k, const int start_expert, const int end_expert)
 {
@@ -209,12 +209,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__

    // Restrictions based on previous section.
    static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
-    static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
+    static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
    static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2");
-    static_assert(THREADS_PER_ROW <= WARP_SIZE, "THREADS_PER_ROW can be at most warp size");
+    static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM, "THREADS_PER_ROW can be at most warp size");

    // We have NUM_EXPERTS elements per row. We specialize for small #experts
-    static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+    static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
    static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
    static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;

@@ -393,41 +393,51 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
 namespace detail
 {
 // Constructs some constants needed to partition the work across threads at compile time.
-template <int EXPERTS, int BYTES_PER_LDG>
+template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM>
 struct TopkConstants
 {
    static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(float);
-    static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, "");
-    static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
+    static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0, "");
+    static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
    static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
    static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
-    static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
+    static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
 };
 } // namespace detail

-template <int EXPERTS, int WARPS_PER_TB, typename IndType>
+template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, typename IndType>
 void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, IndType* indices,
    int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, cudaStream_t stream)
 {
    static constexpr std::size_t MAX_BYTES_PER_LDG = 16;

    static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(float) * EXPERTS);
-    using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG>;
+    using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM>;
    static constexpr int VPT = Constants::VPT;
    static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
    const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
    const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;

-    dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
-    topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG><<<num_blocks, block_dim, 0, stream>>>(
+    dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
+    topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM><<<num_blocks, block_dim, 0, stream>>>(
        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert);
 }

-#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB)                       \
-    topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB>(         \
-        gating_output, nullptr, topk_weights, topk_indices,            \
-        token_expert_indices, num_tokens, topk, 0, num_experts,         \
-        stream);
+#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB)                                \
+    switch (warpSize) {                                                          \
+        case 32:                                                                 \
+            topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32>(      \
+                gating_output, nullptr, topk_weights, topk_indices,              \
+                token_expert_indices, num_tokens, topk, 0, num_experts, stream); \
+            break;                                                               \
+        case 64:                                                                 \
+            topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64>(      \
+                gating_output, nullptr, topk_weights, topk_indices,              \
+                token_expert_indices, num_tokens, topk, 0, num_experts, stream); \
+            break;                                                               \
+        default:                                                                 \
+            TORCH_CHECK(false, "Unsupported warp size: ", warpSize);             \
+    }

 template <typename IndType>
 void topkGatingSoftmaxKernelLauncher(
@@ -441,6 +451,7 @@ void topkGatingSoftmaxKernelLauncher(
    const int topk,
    cudaStream_t stream) {
    static constexpr int WARPS_PER_TB = 4;
+    auto warpSize = WARP_SIZE;
    switch (num_experts) {
        case 1:
            LAUNCH_SOFTMAX(1, WARPS_PER_TB);
@@ -4,7 +4,7 @@

 #include <cmath>
 #include "core/math.hpp"
-#include "cuda_compat.h"
+#include "../cuda_compat.h"
 #include "dispatch_utils.h"

 #include "quantization/fp8/common.cuh"
@@ -47,13 +47,12 @@ __global__ void compute_problem_sizes(const int32_t* __restrict__ topk_ids,

 __global__ void compute_expert_offsets(
    const int32_t* __restrict__ problem_sizes1, int32_t* expert_offsets,
-    int32_t* atomic_buffer, const int num_experts, const int topk_length) {
+    int32_t* atomic_buffer, const int num_experts, const bool swap_ab) {
  int32_t tot_offset = 0;
  expert_offsets[0] = 0;
  for (int i = 0; i < num_experts; ++i) {
    atomic_buffer[i] = tot_offset;
-    tot_offset += topk_length > SWAP_AB_THRESHOLD ? problem_sizes1[i * 3]
-                                                  : problem_sizes1[i * 3 + 1];
+    tot_offset += swap_ab ? problem_sizes1[i * 3 + 1] : problem_sizes1[i * 3];
    expert_offsets[i + 1] = tot_offset;
  }
 }
@@ -61,15 +60,14 @@ __global__ void compute_expert_offsets(
 __global__ void compute_expert_blockscale_offsets(
    const int32_t* __restrict__ problem_sizes1, int32_t* expert_offsets,
    int32_t* blockscale_offsets, int32_t* atomic_buffer, const int num_experts,
-    const int topk_length) {
+    const bool swap_ab) {
  int32_t tot_offset = 0;
  int32_t tot_offset_round = 0;
  expert_offsets[0] = 0;
  blockscale_offsets[0] = 0;
  for (int i = 0; i < num_experts; ++i) {
-    int32_t cur_offset = topk_length > SWAP_AB_THRESHOLD
-                             ? problem_sizes1[i * 3]
-                             : problem_sizes1[i * 3 + 1];
+    int32_t cur_offset =
+        swap_ab ? problem_sizes1[i * 3 + 1] : problem_sizes1[i * 3];
    atomic_buffer[i] = tot_offset;
    tot_offset += cur_offset;
    expert_offsets[i + 1] = tot_offset;
@@ -119,15 +117,19 @@ void get_cutlass_moe_mm_data_caller(

  int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());

-  if (topk_ids.numel() > SWAP_AB_THRESHOLD) {
-    compute_problem_sizes<false><<<num_experts, num_threads, 0, stream>>>(
+  // Swap-AB should be disabled for FP4 path
+  bool may_swap_ab = (!blockscale_offsets.has_value()) &&
+                     (topk_ids.numel() <= SWAP_AB_THRESHOLD);
+
+  if (may_swap_ab) {
+    compute_problem_sizes<true><<<num_experts, num_threads, 0, stream>>>(
        static_cast<const int32_t*>(topk_ids.data_ptr()),
        static_cast<int32_t*>(problem_sizes1.data_ptr()),
        static_cast<int32_t*>(problem_sizes2.data_ptr()),
        static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n,
        k);
  } else {
-    compute_problem_sizes<true><<<num_experts, num_threads, 0, stream>>>(
+    compute_problem_sizes<false><<<num_experts, num_threads, 0, stream>>>(
        static_cast<const int32_t*>(topk_ids.data_ptr()),
        static_cast<int32_t*>(problem_sizes1.data_ptr()),
        static_cast<int32_t*>(problem_sizes2.data_ptr()),
@@ -136,18 +138,19 @@ void get_cutlass_moe_mm_data_caller(
  }

  if (blockscale_offsets.has_value()) {
+    // fp4 path
    compute_expert_blockscale_offsets<<<1, 1, 0, stream>>>(
        static_cast<const int32_t*>(problem_sizes1.data_ptr()),
        static_cast<int32_t*>(expert_offsets.data_ptr()),
        static_cast<int32_t*>(blockscale_offsets.value().data_ptr()),
        static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts,
-        topk_ids.numel());
+        may_swap_ab);
  } else {
    compute_expert_offsets<<<1, 1, 0, stream>>>(
        static_cast<const int32_t*>(problem_sizes1.data_ptr()),
        static_cast<int32_t*>(expert_offsets.data_ptr()),
        static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts,
-        topk_ids.numel());
+        may_swap_ab);
  }
  compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
      static_cast<const int32_t*>(topk_ids.data_ptr()),
@@ -4,7 +4,7 @@
 #include <torch/all.h>
 #include <c10/cuda/CUDAGuard.h>

-#include "cuda_compat.h"
+#include "../../cuda_compat.h"
 #include "dispatch_utils.h"

 #include "ggml-common.h"
@@ -19,7 +19,7 @@
 #include <c10/cuda/CUDAGuard.h>
 #include <hip/hip_fp8.h>
 #include <hip/hip_bf16.h>
-#include "cuda_compat.h"
+#include "../cuda_compat.h"

 #include <algorithm>
 #include "../attention/dtype_fp8.cuh"
@@ -9,7 +9,7 @@
 #include <stdexcept>
 #include <algorithm>

-#include "cuda_compat.h"
+#include "../cuda_compat.h"
 #include "dispatch_utils.h"
 #include "quantization/fp8/common.cuh"

@@ -265,7 +265,7 @@ RUN if [ "$RUN_WHEEL_CHECK" = "true" ]; then \
 #################### EXTENSION Build IMAGE ####################

 #################### DEV IMAGE ####################
-FROM base as dev
+FROM base AS dev

 ARG PIP_INDEX_URL UV_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
@@ -276,10 +276,6 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"

-# Workaround for #17068
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
-
 COPY requirements/lint.txt requirements/lint.txt
 COPY requirements/test.txt requirements/test.txt
 COPY requirements/dev.txt requirements/dev.txt
@@ -390,7 +386,7 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist

 # Install FlashInfer from source
 ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
-ARG FLASHINFER_GIT_REF="v0.2.8rc1"
+ARG FLASHINFER_GIT_REF="v0.2.9rc1"
 RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
  . /etc/environment
    git clone --depth 1 --recursive --shallow-submodules \
@@ -452,10 +448,6 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
 ENV UV_HTTP_TIMEOUT=500
 ENV UV_INDEX_STRATEGY="unsafe-best-match"

-# Workaround for #17068
-RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install --system --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
-
 # install development dependencies (for testing)
 RUN --mount=type=cache,target=/root/.cache/uv \
    CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
@@ -47,7 +47,7 @@ FROM vllm-base AS vllm-openai

 # install additional dependencies for openai api server
 RUN --mount=type=cache,target=/root/.cache/pip \
-    pip install accelerate hf_transfer pytest modelscope
+    pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope

 ENV VLLM_USAGE_SOURCE production-docker-image \
    TRITON_XPU_PROFILE 1
@@ -14,7 +14,6 @@ API documentation for vLLM's configuration classes.
 - [vllm.config.DeviceConfig][]
 - [vllm.config.SpeculativeConfig][]
 - [vllm.config.LoRAConfig][]
- [vllm.config.PromptAdapterConfig][]
 - [vllm.config.MultiModalConfig][]
 - [vllm.config.PoolerConfig][]
 - [vllm.config.DecodingConfig][]
@@ -98,7 +98,7 @@ For additional features and advanced configurations, refer to the official [MkDo
 ??? console "Commands"

    ```bash
-    pip install -r requirements/dev.txt
+    pip install -r requirements/common.txt -r requirements/dev.txt

    # Linting, formatting and static type checking
    pre-commit install --hook-type pre-commit --hook-type commit-msg
@@ -134,7 +134,7 @@ MAX_JOBS=16 uv pip install --system \

 ```bash
 uv pip install --system \
-    --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.4"
+    --no-build-isolation "git+https://github.com/state-spaces/mamba@v2.2.5"
 ```

 ### causal-conv1d
@@ -34,23 +34,22 @@ th:not(:first-child) {
 }
 </style>

-| Feature | [CP][chunked-prefill] | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | <abbr title="Prompt Adapter">prmpt adptr</abbr> | [SD](spec_decode.md) | CUDA graph | <abbr title="Pooling Models">pooling</abbr> | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search |
-|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
+| Feature | [CP][chunked-prefill] | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](spec_decode.md) | CUDA graph | <abbr title="Pooling Models">pooling</abbr> | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search |
+|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
 | [CP][chunked-prefill] | ✅ | | | | | | | | | | | | | | |
 | [APC](automatic_prefix_caching.md) | ✅ | ✅ | | | | | | | | | | | | | |
 | [LoRA](lora.md) | ✅ | ✅ | ✅ | | | | | | | | | | | | |
-| <abbr title="Prompt Adapter">prmpt adptr</abbr> | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | | |
-| [SD](spec_decode.md) | ✅ | ✅ | ❌ | ✅ | ✅ | | | | | | | | | | |
-| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
-| <abbr title="Pooling Models">pooling</abbr> | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | | | | | | | | |
-| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](gh-issue:7366) | ❌ | ❌ | [❌](gh-issue:7366) | ✅ | ✅ | ✅ | | | | | | | |
-| <abbr title="Logprobs">logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | | | | | | |
-| <abbr title="Prompt Logprobs">prmpt logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | | | | | |
-| <abbr title="Async Output Processing">async output</abbr> | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | | | | |
-| multi-step | ❌ | ✅ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | | | |
-| <abbr title="Multimodal Inputs">mm</abbr> | ✅ | [🟠](gh-pr:8348) | [🟠](gh-pr:4194) | ❔ | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | |
-| best-of | ✅ | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ✅ | ✅ | |
-| beam-search | ✅ | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ❔ | ✅ | ✅ |
+| [SD](spec_decode.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | |
+| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
+| <abbr title="Pooling Models">pooling</abbr> | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | | | | | | | | |
+| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](gh-issue:7366) | ❌ | [❌](gh-issue:7366) | ✅ | ✅ | ✅ | | | | | | | |
+| <abbr title="Logprobs">logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | | | | | | |
+| <abbr title="Prompt Logprobs">prmpt logP</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | | | | | |
+| <abbr title="Async Output Processing">async output</abbr> | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | | | | |
+| multi-step | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | | | |
+| <abbr title="Multimodal Inputs">mm</abbr> | ✅ | [🟠](gh-pr:8348) | [🟠](gh-pr:4194) | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | |
+| best-of | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ✅ | ✅ | |
+| beam-search | ✅ | ✅ | ✅ | [❌](gh-issue:6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](gh-issue:7968) | ❔ | ✅ | ✅ |

 [](){ #feature-x-hardware }

@@ -59,10 +58,9 @@ th:not(:first-child) {
 | Feature                                                   | Volta               | Turing    | Ampere    | Ada    | Hopper     | CPU                | AMD    | TPU |
 |-----------------------------------------------------------|---------------------|-----------|-----------|--------|------------|--------------------|--------|-----|
 | [CP][chunked-prefill]                                     | [❌](gh-issue:2729) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅ |
-| [APC](automatic_prefix_caching.md)                           | [❌](gh-issue:3687) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅ |
-| [LoRA](lora.md)                                      | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅ |
-| <abbr title="Prompt Adapter">prmpt adptr</abbr>           | ✅                  | ✅        | ✅        | ✅     | ✅        | [❌](gh-issue:8475) | ✅     | ❌ |
-| [SD](spec_decode.md)                                         | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
+| [APC](automatic_prefix_caching.md)                        | [❌](gh-issue:3687) | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅ |
+| [LoRA](lora.md)                                           | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ✅ |
+| [SD](spec_decode.md)                                      | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ✅     | ❌ |
 | CUDA graph                                                | ✅                  | ✅        | ✅        | ✅     | ✅        | ❌                  | ✅     | ❌ |
 | <abbr title="Pooling Models">pooling</abbr>               | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ❔     | ❌ |
 | <abbr title="Encoder-Decoder Models">enc-dec</abbr>       | ✅                  | ✅        | ✅        | ✅     | ✅        | ✅                  | ❌     | ❌ |
@@ -103,7 +103,8 @@ When tool_choice='required' is set, the model is guaranteed to generate one or m

 vLLM supports the `tool_choice='none'` option in the chat completion API. When this option is set, the model will not generate any tool calls and will respond with regular text content only, even if tools are defined in the request.

-However, when `tool_choice='none'` is specified, vLLM includes tool definitions from the prompt.
+!!! note
+    When tools are specified in the request, vLLM includes tool definitions in the prompt by default, regardless of the `tool_choice` setting. To exclude tool definitions when `tool_choice='none'`, use the `--exclude-tools-when-tool-choice-none` option.

 ## Automatic Function Calling

@@ -5,9 +5,98 @@ vLLM model tensors that have been serialized to disk, an HTTP/HTTPS endpoint, or
 at runtime extremely quickly directly to the GPU, resulting in significantly
 shorter Pod startup times and CPU memory usage. Tensor encryption is also supported.

-For more information on CoreWeave's Tensorizer, please refer to
-[CoreWeave's Tensorizer documentation](https://github.com/coreweave/tensorizer). For more information on serializing a vLLM model, as well a general usage guide to using Tensorizer with vLLM, see
-the [vLLM example script](../../examples/others/tensorize_vllm_model.md).
+vLLM fully integrates Tensorizer in to its model loading machinery. The following will give a brief overview on how to get started with using Tensorizer on vLLM.

-!!! note
-    Note that to use this feature you will need to install `tensorizer` by running `pip install vllm[tensorizer]`.
+## Installing Tensorizer
+
+To install `tensorizer`, run `pip install vllm[tensorizer]`.
+
+## The basics
+
+To load a model using Tensorizer, the model first needs to be serialized by
+Tensorizer. [The example script](../../examples/others/tensorize_vllm_model.md) takes care of this process.
+
+Let's walk through a basic example by serializing `facebook/opt-125m` using the script, and then loading it for inference.
+
+## Serializing a vLLM model with Tensorizer
+
+To serialize a model with Tensorizer, call the example script with the necessary
+CLI arguments. The docstring for the script itself explains the CLI args
+and how to use it properly in great detail, and we'll use one of the examples from the docstring directly, assuming we want to serialize and save our model at our S3 bucket example `s3://my-bucket`:
+
+```bash
+python examples/others/tensorize_vllm_model.py \
+   --model facebook/opt-125m \
+   serialize \
+   --serialized-directory s3://my-bucket \
+   --suffix v1
+```
+
+This saves the model tensors at `s3://my-bucket/vllm/facebook/opt-125m/v1`. If you intend on applying a LoRA adapter to your tensorized model, you can pass the HF id of the LoRA adapter in the above command, and the artifacts will be saved there too:
+
+```bash
+python examples/others/tensorize_vllm_model.py \
+   --model facebook/opt-125m \
+   --lora-path <lora_id> \
+   serialize \
+   --serialized-directory s3://my-bucket \
+   --suffix v1
+```
+
+## Serving the model using Tensorizer
+
+Once the model is serialized where you want it, you can load the model using `vllm serve` or the `LLM` entrypoint. You can pass the directory where you saved the model to the `model` argument for `LLM()` and `vllm serve`. For example, to serve the tensorized model saved previously with the LoRA adapter, you'd do:
+
+```bash
+vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
+    --load-format tensorizer \
+    --enable-lora 
+```
+
+Or, with `LLM()`:
+
+```python
+from vllm import LLM
+llm = LLM(
+    "s3://my-bucket/vllm/facebook/opt-125m/v1", 
+    load_format="tensorizer",
+    enable_lora=True
+)
+```
+
+## Options for configuring Tensorizer
+
+`tensorizer`'s core objects that serialize and deserialize models are `TensorSerializer` and `TensorDeserializer` respectively. In order to pass arbitrary kwargs to these, which will configure the serialization and deserialization processes, you can provide them as keys to `model_loader_extra_config` with `serialization_kwargs` and `deserialization_kwargs` respectively. Full docstrings detailing all parameters for the aforementioned objects can be found in `tensorizer`'s [serialization.py](https://github.com/coreweave/tensorizer/blob/main/tensorizer/serialization.py) file.
+
+As an example, CPU concurrency can be limited when serializing with `tensorizer` via the `limit_cpu_concurrency` parameter in the initializer for `TensorSerializer`. To set `limit_cpu_concurrency` to some arbitrary value, you would do so like this when serializing:
+
+```bash
+python examples/others/tensorize_vllm_model.py \
+   --model facebook/opt-125m \
+   --lora-path <lora_id> \
+   serialize \
+   --serialized-directory s3://my-bucket \
+   --serialization-kwargs '{"limit_cpu_concurrency": 2}' \
+   --suffix v1
+```
+
+As an example when customizing the loading process via `TensorDeserializer`, you could limit the number of concurrency readers during deserialization with the `num_readers` parameter in the initializer via `model_loader_extra_config` like so:
+
+```bash
+vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
+    --load-format tensorizer \
+    --enable-lora \
+    --model-loader-extra-config '{"deserialization_kwargs": {"num_readers": 2}}'
+```
+
+Or with `LLM()`:
+
+```python
+from vllm import LLM
+llm = LLM(
+    "s3://my-bucket/vllm/facebook/opt-125m/v1", 
+    load_format="tensorizer",
+    enable_lora=True,
+    model_loader_extra_config={"deserialization_kwargs": {"num_readers": 2}}
+)
+```
@@ -623,6 +623,12 @@ Specified using `--task generate`.
 | `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ | ✅︎ |
 | `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ | ✅︎ |

+Some models are supported only via the [Transformers backend](#transformers). The purpose of the table below is to acknowledge models which we officially support in this way. The logs will say that the Transformers backend is being used, and you will see no warning that this is fallback behaviour. This means that, if you have issues with any of the models listed below, please [make an issue](https://github.com/vllm-project/vllm/issues/new/choose) and we'll do our best to fix it!
+
+| Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/distributed_serving.md) | [V1](gh-issue:8779) |
+|--------------|--------|--------|-------------------|-----------------------------|-----------------------------------------|---------------------|
+| `Emu3ForConditionalGeneration` | Emu3 | T + I | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ | ✅︎ |
+
 <sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.  
 &nbsp;&nbsp;&nbsp;&nbsp;• For example, to use DeepSeek-VL2 series models:  
 &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;`--hf-overrides '{"architectures": ["DeepseekVLV2ForCausalLM"]}'`  
@@ -1,31 +1,38 @@
-# Distributed Inference and Serving
+# Distributed inference and serving

-## How to decide the distributed inference strategy?
+## Distributed inference strategies for a single-model replica

-Before going into the details of distributed inference and serving, let's first make it clear when to use distributed inference and what are the strategies available. The common practice is:
+To choose a distributed inference strategy for a single-model replica, use the following guidelines:

- **Single GPU (no distributed inference)**: If your model fits in a single GPU, you probably don't need to use distributed inference. Just use the single GPU to run the inference.
- **Single-Node Multi-GPU (tensor parallel inference)**: If your model is too large to fit in a single GPU, but it can fit in a single node with multiple GPUs, you can use tensor parallelism. The tensor parallel size is the number of GPUs you want to use. For example, if you have 4 GPUs in a single node, you can set the tensor parallel size to 4.
- **Multi-Node Multi-GPU (tensor parallel plus pipeline parallel inference)**: If your model is too large to fit in a single node, you can use tensor parallel together with pipeline parallelism. The tensor parallel size is the number of GPUs you want to use in each node, and the pipeline parallel size is the number of nodes you want to use. For example, if you have 16 GPUs in 2 nodes (8 GPUs per node), you can set the tensor parallel size to 8 and the pipeline parallel size to 2.
+- **Single GPU (no distributed inference):** if the model fits on a single GPU, distributed inference is probably unnecessary. Run inference on that GPU.
+- **Single-node multi-GPU using tensor parallel inference:** if the model is too large for a single GPU but fits on a single node with multiple GPUs, use *tensor parallelism*. For example, set `tensor_parallel_size=4` when using a node with 4 GPUs.
+- **Multi-node multi-GPU using tensor parallel and pipeline parallel inference:** if the model is too large for a single node, combine *tensor parallelism* with *pipeline parallelism*. Set `tensor_parallel_size` to the number of GPUs per node and `pipeline_parallel_size` to the number of nodes. For example, set `tensor_parallel_size=8` and `pipeline_parallel_size=2` when using 2 nodes with 8 GPUs per node.

-In short, you should increase the number of GPUs and the number of nodes until you have enough GPU memory to hold the model. The tensor parallel size should be the number of GPUs in each node, and the pipeline parallel size should be the number of nodes.
+Increase the number of GPUs and nodes until there is enough GPU memory for the model. Set `tensor_parallel_size` to the number of GPUs per node and `pipeline_parallel_size` to the number of nodes.

-After adding enough GPUs and nodes to hold the model, you can run vLLM first, which will print some logs like `# GPU blocks: 790`. Multiply the number by `16` (the block size), and you can get roughly the maximum number of tokens that can be served on the current configuration. If this number is not satisfying, e.g. you want higher throughput, you can further increase the number of GPUs or nodes, until the number of blocks is enough.
+After you provision sufficient resources to fit the model, run `vllm`. Look for log messages like:

-!!! note
-    There is one edge case: if the model fits in a single node with multiple GPUs, but the number of GPUs cannot divide the model size evenly, you can use pipeline parallelism, which splits the model along layers and supports uneven splits. In this case, the tensor parallel size should be 1 and the pipeline parallel size should be the number of GPUs.
+```text
+INFO 07-23 13:56:04 [kv_cache_utils.py:775] GPU KV cache size: 643,232 tokens
+INFO 07-23 13:56:04 [kv_cache_utils.py:779] Maximum concurrency for 40,960 tokens per request: 15.70x
+```

-### Distributed serving of MoE (Mixture of Experts) models
+The `GPU KV cache size` line reports the total number of tokens that can be stored in the GPU KV cache at once. The `Maximum concurrency` line provides an estimate of how many requests can be served concurrently if each request requires the specified number of tokens (40,960 in the example above). The tokens-per-request number is taken from the model configuration's maximum sequence length, `ModelConfig.max_model_len`. If these numbers are lower than your throughput requirements, add more GPUs or nodes to your cluster.

-It is often advantageous to exploit the inherent parallelism of experts by using a separate parallelism strategy for the expert layers. vLLM supports large-scale deployment combining Data Parallel attention with Expert or Tensor Parallel MoE layers. See the page on [Data Parallel Deployment](data_parallel_deployment.md) for more information.
+!!! note "Edge case: uneven GPU splits"
+    If the model fits within a single node but the GPU count doesn't evenly divide the model size, enable pipeline parallelism, which splits the model along layers and supports uneven splits. In this scenario, set `tensor_parallel_size=1` and `pipeline_parallel_size` to the number of GPUs. Furthermore, if the GPUs on the node do not have NVLINK interconnect (e.g. L40S), leverage pipeline parallelism instead of tensor parallelism for higher throughput and lower communication overhead.

-## Running vLLM on a single node
+### Distributed serving of *Mixture of Experts* (*MoE*) models

-vLLM supports distributed tensor-parallel and pipeline-parallel inference and serving. Currently, we support [Megatron-LM's tensor parallel algorithm](https://arxiv.org/pdf/1909.08053.pdf). We manage the distributed runtime with either [Ray](https://github.com/ray-project/ray) or python native multiprocessing. Multiprocessing can be used when deploying on a single node, multi-node inference currently requires Ray.
+It's often advantageous to exploit the inherent parallelism of experts by using a separate parallelism strategy for the expert layers. vLLM supports large-scale deployment combining Data Parallel attention with Expert or Tensor Parallel MoE layers. For more information, see [Data Parallel Deployment](data_parallel_deployment.md).

-Multiprocessing will be used by default when not running in a Ray placement group and if there are sufficient GPUs available on the same node for the configured `tensor_parallel_size`, otherwise Ray will be used. This default can be overridden via the `LLM` class `distributed_executor_backend` argument or `--distributed-executor-backend` API server argument. Set it to `mp` for multiprocessing or `ray` for Ray. It's not required for Ray to be installed for the multiprocessing case.
+## Single-node deployment

-To run multi-GPU inference with the `LLM` class, set the `tensor_parallel_size` argument to the number of GPUs you want to use. For example, to run inference on 4 GPUs:
+vLLM supports distributed tensor-parallel and pipeline-parallel inference and serving. The implementation includes [Megatron-LM's tensor parallel algorithm](https://arxiv.org/pdf/1909.08053.pdf).
+
+The default distributed runtimes are [Ray](https://github.com/ray-project/ray) for multi-node inference and native Python `multiprocessing` for single-node inference. You can override the defaults by setting `distributed_executor_backend` in the `LLM` class or `--distributed-executor-backend` in the API server. Use `mp` for `multiprocessing` or `ray` for Ray.
+
+For multi-GPU inference, set `tensor_parallel_size` in the `LLM` class to the desired GPU count. For example, to run inference on 4 GPUs:

 ```python
 from vllm import LLM
@@ -33,84 +40,96 @@ llm = LLM("facebook/opt-13b", tensor_parallel_size=4)
 output = llm.generate("San Francisco is a")
 ```

-To run multi-GPU serving, pass in the `--tensor-parallel-size` argument when starting the server. For example, to run API server on 4 GPUs:
+For multi-GPU serving, include `--tensor-parallel-size` when starting the server. For example, to run the API server on 4 GPUs:

 ```bash
 vllm serve facebook/opt-13b \
     --tensor-parallel-size 4
 ```

-You can also additionally specify `--pipeline-parallel-size` to enable pipeline parallelism. For example, to run API server on 8 GPUs with pipeline parallelism and tensor parallelism:
+To enable pipeline parallelism, add `--pipeline-parallel-size`. For example, to run the API server on 8 GPUs with pipeline parallelism and tensor parallelism:

 ```bash
+# Eight GPUs total
 vllm serve gpt2 \
     --tensor-parallel-size 4 \
     --pipeline-parallel-size 2
 ```

-## Running vLLM on multiple nodes
+## Multi-node deployment

-If a single node does not have enough GPUs to hold the model, you can run the model using multiple nodes. It is important to make sure the execution environment is the same on all nodes, including the model path, the Python environment. The recommended way is to use docker images to ensure the same environment, and hide the heterogeneity of the host machines via mapping them into the same docker configuration.
+If a single node lacks sufficient GPUs to hold the model, deploy vLLM across multiple nodes. Multi-node deployments require Ray as the runtime engine. Ensure that every node provides an identical execution environment, including the model path and Python packages. Using container images is recommended because they provide a convenient way to keep environments consistent and to hide host heterogeneity.

-The first step, is to start containers and organize them into a cluster. We have provided the helper script <gh-file:examples/online_serving/run_cluster.sh> to start the cluster. Please note, this script launches docker without administrative privileges that would be required to access GPU performance counters when running profiling and tracing tools. For that purpose, the script can have `CAP_SYS_ADMIN` to the docker container by using the `--cap-add` option in the docker run command.
+### Ray cluster setup with containers

-Pick a node as the head node, and run the following command:
+The helper script <gh-file:examples/online_serving/run_cluster.sh> starts containers across nodes and initializes Ray. By default, the script runs Docker without administrative privileges, which prevents access to the GPU performance counters when profiling or tracing. To enable admin privileges, add the `--cap-add=CAP_SYS_ADMIN` flag to the Docker command.
+
+Choose one node as the head node and run:

 ```bash
 bash run_cluster.sh \
                vllm/vllm-openai \
-                ip_of_head_node \
+                <HEAD_NODE_IP> \
                --head \
                /path/to/the/huggingface/home/in/this/node \
-                -e VLLM_HOST_IP=ip_of_this_node
+                -e VLLM_HOST_IP=<HEAD_NODE_IP>
 ```

-On the rest of the worker nodes, run the following command:
+On each worker node, run:

 ```bash
 bash run_cluster.sh \
                vllm/vllm-openai \
-                ip_of_head_node \
+                <HEAD_NODE_IP> \
                --worker \
                /path/to/the/huggingface/home/in/this/node \
-                -e VLLM_HOST_IP=ip_of_this_node
+                -e VLLM_HOST_IP=<WORKER_NODE_IP>
 ```

-Then you get a ray cluster of **containers**. Note that you need to keep the shells running these commands alive to hold the cluster. Any shell disconnect will terminate the cluster. In addition, please note that the argument `ip_of_head_node` should be the IP address of the head node, which is accessible by all the worker nodes. The IP addresses of each worker node should be specified in the `VLLM_HOST_IP` environment variable, and should be different for each worker node. Please check the network configuration of your cluster to make sure the nodes can communicate with each other through the specified IP addresses.
+Note that `VLLM_HOST_IP` is unique for each worker. Keep the shells running these commands open; closing any shell terminates the cluster. Ensure that all nodes can communicate with each other through their IP addresses.

-!!! warning
-    It is considered best practice to set `VLLM_HOST_IP` to an address on a private network segment for the vLLM cluster. The traffic sent here is not encrypted. The endpoints are also exchanging data in a format that could be exploited to execute arbitrary code should a malicious party gain access to the network. Please ensure that this network is not reachable by any untrusted parties.
+!!! warning "Network security"
+    For security, set `VLLM_HOST_IP` to an address on a private network segment. Traffic sent over this network is unencrypted, and the endpoints exchange data in a format that can be exploited to execute arbitrary code if an adversary gains network access. Ensure that untrusted parties cannot reach the network.

-!!! warning
-    Since this is a ray cluster of **containers**, all the following commands should be executed in the **containers**, otherwise you are executing the commands on the host machine, which is not connected to the ray cluster. To enter the container, you can use `docker exec -it node /bin/bash`.
+From any node, enter a container and run `ray status` and `ray list nodes` to verify that Ray finds the expected number of nodes and GPUs.

-Then, on any node, use `docker exec -it node /bin/bash` to enter the container, execute `ray status` and `ray list nodes` to check the status of the Ray cluster. You should see the right number of nodes and GPUs.
+!!! tip
+    Alternatively, set up the Ray cluster using KubeRay. For more information, see [KubeRay vLLM documentation](https://docs.ray.io/en/latest/cluster/kubernetes/examples/vllm-rayservice.html).

-After that, on any node, use `docker exec -it node /bin/bash` to enter the container again. **In the container**, you can use vLLM as usual, just as you have all the GPUs on one node: vLLM will be able to leverage GPU resources of all nodes in the Ray cluster, and therefore, only run the `vllm` command on this node but not other nodes. The common practice is to set the tensor parallel size to the number of GPUs in each node, and the pipeline parallel size to the number of nodes. For example, if you have 16 GPUs in 2 nodes (8 GPUs per node), you can set the tensor parallel size to 8 and the pipeline parallel size to 2:
+### Running vLLM on a Ray cluster
+
+!!! tip
+     If Ray is running inside containers, run the commands in the remainder of this guide _inside the containers_, not on the host. To open a shell inside a container, connect to a node and use `docker exec -it <container_name> /bin/bash`.
+
+Once a Ray cluster is running, use vLLM as you would in a single-node setting. All resources across the Ray cluster are visible to vLLM, so a single `vllm` command on a single node is sufficient.
+
+The common practice is to set the tensor parallel size to the number of GPUs in each node, and the pipeline parallel size to the number of nodes. For example, if you have 16 GPUs across 2 nodes (8 GPUs per node), set the tensor parallel size to 8 and the pipeline parallel size to 2:

 ```bash
- vllm serve /path/to/the/model/in/the/container \
-     --tensor-parallel-size 8 \
-     --pipeline-parallel-size 2
+vllm serve /path/to/the/model/in/the/container \
+    --tensor-parallel-size 8 \
+    --pipeline-parallel-size 2
 ```

-You can also use tensor parallel without pipeline parallel, just set the tensor parallel size to the number of GPUs in the cluster. For example, if you have 16 GPUs in 2 nodes (8 GPUs per node), you can set the tensor parallel size to 16:
+Alternatively, you can set `tensor_parallel_size` to the total number of GPUs in the cluster:

 ```bash
 vllm serve /path/to/the/model/in/the/container \
     --tensor-parallel-size 16
 ```

-To make tensor parallel performant, you should make sure the communication between nodes is efficient, e.g. using high-speed network cards like InfiniBand. To correctly set up the cluster to use InfiniBand, append additional arguments like `--privileged -e NCCL_IB_HCA=mlx5` to the `run_cluster.sh` script. Please contact your system administrator for more information on how to set up the flags. One way to confirm if the InfiniBand is working is to run vLLM with `NCCL_DEBUG=TRACE` environment variable set, e.g. `NCCL_DEBUG=TRACE vllm serve ...` and check the logs for the NCCL version and the network used. If you find `[send] via NET/Socket` in the logs, it means NCCL uses raw TCP Socket, which is not efficient for cross-node tensor parallel. If you find `[send] via NET/IB/GDRDMA` in the logs, it means NCCL uses InfiniBand with GPUDirect RDMA, which is efficient.
+## Troubleshooting distributed deployments

-### GPUDirect RDMA
+To make tensor parallelism performant, ensure that communication between nodes is efficient, for example, by using high-speed network cards such as InfiniBand. To set up the cluster to use InfiniBand, append additional arguments like `--privileged -e NCCL_IB_HCA=mlx5` to the `run_cluster.sh` script. Contact your system administrator for more information about the required flags. One way to confirm if InfiniBand is working is to run `vllm` with the `NCCL_DEBUG=TRACE` environment variable set, for example `NCCL_DEBUG=TRACE vllm serve ...`, and check the logs for the NCCL version and the network used. If you find `[send] via NET/Socket` in the logs, NCCL uses a raw TCP socket, which is not efficient for cross-node tensor parallelism. If you find `[send] via NET/IB/GDRDMA` in the logs, NCCL uses InfiniBand with GPUDirect RDMA, which is efficient.

-To enable GPUDirect RDMA with vLLM, specific configuration tweaks are needed. This setup ensures:
+## Enabling GPUDirect RDMA

- `IPC_LOCK` Security Context: Add the `IPC_LOCK` capability to the container’s security context to lock memory pages and prevent swapping to disk.
- Shared Memory with `/dev/shm`: Mount `/dev/shm` in the pod spec to provide shared memory for IPC.
+To enable GPUDirect RDMA with vLLM, configure the following settings:

-When using Docker, you can set up the container as follows:
+- `IPC_LOCK` security context: add the `IPC_LOCK` capability to the container's security context to lock memory pages and prevent swapping to disk.
+- Shared memory with `/dev/shm`: mount `/dev/shm` in the pod spec to provide shared memory for interprocess communication (IPC).
+
+If you use Docker, set up the container as follows:

 ```bash
 docker run --gpus all \
@@ -120,7 +139,7 @@ docker run --gpus all \
    vllm/vllm-openai
 ```

-When using Kubernetes, you can set up the pod spec as follows:
+If you use Kubernetes, set up the pod spec as follows:

 ```yaml
 ...
@@ -146,13 +165,21 @@ spec:
 ...
 ```

-!!! warning
-    After you start the Ray cluster, you'd better also check the GPU-GPU communication between nodes. It can be non-trivial to set up. Please refer to the [sanity check script][troubleshooting-incorrect-hardware-driver] for more information. If you need to set some environment variables for the communication configuration, you can append them to the `run_cluster.sh` script, e.g. `-e NCCL_SOCKET_IFNAME=eth0`. Note that setting environment variables in the shell (e.g. `NCCL_SOCKET_IFNAME=eth0 vllm serve ...`) only works for the processes in the same node, not for the processes in the other nodes. Setting environment variables when you create the cluster is the recommended way. See <gh-issue:6803> for more information.
+Efficient tensor parallelism requires fast inter-node communication, preferably through high-speed network adapters such as InfiniBand. To enable InfiniBand, append flags such as `--privileged -e NCCL_IB_HCA=mlx5` to `run_cluster.sh`. For cluster-specific settings, consult your system administrator.

-!!! warning
-    Please make sure you downloaded the model to all the nodes (with the same path), or the model is downloaded to some distributed file system that is accessible by all nodes.
+To confirm InfiniBand operation, enable detailed NCCL logs:

-    When you use huggingface repo id to refer to the model, you should append your huggingface token to the `run_cluster.sh` script, e.g. `-e HF_TOKEN=`. The recommended way is to download the model first, and then use the path to refer to the model.
+```bash
+NCCL_DEBUG=TRACE vllm serve ...
+```

-!!! warning
-    If you keep receiving the error message `Error: No available node types can fulfill resource request` but you have enough GPUs in the cluster, chances are your nodes have multiple IP addresses and vLLM cannot find the right one, especially when you are using multi-node inference. Please make sure vLLM and ray use the same IP address. You can set the `VLLM_HOST_IP` environment variable to the right IP address in the `run_cluster.sh` script (different for each node!), and check `ray status` and `ray list nodes` to see the IP address used by Ray. See <gh-issue:7815> for more information.
+Search the logs for the transport method. Entries containing `[send] via NET/Socket` indicate raw TCP sockets, which perform poorly for cross-node tensor parallelism. Entries containing `[send] via NET/IB/GDRDMA` indicate InfiniBand with GPUDirect RDMA, which provides high performance.
+
+!!! tip "Verify inter-node GPU communication"
+    After you start the Ray cluster, verify GPU-to-GPU communication across nodes. Proper configuration can be non-trivial. For more information, see [troubleshooting script][troubleshooting-incorrect-hardware-driver]. If you need additional environment variables for communication configuration, append them to `run_cluster.sh`, for example `-e NCCL_SOCKET_IFNAME=eth0`. Setting environment variables during cluster creation is recommended because the variables propagate to all nodes. In contrast, setting environment variables in the shell affects only the local node. For more information, see <gh-issue:6803>.
+
+!!! tip "Pre-download Hugging Face models"
+    If you use Hugging Face models, downloading the model before starting vLLM is recommended. Download the model on every node to the same path, or store the model on a distributed file system accessible by all nodes. Then pass the path to the model in place of the repository ID. Otherwise, supply a Hugging Face token by appending `-e HF_TOKEN=<TOKEN>` to `run_cluster.sh`.
+
+!!! tip
+    The error message `Error: No available node types can fulfill resource request` can appear even when the cluster has enough GPUs. The issue often occurs when nodes have multiple IP addresses and vLLM can't select the correct one. Ensure that vLLM and Ray use the same IP address by setting `VLLM_HOST_IP` in `run_cluster.sh` (with a different value on each node). Use `ray status` and `ray list nodes` to verify the chosen IP address. For more information, see <gh-issue:7815>.
@@ -0,0 +1,244 @@
+# Expert Parallel Deployment
+
+vLLM supports Expert Parallelism (EP), which allows experts in Mixture-of-Experts (MoE) models to be deployed on separate GPUs, increasing locality, efficiency, and throughput overall.
+
+EP is typically coupled with Data Parallelism (DP). While DP can be used independently of EP, EP is more efficient when used in conjunction with DP. You can read more about data parallelism [here](data_parallel_deployment.md).
+
+## Prerequisites
+
+Before using EP, you need to install the necessary dependencies. We are actively working on making this easier in the future:
+
+1. **Install DeepEP and pplx-kernels**: Set up host environment following vLLM's guide for EP kernels [here](gh-file:tools/ep_kernels).
+2. **Install DeepGEMM library**: Follow the [official instructions](https://github.com/deepseek-ai/DeepGEMM#installation).
+3. **For disaggregated serving**: Install UCX and NIXL following the [script](gh-file:tools/install_nixl.sh).
+
+### Backend Selection Guide
+
+vLLM provides three communication backends for EP:
+
+| Backend | Use Case | Features | Best For |
+|---------|----------|----------|----------|
+| `pplx` | Single node | Chunked prefill support | Development, best for intra-node deployments |
+| `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout | High-throughput scenarios, prefill-dominated workloads |
+| `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout | Low-latency scenarios, decode-dominated workloads |
+
+## Single Node Deployment
+
+!!! warning
+    EP is an experimental feature. Argument names and default values may change in the future.
+
+### Configuration
+
+Enable EP by setting the `--enable-expert-parallel` flag. The EP size is automatically calculated as:
+
+```
+EP_SIZE = TP_SIZE × DP_SIZE
+```
+
+Where:
+- `TP_SIZE`: Tensor parallel size (always 1 for now)
+- `DP_SIZE`: Data parallel size
+- `EP_SIZE`: Expert parallel size (computed automatically)
+
+### Example Command
+
+The following command serves a `DeepSeek-V3-0324` model with 1-way tensor parallel, 8-way (attention) data parallel, and 8-way expert parallel. The attention weights are replicated across all GPUs, while the expert weights are split across GPUs. It will work on a H200 (or H20) node with 8 GPUs. For H100, you can try to serve a smaller model or refer to the multi-node deployment section.
+
+```bash
+# Single node EP deployment with pplx backend
+VLLM_ALL2ALL_BACKEND=pplx VLLM_USE_DEEP_GEMM=1 \
+    vllm serve deepseek-ai/DeepSeek-V3-0324 \
+    --tensor-parallel-size 1 \      # Tensor parallelism across 1 GPU
+    --data-parallel-size 8 \         # Data parallelism across 8 processes
+    --enable-expert-parallel         # Enable expert parallelism
+```
+
+## Multi-Node Deployment
+
+For multi-node deployment, use the DeepEP communication kernel with one of two modes (see [Backend Selection Guide](#backend-selection-guide) above).
+
+### Deployment Steps
+
+1. **Run one command per node** - Each node requires its own launch command
+2. **Configure networking** - Ensure proper IP addresses and port configurations
+3. **Set node roles** - First node handles requests, additional nodes run in headless mode
+
+### Example: 2-Node Deployment
+
+The following example deploys `DeepSeek-V3-0324` across 2 nodes using `deepep_low_latency` mode:
+
+```bash
+# Node 1 (Primary - handles incoming requests)
+VLLM_ALL2ALL_BACKEND=deepep_low_latency VLLM_USE_DEEP_GEMM=1 \
+    vllm serve deepseek-ai/DeepSeek-V3-0324 \
+    --tensor-parallel-size 1 \               # TP size per node
+    --enable-expert-parallel \               # Enable EP
+    --data-parallel-size 16 \                # Total DP size across all nodes
+    --data-parallel-size-local 8 \           # Local DP size on this node (8 GPUs per node)
+    --data-parallel-address 192.168.1.100 \  # Replace with actual IP of Node 1
+    --data-parallel-rpc-port 13345 \         # RPC communication port, can be any port as long as reachable by all nodes
+    --api-server-count=8                     # Number of API servers for load handling (scaling this out to total ranks are recommended)
+
+# Node 2 (Secondary - headless mode, no API server)
+VLLM_ALL2ALL_BACKEND=deepep_low_latency VLLM_USE_DEEP_GEMM=1 \
+    vllm serve deepseek-ai/DeepSeek-V3-0324 \
+    --tensor-parallel-size 1 \               # TP size per node
+    --enable-expert-parallel \               # Enable EP
+    --data-parallel-size 16 \                # Total DP size across all nodes
+    --data-parallel-size-local 8 \           # Local DP size on this node
+    --data-parallel-start-rank 8 \           # Starting rank offset for this node
+    --data-parallel-address 192.168.1.100 \  # IP of primary node (Node 1)
+    --data-parallel-rpc-port 13345 \         # Same RPC port as primary
+    --headless                               # No API server, worker only
+```
+
+### Key Configuration Notes
+
+- **Headless mode**: Secondary nodes run with `--headless` flag, meaning all client requests are handled by the primary node
+- **Rank calculation**: `--data-parallel-start-rank` should equal the cumulative local DP size of previous nodes
+- **Load scaling**: Adjust `--api-server-count` on the primary node to handle higher request loads
+
+### Network Configuration
+
+!!! important "InfiniBand Clusters"
+    On InfiniBand networked clusters, set this environment variable to prevent initialization hangs:
+    ```bash
+    export GLOO_SOCKET_IFNAME=eth0
+    ```
+    This ensures torch distributed group discovery uses Ethernet instead of InfiniBand for initial setup.
+
+## Expert Parallel Load Balancer (EPLB)
+
+While MoE models are typically trained so that each expert receives a similar number of tokens, in practice the distribution of tokens across experts can be highly skewed. vLLM provides an Expert Parallel Load Balancer (EPLB) to redistribute expert mappings across EP ranks, evening the load across experts.
+
+### Configuration
+
+Enable EPLB with the `--enable-eplb` flag.
+
+!!! note "Model Support"
+    Currently only DeepSeek V3 architecture is supported.
+
+When enabled, vLLM collects load statistics with every forward pass and periodically rebalances expert distribution.
+
+### EPLB Parameters
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--eplb-window-size` | Number of engine steps to track for rebalancing decisions | - |
+| `--eplb-step-interval` | Frequency of rebalancing (every N engine steps) | - |
+| `--eplb-log-balancedness` | Log balancedness metrics (avg tokens per expert ÷ max tokens per expert) | `false` |
+| `--num-redundant-experts` | Additional global experts per EP rank beyond equal distribution | `0` |
+
+### Expert Distribution Formula
+
+- **Default**: Each EP rank has `NUM_TOTAL_EXPERTS ÷ NUM_EP_RANKS` experts
+- **With redundancy**: Each EP rank has `(NUM_TOTAL_EXPERTS + NUM_REDUNDANT_EXPERTS) ÷ NUM_EP_RANKS` experts
+
+### Example Command
+
+Single node deployment with EPLB enabled:
+
+```bash
+# Single node with EPLB load balancing
+VLLM_ALL2ALL_BACKEND=pplx VLLM_USE_DEEP_GEMM=1 vllm serve deepseek-ai/DeepSeek-V3-0324 \
+    --tensor-parallel-size 1 \     # Tensor parallelism
+    --data-parallel-size 8 \        # Data parallelism  
+    --enable-expert-parallel \      # Enable EP
+    --enable-eplb \                 # Enable load balancer
+    --eplb-log-balancedness \       # Log balancing metrics
+    --eplb-window-size 1000 \       # Track last 1000 engine steps
+    --eplb-step-interval 3000       # Rebalance every 3000 steps
+```
+
+For multi-node deployment, add these EPLB flags to each node's command. We recommend setting `--num-redundant-experts` to 32 in large scale use cases so the most popular experts are always available.
+
+## Disaggregated Serving (Prefill/Decode Split)
+
+For production deployments requiring strict SLA guarantees for time-to-first-token and inter-token latency, disaggregated serving allows independent scaling of prefill and decode operations.
+
+### Architecture Overview
+
+- **Prefill Instance**: Uses `deepep_high_throughput` backend for optimal prefill performance
+- **Decode Instance**: Uses `deepep_low_latency` backend for minimal decode latency  
+- **KV Cache Transfer**: Connects instances via NIXL or other KV connectors
+
+### Setup Steps
+
+1. **Install KV Connector**: Install NIXL using the [installation script](gh-file:tools/install_nixl.sh)
+
+2. **Configure Both Instances**: Add this flag to both prefill and decode instances `--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}`
+
+3. **Client Orchestration**: Use the client-side script below to coordinate prefill/decode operations. We are actively working on routing solutions.
+
+### Client Orchestration Example
+
+```python
+from openai import OpenAI
+import uuid
+
+try:
+    # 1: Set up clients for prefill and decode instances
+    openai_api_key = "EMPTY"  # vLLM doesn't require a real API key
+    
+    # Replace these IP addresses with your actual instance addresses
+    prefill_client = OpenAI(
+        api_key=openai_api_key,
+        base_url="http://192.168.1.100:8000/v1",  # Prefill instance URL
+    )
+    decode_client = OpenAI(
+        api_key=openai_api_key,
+        base_url="http://192.168.1.101:8001/v1",  # Decode instance URL  
+    )
+    
+    # Get model name from prefill instance
+    models = prefill_client.models.list()
+    model = models.data[0].id
+    print(f"Using model: {model}")
+
+    # 2: Prefill Phase
+    # Generate unique request ID to link prefill and decode operations
+    request_id = str(uuid.uuid4())
+    print(f"Request ID: {request_id}")
+    
+    prefill_response = prefill_client.completions.create(
+        model=model,
+        # Prompt must exceed vLLM's block size (16 tokens) for PD to work
+        prompt="Write a detailed explanation of Paged Attention for Transformers works including the management of KV cache for multi-turn conversations",
+        max_tokens=1,  # Force prefill-only operation
+        extra_body={
+            "kv_transfer_params": {
+                "do_remote_decode": True,     # Enable remote decode
+                "do_remote_prefill": False,   # This is the prefill instance
+                "remote_engine_id": None,     # Will be populated by vLLM
+                "remote_block_ids": None,     # Will be populated by vLLM
+                "remote_host": None,          # Will be populated by vLLM
+                "remote_port": None           # Will be populated by vLLM
+            }
+        },
+        extra_headers={"X-Request-Id": request_id}
+    )
+    
+    print("-" * 50)
+    print("✓ Prefill completed successfully")
+    print(f"Prefill response: {prefill_response.choices[0].text}")
+    
+    # 3: Decode Phase
+    # Transfer KV cache parameters from prefill to decode instance
+    decode_response = decode_client.completions.create(
+        model=model,
+        prompt="This prompt is ignored during decode",  # Original prompt not needed
+        max_tokens=150,  # Generate up to 150 tokens
+        extra_body={
+            "kv_transfer_params": prefill_response.kv_transfer_params  # Pass KV cache info
+        },
+        extra_headers={"X-Request-Id": request_id}  # Same request ID
+    )
+    
+    print("-" * 50)
+    print("✓ Decode completed successfully")
+    print(f"Final response: {decode_response.choices[0].text}")
+
+except Exception as e:
+    print(f"❌ Error during disaggregated serving: {e}")
+    print("Check that both prefill and decode instances are running and accessible")
+```
@@ -351,6 +351,11 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
 Code example: <gh-file:examples/online_serving/openai_transcription_client.py>
 <!-- TODO: api enforced limits + uploading audios -->

+#### API Enforced Limits
+
+Set the maximum audio file size (in MB) that VLLM will accept, via the
+`VLLM_MAX_AUDIO_CLIP_FILESIZE_MB` environment variable. Default is 25 MB.
+
 #### Extra Parameters

 The following [sampling parameters][sampling-params] are supported.
@@ -1,122 +1,27 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This is a demo script showing how to use the
-PrithviGeospatialMAE model with vLLM
-This script is based on: https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11/blob/main/inference.py # noqa
-
-Target model weights: https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11/resolve/main/Prithvi-EO-V2-300M-TL-Sen1Floods11.pt # noqa
-
-The requirements for running this script are:
- Installing [terratorch, albumentations, rasterio] in your python environment
- downloading the model weights in a 'model' folder local to the script
-  (temporary measure until the proper config.json file is uploaded to HF)
- download an input example image (India_900498_S2Hand.tif) and place it in
-  the same folder with the script (or specify with the --data_file argument)
-
-Run the example:
-python prithvi_geospatial_mae.py
-
-"""  # noqa: E501
-
 import argparse
 import datetime
 import os
+import re
 from typing import Union

 import albumentations
 import numpy as np
 import rasterio
-import regex as re
 import torch
 from einops import rearrange
 from terratorch.datamodules import Sen1Floods11NonGeoDataModule

 from vllm import LLM

+torch.set_default_dtype(torch.float16)
+
 NO_DATA = -9999
 NO_DATA_FLOAT = 0.0001
 OFFSET = 0
 PERCENTILE = 99

-model_config = """{
-  "architectures": ["PrithviGeoSpatialMAE"],
-  "num_classes": 0,
-  "pretrained_cfg": {
-    "task_args": {
-      "task": "SemanticSegmentationTask",
-      "model_factory": "EncoderDecoderFactory",
-      "loss": "ce",
-      "ignore_index": -1,
-      "lr": 0.001,
-      "freeze_backbone": false,
-      "freeze_decoder": false,
-      "plot_on_val": 10,
-      "optimizer": "AdamW",
-      "scheduler": "CosineAnnealingLR"
-    },
-    "model_args": {
-      "backbone_pretrained": false,
-      "backbone": "prithvi_eo_v2_300_tl",
-      "decoder": "UperNetDecoder",
-      "decoder_channels": 256,
-      "decoder_scale_modules": true,
-      "num_classes": 2,
-      "rescale": true,
-      "backbone_bands": [
-        "BLUE",
-        "GREEN",
-        "RED",
-        "NIR_NARROW",
-        "SWIR_1",
-        "SWIR_2"
-      ],
-      "head_dropout": 0.1,
-      "necks": [
-        {
-          "name": "SelectIndices",
-          "indices": [
-            5,
-            11,
-            17,
-            23
-          ]
-        },
-        {
-          "name": "ReshapeTokensToImage"
-        }
-      ]
-    },
-    "optimizer_params" : {
-      "lr": 5.0e-05,
-      "betas": [0.9, 0.999],
-      "eps": [1.0e-08],
-      "weight_decay": 0.05,
-      "amsgrad": false,
-      "maximize": false,
-      "capturable": false,
-      "differentiable": false
-    },
-    "scheduler_params" : {
-        "T_max": 50,
-        "eta_min": 0,
-        "last_epoch": -1,
-        "verbose": "deprecated"
-    }
-  },
-
-
-  "torch_dtype": "float32"
-}
-"""
-
-# Temporarily creating the "config.json" for the model.
-# This is going to disappear once the correct config.json is available on HF
-with open(
-    os.path.join(os.path.dirname(__file__), "./model/config.json"), "w"
-) as config_file:
-    config_file.write(model_config)
-
 datamodule_config = {
    "bands": ["BLUE", "GREEN", "RED", "NIR_NARROW", "SWIR_1", "SWIR_2"],
    "batch_size": 16,
@@ -138,28 +43,24 @@ datamodule_config = {


 class PrithviMAE:
-    def __init__(self):
-        print("Initializing PrithviMAE model")
-        self.llm = LLM(
-            model=os.path.join(os.path.dirname(__file__), "./model"),
-            skip_tokenizer_init=True,
-            dtype="float32",
+    def __init__(self, model):
+        self.model = LLM(
+            model=model, skip_tokenizer_init=True, dtype="float16", enforce_eager=True
        )

    def run(self, input_data, location_coords):
-        print("################ Running inference on vLLM ##############")
        # merge the inputs into one data structure
+        if input_data is not None and input_data.dtype == torch.float32:
+            input_data = input_data.to(torch.float16)
+            input_data = input_data[0]
+
        mm_data = {
-            "pixel_values": torch.empty(0) if input_data is None else input_data,
-            "location_coords": torch.empty(0)
-            if location_coords is None
-            else location_coords,
+            "pixel_values": input_data,
+            "location_coords": location_coords,
        }

        prompt = {"prompt_token_ids": [1], "multi_modal_data": mm_data}
-
-        outputs = self.llm.encode(prompt, use_tqdm=False)
-        print("################ Inference done (it took seconds)  ##############")
+        outputs = self.model.encode(prompt, use_tqdm=False)

        return outputs[0].outputs.data

@@ -181,11 +82,12 @@ def process_channel_group(orig_img, channels):
    """
    Args:
        orig_img: torch.Tensor representing original image (reference)
-                  with shape = (bands, H, W).
+        with shape = (bands, H, W).
        channels: list of indices representing RGB channels.

    Returns:
-        torch.Tensor with shape (num_channels, height, width) for original image
+        torch.Tensor with shape (num_channels, height, width)
+        for original image
    """

    orig_img = orig_img[channels, ...]
@@ -260,10 +162,10 @@ def load_example(

    Args:
        file_paths: list of file paths .
-        mean: list containing mean values for each band in the images
-              in *file_paths*.
-        std: list containing std values for each band in the images
-             in *file_paths*.
+        mean: list containing mean values for each band in the
+              images in *file_paths*.
+        std: list containing std values for each band in the
+             images in *file_paths*.

    Returns:
        np.array containing created example
@@ -308,7 +210,7 @@ def load_example(
            print(f"Could not extract timestamp for {file} ({e})")

    imgs = np.stack(imgs, axis=0)  # num_frames, H, W, C
-    imgs = np.moveaxis(imgs, -1, 0).astype("float32")
+    imgs = np.moveaxis(imgs, -1, 0).astype("float32")  # C, num_frames, H, W
    imgs = np.expand_dims(imgs, axis=0)  # add batch di

    return imgs, temporal_coords, location_coords, metas
@@ -332,8 +234,10 @@ def run_model(
    )

    # Build sliding window
+
    batch_size = 1
-    batch = torch.tensor(input_data, device="cpu")
+    # batch = torch.tensor(input_data, device="cpu")
+    batch = torch.tensor(input_data)
    windows = batch.unfold(3, img_size, img_size).unfold(4, img_size, img_size)
    h1, w1 = windows.shape[3:5]
    windows = rearrange(
@@ -344,18 +248,16 @@ def run_model(
    num_batches = windows.shape[0] // batch_size if windows.shape[0] > batch_size else 1
    windows = torch.tensor_split(windows, num_batches, dim=0)

-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-
    if temporal_coords:
-        temporal_coords = torch.tensor(temporal_coords, device=device).unsqueeze(0)
+        temporal_coords = torch.tensor(temporal_coords).unsqueeze(0)
    else:
        temporal_coords = None
    if location_coords:
-        location_coords = torch.tensor(location_coords[0], device=device).unsqueeze(0)
+        location_coords = torch.tensor(location_coords[0]).unsqueeze(0)
    else:
        location_coords = None

-    # Run model
+    # Run Prithvi-EO-V2-300M-TL-Sen1Floods11
    pred_imgs = []
    for x in windows:
        # Apply standardization
@@ -363,15 +265,7 @@ def run_model(
        x = datamodule.aug(x)["image"]

        with torch.no_grad():
-            x = x.to(device)
            pred = model.run(x, location_coords=location_coords)
-            if lightning_model:
-                pred_lightning = lightning_model(
-                    x, temporal_coords=temporal_coords, location_coords=location_coords
-                )
-                pred_lightning = pred_lightning.output.detach().cpu()
-                if not torch.equal(pred, pred_lightning):
-                    print("Inference output is not equal")
        y_hat = pred.argmax(dim=1)

        y_hat = torch.nn.functional.interpolate(
@@ -403,52 +297,18 @@ def run_model(
    return pred_imgs


-def parse_args():
-    parser = argparse.ArgumentParser("MAE run inference", add_help=False)
-
-    parser.add_argument(
-        "--data_file",
-        type=str,
-        default="./India_900498_S2Hand.tif",
-        help="Path to the file.",
-    )
-    parser.add_argument(
-        "--output_dir",
-        type=str,
-        default="output",
-        help="Path to the directory where to save outputs.",
-    )
-    parser.add_argument(
-        "--input_indices",
-        default=[1, 2, 3, 8, 11, 12],
-        type=int,
-        nargs="+",
-        help="0-based indices of the six Prithvi channels to be selected from the  "
-        "input. By default selects [1,2,3,8,11,12] for S2L1C data.",
-    )
-    parser.add_argument(
-        "--rgb_outputs",
-        action="store_true",
-        help="If present, output files will only contain RGB channels. "
-        "Otherwise, all bands will be saved.",
-    )
-
-
 def main(
    data_file: str,
+    model: str,
    output_dir: str,
    rgb_outputs: bool,
    input_indices: list[int] = None,
 ):
    os.makedirs(output_dir, exist_ok=True)

-    # Load model ---------------------------------------------------------------
-
-    model_obj = PrithviMAE()
+    model_obj = PrithviMAE(model=model)
    datamodule = generate_datamodule()
-    img_size = 256  # Size of Sen1Floods11
-
-    # Loading data -------------------------------------------------------------
+    img_size = 512  # Size of Sen1Floods11

    input_data, temporal_coords, location_coords, meta_data = load_example(
        file_paths=[data_file],
@@ -460,8 +320,6 @@ def main(
    if input_data.mean() > 1:
        input_data = input_data / 10000  # Convert to range 0-1

-    # Running model ------------------------------------------------------------
-
    channels = [
        datamodule_config["bands"].index(b) for b in ["RED", "GREEN", "BLUE"]
    ]  # BGR -> RGB
@@ -469,7 +327,6 @@ def main(
    pred = run_model(
        input_data, temporal_coords, location_coords, model_obj, datamodule, img_size
    )
-
    # Save pred
    meta_data.update(count=1, dtype="uint8", compress="lzw", nodata=0)
    pred_file = os.path.join(
@@ -487,6 +344,7 @@ def main(
        orig_img=torch.Tensor(input_data[0, :, 0, ...]),
        channels=channels,
    )
+    rgb_orig = rgb_orig.to(torch.float32)

    pred[pred == 0.0] = np.nan
    img_pred = rgb_orig * 0.7 + pred * 0.3
@@ -503,9 +361,10 @@ def main(

    # Save image rgb
    if rgb_outputs:
+        name_suffix = os.path.splitext(os.path.basename(data_file))[0]
        rgb_file = os.path.join(
            output_dir,
-            f"original_rgb_{os.path.splitext(os.path.basename(data_file))[0]}.tiff",
+            f"original_rgb_{name_suffix}.tiff",
        )
        save_geotiff(
            image=_convert_np_uint8(rgb_orig),
@@ -515,6 +374,42 @@ def main(


 if __name__ == "__main__":
-    args = parse_args()
+    parser = argparse.ArgumentParser("MAE run inference", add_help=False)
+
+    parser.add_argument(
+        "--data_file",
+        type=str,
+        default="./India_900498_S2Hand.tif",
+        help="Path to the file.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
+        help="Path to a checkpoint file to load from.",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="output",
+        help="Path to the directory where to save outputs.",
+    )
+    parser.add_argument(
+        "--input_indices",
+        default=[1, 2, 3, 8, 11, 12],
+        type=int,
+        nargs="+",
+        help="""
+        0-based indices of the six Prithvi channels to be selected from the input.
+        By default selects [1,2,3,8,11,12] for S2L1C data.
+        """,
+    )
+    parser.add_argument(
+        "--rgb_outputs",
+        action="store_true",
+        help="If present, output files will only contain RGB channels. "
+        "Otherwise, all bands will be saved.",
+    )
+    args = parser.parse_args()

    main(**vars(args))
@@ -93,6 +93,7 @@ ensure_python_library_installed() {
 cleanup() {
    echo "Stopping everything…"
    trap - INT TERM        # prevent re-entrancy
+    pkill -9 -f "disagg_proxy_p2p_nccl_xpyd.py"
    kill -- -$$            # negative PID  ==  "this whole process-group"
    wait                   # reap children so we don't leave zombies
    exit 0
@@ -84,18 +84,22 @@ Or for deserializing:
 Once a model is serialized, tensorizer can be invoked with the `LLM` class 
 directly to load models:

-    llm = LLM(model="facebook/opt-125m",
-              load_format="tensorizer",
-              model_loader_extra_config=TensorizerConfig(
-                    tensorizer_uri = path_to_tensors,
-                    num_readers=3,
-                    )
-              )
+```python
+from vllm import LLM
+llm = LLM(
+    "s3://my-bucket/vllm/facebook/opt-125m/v1", 
+    load_format="tensorizer"
+)
+```
+
            
 A serialized model can be used during model loading for the vLLM OpenAI
-inference server. `model_loader_extra_config` is exposed as the CLI arg
-`--model-loader-extra-config`, and accepts a JSON string literal of the
-TensorizerConfig arguments desired.
+inference server:
+
+```
+vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
+    --load-format tensorizer
+```

 In order to see all of the available arguments usable to configure 
 loading with tensorizer that are given to `TensorizerConfig`, run:
@@ -116,10 +120,9 @@ the LoRA artifacts are in your model artifacts directory and specifying
 `--enable-lora`. For instance:

 ```
-vllm serve <model_path> \
+vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
    --load-format tensorizer \
-    --model-loader-extra-config '{"tensorizer_uri": "<model_path>.tensors"}' \
-    --enable-lora
+    --enable-lora 
 ```
 """

@@ -1,5 +1,5 @@
 site_name: vLLM
-site_url: https://docs.vllm.ai
+site_url: !ENV READTHEDOCS_CANONICAL_URL
 repo_url: https://github.com/vllm-project/vllm
 edit_uri: edit/main/docs/
 exclude_docs: |
@@ -72,7 +72,6 @@ line-length = 80
 "vllm/core/**/*.py" = ["UP006", "UP035"]
 "vllm/engine/**/*.py" = ["UP006", "UP035"]
 "vllm/executor/**/*.py" = ["UP006", "UP035"]
-"vllm/prompt_adapter/**/*.py" = ["UP006", "UP035"]
 "vllm/worker/**/*.py" = ["UP006", "UP035"]
 # Python 3.8 typing - skip utils for ROCm
 "vllm/utils/__init__.py" = ["UP006", "UP035"]
@@ -33,7 +33,7 @@ pyzmq >= 25.0.0
 msgspec
 gguf >= 0.13.0
 importlib_metadata; python_version < '3.10'
-mistral_common[opencv] >= 1.8.0
+mistral_common[image,audio] >= 1.8.2
 opencv-python-headless >= 4.11.0    # required for video IO
 pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12
@@ -48,3 +48,4 @@ scipy # Required for phi-4-multimodal-instruct
 ninja # Required for xgrammar, rocm, tpu, xpu
 pybase64 # fast base64 implementation
 cbor2 # Required for cross-language serialization of hashable objects
+setproctitle # Used to set process names for better debugging and monitoring
@@ -22,6 +22,7 @@ pillow
 psutil
 pybase64
 pydantic
+setproctitle
 torch
 transformers
 zmq
@@ -23,7 +23,7 @@ jiwer # required for audio tests
 timm # required for internvl test
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[opencv] >= 1.8.0 # required for voxtral test
+mistral_common[image,audio] >= 1.8.2 # required for voxtral test
 num2words # required for smolvlm test
 opencv-python-headless >= 4.11.0 # required for video test
 datamodel_code_generator # required for minicpm3 test
@@ -26,9 +26,9 @@ torch==2.7.1
 torchaudio==2.7.1
 torchvision==0.22.1
 transformers_stream_generator # required for qwen-vl test
-mamba_ssm # required for plamo2 test
+mamba_ssm==2.2.5 # required for plamo2 test
 matplotlib # required for qwen-vl test
-mistral_common[opencv] >= 1.8.0 # required for voxtral test
+mistral_common[image,audio] >= 1.8.2 # required for voxtral test
 num2words # required for smolvlm test
 open_clip_torch==2.32.0 # Required for nemotron_vl test
 opencv-python-headless >= 4.11.0 # required for video test
@@ -54,3 +54,4 @@ runai-model-streamer==0.11.0
 runai-model-streamer-s3==0.11.0
 fastsafetensors>=0.1.10
 pydantic>=2.10 # 2.9 leads to error on python 3.10
+terratorch==1.1rc2 # required for PrithviMAE test
@@ -6,6 +6,10 @@ accelerate==1.0.1
    # via
    #   lm-eval
    #   peft
+aenum==3.1.16
+    # via lightly
+affine==2.4.0
+    # via rasterio
 aiohappyeyeballs==2.4.3
    # via aiohttp
 aiohttp==3.10.11
@@ -21,8 +25,18 @@ aiosignal==1.3.1
    # via
    #   aiohttp
    #   ray
+albucore==0.0.16
+    # via terratorch
+albumentations==1.4.6
+    # via terratorch
+alembic==1.16.4
+    # via mlflow
 annotated-types==0.7.0
    # via pydantic
+antlr4-python3-runtime==4.9.3
+    # via
+    #   hydra-core
+    #   omegaconf
 anyio==4.6.2.post1
    # via
    #   httpx
@@ -34,10 +48,12 @@ arrow==1.3.0
 attrs==24.2.0
    # via
    #   aiohttp
+    #   fiona
    #   hypothesis
    #   jsonlines
    #   jsonschema
    #   pytest-subtests
+    #   rasterio
    #   referencing
 audioread==3.0.1
    # via librosa
@@ -46,9 +62,13 @@ backoff==2.2.1
    #   -r requirements/test.in
    #   schemathesis
 bitsandbytes==0.46.1
-    # via -r requirements/test.in
+    # via
+    #   -r requirements/test.in
+    #   lightning
 black==24.10.0
    # via datamodel-code-generator
+blinker==1.9.0
+    # via flask
 blobfile==3.0.0
    # via -r requirements/test.in
 bm25s==0.2.13
@@ -64,11 +84,18 @@ bounded-pool-executor==0.0.3
 buildkite-test-collector==0.1.9
    # via -r requirements/test.in
 cachetools==5.5.2
-    # via google-auth
+    # via
+    #   google-auth
+    #   mlflow-skinny
 certifi==2024.8.30
    # via
+    #   fiona
    #   httpcore
    #   httpx
+    #   lightly
+    #   pyogrio
+    #   pyproj
+    #   rasterio
    #   requests
 cffi==1.17.1
    # via soundfile
@@ -79,11 +106,28 @@ charset-normalizer==3.4.0
 click==8.1.7
    # via
    #   black
+    #   click-plugins
+    #   cligj
+    #   fiona
+    #   flask
    #   jiwer
+    #   mlflow-skinny
    #   nltk
+    #   rasterio
    #   ray
    #   schemathesis
    #   typer
+    #   uvicorn
+click-plugins==1.1.1.2
+    # via
+    #   fiona
+    #   rasterio
+cligj==0.7.2
+    # via
+    #   fiona
+    #   rasterio
+cloudpickle==3.1.1
+    # via mlflow-skinny
 colorama==0.4.6
    # via
    #   sacrebleu
@@ -99,6 +143,8 @@ cupy-cuda12x==13.3.0
    # via ray
 cycler==0.12.1
    # via matplotlib
+databricks-sdk==0.59.0
+    # via mlflow-skinny
 datamodel-code-generator==0.26.3
    # via -r requirements/test.in
 dataproperty==1.0.1
@@ -122,13 +168,21 @@ distlib==0.3.9
    # via virtualenv
 dnspython==2.7.0
    # via email-validator
+docker==7.1.0
+    # via mlflow
 docopt==0.6.2
    # via num2words
-einops==0.8.0
+docstring-parser==0.17.0
+    # via jsonargparse
+efficientnet-pytorch==0.7.1
+    # via segmentation-models-pytorch
+einops==0.8.1
    # via
    #   -r requirements/test.in
    #   encodec
    #   mamba-ssm
+    #   terratorch
+    #   torchgeo
    #   vector-quantize-pytorch
    #   vocos
 einx==0.3.0
@@ -141,6 +195,8 @@ eval-type-backport==0.2.2
    # via mteb
 evaluate==0.4.3
    # via lm-eval
+fastapi==0.116.1
+    # via mlflow-skinny
 fastparquet==2024.11.0
    # via genai-perf
 fastrlock==0.8.2
@@ -156,6 +212,10 @@ filelock==3.16.1
    #   torch
    #   transformers
    #   virtualenv
+fiona==1.10.1
+    # via torchgeo
+flask==3.1.1
+    # via mlflow
 fonttools==4.54.1
    # via matplotlib
 fqdn==1.5.1
@@ -173,6 +233,8 @@ fsspec==2024.9.0
    #   evaluate
    #   fastparquet
    #   huggingface-hub
+    #   lightning
+    #   pytorch-lightning
    #   torch
 ftfy==6.3.1
    # via open-clip-torch
@@ -180,18 +242,41 @@ genai-perf==0.0.8
    # via -r requirements/test.in
 genson==1.3.0
    # via datamodel-code-generator
+geopandas==1.0.1
+    # via terratorch
+gitdb==4.0.12
+    # via gitpython
+gitpython==3.1.44
+    # via mlflow-skinny
 google-api-core==2.24.2
    # via opencensus
 google-auth==2.40.2
-    # via google-api-core
+    # via
+    #   databricks-sdk
+    #   google-api-core
 googleapis-common-protos==1.70.0
    # via google-api-core
+graphene==3.4.3
+    # via mlflow
 graphql-core==3.2.6
-    # via hypothesis-graphql
+    # via
+    #   graphene
+    #   graphql-relay
+    #   hypothesis-graphql
+graphql-relay==3.2.0
+    # via graphene
+greenlet==3.2.3
+    # via sqlalchemy
 grpcio==1.71.0
    # via ray
+gunicorn==23.0.0
+    # via mlflow
 h11==0.14.0
-    # via httpcore
+    # via
+    #   httpcore
+    #   uvicorn
+h5py==3.13.0
+    # via terratorch
 harfile==0.3.0
    # via schemathesis
 hf-xet==1.1.3
@@ -204,7 +289,7 @@ httpx==0.27.2
    # via
    #   -r requirements/test.in
    #   schemathesis
-huggingface-hub==0.33.0
+huggingface-hub==0.33.1
    # via
    #   -r requirements/test.in
    #   accelerate
@@ -212,13 +297,19 @@ huggingface-hub==0.33.0
    #   evaluate
    #   open-clip-torch
    #   peft
+    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   terratorch
    #   timm
    #   tokenizers
    #   transformers
    #   vocos
 humanize==4.11.0
    # via runai-model-streamer
+hydra-core==1.3.2
+    # via
+    #   lightly
+    #   lightning
 hypothesis==6.131.0
    # via
    #   hypothesis-graphql
@@ -236,6 +327,14 @@ idna==3.10
    #   jsonschema
    #   requests
    #   yarl
+imageio==2.37.0
+    # via scikit-image
+importlib-metadata==8.7.0
+    # via
+    #   mlflow-skinny
+    #   opentelemetry-api
+importlib-resources==6.5.2
+    # via typeshed-client
 inflect==5.6.2
    # via datamodel-code-generator
 iniconfig==2.0.0
@@ -244,9 +343,13 @@ isoduration==20.11.0
    # via jsonschema
 isort==5.13.2
    # via datamodel-code-generator
+itsdangerous==2.2.0
+    # via flask
 jinja2==3.1.6
    # via
    #   datamodel-code-generator
+    #   flask
+    #   mlflow
    #   torch
 jiwer==3.0.5
    # via -r requirements/test.in
@@ -259,6 +362,10 @@ joblib==1.4.2
    #   librosa
    #   nltk
    #   scikit-learn
+jsonargparse==4.35.0
+    # via
+    #   lightning
+    #   terratorch
 jsonlines==4.0.0
    # via lm-eval
 jsonpointer==3.0.0
@@ -277,12 +384,33 @@ kaleido==0.2.1
    # via genai-perf
 kiwisolver==1.4.7
    # via matplotlib
+kornia==0.8.1
+    # via torchgeo
+kornia-rs==0.1.9
+    # via kornia
 lazy-loader==0.4
-    # via librosa
+    # via
+    #   librosa
+    #   scikit-image
 libnacl==2.1.0
    # via tensorizer
 librosa==0.10.2.post1
    # via -r requirements/test.in
+lightly==1.5.20
+    # via
+    #   terratorch
+    #   torchgeo
+lightly-utils==0.0.2
+    # via lightly
+lightning==2.5.1.post0
+    # via
+    #   terratorch
+    #   torchgeo
+lightning-utilities==0.14.3
+    # via
+    #   lightning
+    #   pytorch-lightning
+    #   torchmetrics
 llvmlite==0.44.0
    # via numba
 lm-eval==0.4.8
@@ -291,16 +419,27 @@ lxml==5.3.0
    # via
    #   blobfile
    #   sacrebleu
-mamba-ssm==2.2.4
+mako==1.3.10
+    # via alembic
+mamba-ssm==2.2.5
    # via -r requirements/test.in
+markdown==3.8.2
+    # via mlflow
 markdown-it-py==3.0.0
    # via rich
 markupsafe==3.0.1
    # via
+    #   flask
    #   jinja2
+    #   mako
    #   werkzeug
 matplotlib==3.9.2
-    # via -r requirements/test.in
+    # via
+    #   -r requirements/test.in
+    #   lightning
+    #   mlflow
+    #   pycocotools
+    #   torchgeo
 mbstrdecoder==1.1.3
    # via
    #   dataproperty
@@ -308,8 +447,12 @@ mbstrdecoder==1.1.3
    #   typepy
 mdurl==0.1.2
    # via markdown-it-py
-mistral-common==1.8.0
+mistral-common==1.8.2
    # via -r requirements/test.in
+mlflow==2.22.0
+    # via terratorch
+mlflow-skinny==2.22.0
+    # via mlflow
 more-itertools==10.5.0
    # via lm-eval
 mpmath==1.3.0
@@ -328,10 +471,14 @@ multiprocess==0.70.16
    # via
    #   datasets
    #   evaluate
+munch==4.0.0
+    # via pretrainedmodels
 mypy-extensions==1.0.0
    # via black
 networkx==3.2.1
-    # via torch
+    # via
+    #   scikit-image
+    #   torch
 ninja==1.11.1.3
    # via mamba-ssm
 nltk==3.9.1
@@ -348,6 +495,8 @@ numpy==1.26.4
    # via
    #   -r requirements/test.in
    #   accelerate
+    #   albucore
+    #   albumentations
    #   bitsandbytes
    #   bm25s
    #   contourpy
@@ -358,9 +507,15 @@ numpy==1.26.4
    #   evaluate
    #   fastparquet
    #   genai-perf
+    #   geopandas
+    #   h5py
+    #   imageio
    #   librosa
+    #   lightly
+    #   lightly-utils
    #   matplotlib
    #   mistral-common
+    #   mlflow
    #   mteb
    #   numba
    #   numexpr
@@ -368,18 +523,30 @@ numpy==1.26.4
    #   pandas
    #   patsy
    #   peft
+    #   pycocotools
+    #   pyogrio
+    #   rasterio
+    #   rioxarray
    #   rouge-score
    #   runai-model-streamer
    #   sacrebleu
+    #   scikit-image
    #   scikit-learn
    #   scipy
+    #   segmentation-models-pytorch
+    #   shapely
    #   soxr
    #   statsmodels
+    #   tensorboardx
    #   tensorizer
+    #   tifffile
+    #   torchgeo
+    #   torchmetrics
    #   torchvision
    #   transformers
    #   tritonclient
    #   vocos
+    #   xarray
 nvidia-cublas-cu12==12.8.3.14
    # via
    #   nvidia-cudnn-cu12
@@ -417,6 +584,10 @@ nvidia-nvjitlink-cu12==12.8.61
    #   torch
 nvidia-nvtx-cu12==12.8.55
    # via torch
+omegaconf==2.3.0
+    # via
+    #   hydra-core
+    #   lightning
 open-clip-torch==2.32.0
    # via -r requirements/test.in
 opencensus==0.11.4
@@ -426,7 +597,18 @@ opencensus-context==0.1.3
 opencv-python-headless==4.11.0.86
    # via
    #   -r requirements/test.in
+    #   albucore
+    #   albumentations
    #   mistral-common
+opentelemetry-api==1.35.0
+    # via
+    #   mlflow-skinny
+    #   opentelemetry-sdk
+    #   opentelemetry-semantic-conventions
+opentelemetry-sdk==1.35.0
+    # via mlflow-skinny
+opentelemetry-semantic-conventions==0.56b0
+    # via opentelemetry-sdk
 packaging==24.2
    # via
    #   accelerate
@@ -435,26 +617,44 @@ packaging==24.2
    #   datasets
    #   evaluate
    #   fastparquet
+    #   geopandas
+    #   gunicorn
    #   huggingface-hub
+    #   hydra-core
+    #   kornia
    #   lazy-loader
+    #   lightning
+    #   lightning-utilities
    #   mamba-ssm
    #   matplotlib
+    #   mlflow-skinny
    #   peft
    #   plotly
    #   pooch
+    #   pyogrio
    #   pytest
    #   pytest-rerunfailures
+    #   pytorch-lightning
    #   ray
+    #   rioxarray
+    #   scikit-image
    #   statsmodels
+    #   tensorboardx
+    #   torchmetrics
    #   transformers
    #   typepy
+    #   xarray
 pandas==2.2.3
    # via
    #   datasets
    #   evaluate
    #   fastparquet
    #   genai-perf
+    #   geopandas
+    #   mlflow
    #   statsmodels
+    #   torchgeo
+    #   xarray
 pathspec==0.12.1
    # via black
 pathvalidate==3.2.1
@@ -468,9 +668,14 @@ peft==0.13.2
 pillow==10.4.0
    # via
    #   genai-perf
+    #   imageio
+    #   lightly-utils
    #   matplotlib
    #   mistral-common
+    #   scikit-image
+    #   segmentation-models-pytorch
    #   sentence-transformers
+    #   torchgeo
    #   torchvision
 platformdirs==4.3.6
    # via
@@ -489,6 +694,8 @@ portalocker==2.10.1
    # via sacrebleu
 pqdm==0.2.0
    # via -r requirements/test.in
+pretrainedmodels==0.7.4
+    # via segmentation-models-pytorch
 prometheus-client==0.22.0
    # via ray
 propcache==0.2.0
@@ -499,8 +706,10 @@ protobuf==5.28.3
    # via
    #   google-api-core
    #   googleapis-common-protos
+    #   mlflow-skinny
    #   proto-plus
    #   ray
+    #   tensorboardx
    #   tensorizer
 psutil==6.1.0
    # via
@@ -515,6 +724,7 @@ pyarrow==18.0.0
    # via
    #   datasets
    #   genai-perf
+    #   mlflow
 pyasn1==0.6.1
    # via
    #   pyasn1-modules
@@ -523,6 +733,8 @@ pyasn1-modules==0.4.2
    # via google-auth
 pybind11==2.13.6
    # via lm-eval
+pycocotools==2.0.8
+    # via terratorch
 pycountry==24.6.1
    # via pydantic-extra-types
 pycparser==2.22
@@ -532,8 +744,12 @@ pycryptodomex==3.22.0
 pydantic==2.11.5
    # via
    #   -r requirements/test.in
+    #   albumentations
    #   datamodel-code-generator
+    #   fastapi
+    #   lightly
    #   mistral-common
+    #   mlflow-skinny
    #   mteb
    #   pydantic-extra-types
    #   ray
@@ -543,15 +759,24 @@ pydantic-extra-types==2.10.5
    # via mistral-common
 pygments==2.18.0
    # via rich
+pyogrio==0.11.0
+    # via geopandas
 pyparsing==3.2.0
-    # via matplotlib
+    # via
+    #   matplotlib
+    #   rasterio
+pyproj==3.7.1
+    # via
+    #   geopandas
+    #   rioxarray
+    #   torchgeo
 pyrate-limiter==3.7.0
    # via schemathesis
 pystemmer==3.0.0
    # via mteb
 pytablewriter==1.2.0
    # via lm-eval
-pytest==8.3.3
+pytest==8.3.5
    # via
    #   -r requirements/test.in
    #   buildkite-test-collector
@@ -564,6 +789,7 @@ pytest==8.3.3
    #   pytest-subtests
    #   pytest-timeout
    #   schemathesis
+    #   terratorch
 pytest-asyncio==0.24.0
    # via -r requirements/test.in
 pytest-forked==1.6.0
@@ -578,15 +804,23 @@ pytest-subtests==0.14.1
    # via schemathesis
 pytest-timeout==2.3.1
    # via -r requirements/test.in
+python-box==7.3.2
+    # via terratorch
 python-dateutil==2.9.0.post0
    # via
    #   arrow
    #   botocore
+    #   graphene
+    #   lightly
    #   matplotlib
    #   pandas
    #   typepy
 python-rapidjson==1.20
    # via tritonclient
+pytorch-lightning==2.5.2
+    # via
+    #   lightly
+    #   lightning
 pytrec-eval-terrier==0.5.7
    # via mteb
 pytz==2024.2
@@ -596,11 +830,17 @@ pytz==2024.2
 pyyaml==6.0.2
    # via
    #   accelerate
+    #   albumentations
    #   datamodel-code-generator
    #   datasets
    #   genai-perf
    #   huggingface-hub
+    #   jsonargparse
+    #   lightning
+    #   mlflow-skinny
+    #   omegaconf
    #   peft
+    #   pytorch-lightning
    #   ray
    #   responses
    #   schemathesis
@@ -609,6 +849,11 @@ pyyaml==6.0.2
    #   vocos
 rapidfuzz==3.12.1
    # via jiwer
+rasterio==1.4.3
+    # via
+    #   rioxarray
+    #   terratorch
+    #   torchgeo
 ray==2.43.0
    # via -r requirements/test.in
 redis==5.2.0
@@ -627,12 +872,16 @@ regex==2024.9.11
 requests==2.32.3
    # via
    #   buildkite-test-collector
+    #   databricks-sdk
    #   datasets
+    #   docker
    #   evaluate
    #   google-api-core
    #   huggingface-hub
+    #   lightly
    #   lm-eval
    #   mistral-common
+    #   mlflow-skinny
    #   mteb
    #   pooch
    #   ray
@@ -650,8 +899,11 @@ rfc3987==1.3.8
 rich==13.9.4
    # via
    #   genai-perf
+    #   lightning
    #   mteb
    #   typer
+rioxarray==0.19.0
+    # via terratorch
 rouge-score==0.1.2
    # via lm-eval
 rpds-py==0.20.1
@@ -660,6 +912,8 @@ rpds-py==0.20.1
    #   referencing
 rsa==4.9.1
    # via google-auth
+rtree==1.4.0
+    # via torchgeo
 runai-model-streamer==0.11.0
    # via -r requirements/test.in
 runai-model-streamer-s3==0.11.0
@@ -677,21 +931,32 @@ safetensors==0.4.5
    #   transformers
 schemathesis==3.39.15
    # via -r requirements/test.in
+scikit-image==0.25.2
+    # via albumentations
 scikit-learn==1.5.2
    # via
+    #   albumentations
    #   librosa
    #   lm-eval
+    #   mlflow
    #   mteb
    #   sentence-transformers
 scipy==1.13.1
    # via
+    #   albumentations
    #   bm25s
    #   librosa
+    #   mlflow
    #   mteb
+    #   scikit-image
    #   scikit-learn
    #   sentence-transformers
    #   statsmodels
    #   vocos
+segmentation-models-pytorch==0.4.0
+    # via
+    #   terratorch
+    #   torchgeo
 sentence-transformers==3.2.1
    # via
    #   -r requirements/test.in
@@ -700,21 +965,30 @@ sentencepiece==0.2.0
    # via mistral-common
 setuptools==77.0.3
    # via
+    #   lightning-utilities
    #   mamba-ssm
    #   pytablewriter
    #   torch
    #   triton
+shapely==2.1.1
+    # via
+    #   geopandas
+    #   torchgeo
 shellingham==1.5.4
    # via typer
 six==1.16.0
    # via
    #   junit-xml
+    #   lightly
    #   opencensus
    #   python-dateutil
    #   rfc3339-validator
    #   rouge-score
+    #   segmentation-models-pytorch
 smart-open==7.1.0
    # via ray
+smmap==5.0.2
+    # via gitdb
 sniffio==1.3.1
    # via
    #   anyio
@@ -725,12 +999,22 @@ soundfile==0.12.1
    # via
    #   -r requirements/test.in
    #   librosa
+    #   mistral-common
 soxr==0.5.0.post1
-    # via librosa
+    # via
+    #   librosa
+    #   mistral-common
+sqlalchemy==2.0.41
+    # via
+    #   alembic
+    #   mlflow
 sqlitedict==2.1.0
    # via lm-eval
+sqlparse==0.5.3
+    # via mlflow-skinny
 starlette==0.46.2
    # via
+    #   fastapi
    #   schemathesis
    #   starlette-testclient
 starlette-testclient==0.4.1
@@ -751,18 +1035,29 @@ tenacity==9.0.0
    # via
    #   lm-eval
    #   plotly
+tensorboardx==2.6.4
+    # via lightning
 tensorizer==2.10.1
    # via -r requirements/test.in
+terratorch==1.1rc2
+    # via -r requirements/test.in
 threadpoolctl==3.5.0
    # via scikit-learn
+tifffile==2025.3.30
+    # via
+    #   scikit-image
+    #   terratorch
 tiktoken==0.7.0
    # via
    #   lm-eval
    #   mistral-common
-timm==1.0.11
+timm==1.0.15
    # via
    #   -r requirements/test.in
    #   open-clip-torch
+    #   segmentation-models-pytorch
+    #   terratorch
+    #   torchgeo
 tokenizers==0.21.1
    # via
    #   -r requirements/test.in
@@ -776,18 +1071,28 @@ torch==2.7.1+cu128
    #   -r requirements/test.in
    #   accelerate
    #   bitsandbytes
+    #   efficientnet-pytorch
    #   encodec
    #   fastsafetensors
+    #   kornia
+    #   lightly
+    #   lightning
    #   lm-eval
    #   mamba-ssm
    #   mteb
    #   open-clip-torch
    #   peft
+    #   pretrainedmodels
+    #   pytorch-lightning
    #   runai-model-streamer
+    #   segmentation-models-pytorch
    #   sentence-transformers
    #   tensorizer
+    #   terratorch
    #   timm
    #   torchaudio
+    #   torchgeo
+    #   torchmetrics
    #   torchvision
    #   vector-quantize-pytorch
    #   vocos
@@ -796,22 +1101,40 @@ torchaudio==2.7.1+cu128
    #   -r requirements/test.in
    #   encodec
    #   vocos
+torchgeo==0.7.0
+    # via terratorch
+torchmetrics==1.7.4
+    # via
+    #   lightning
+    #   pytorch-lightning
+    #   terratorch
+    #   torchgeo
 torchvision==0.22.1+cu128
    # via
    #   -r requirements/test.in
+    #   lightly
    #   open-clip-torch
+    #   pretrainedmodels
+    #   segmentation-models-pytorch
+    #   terratorch
    #   timm
+    #   torchgeo
 tqdm==4.66.6
    # via
    #   datasets
    #   evaluate
    #   huggingface-hub
+    #   lightly
+    #   lightning
    #   lm-eval
    #   mteb
    #   nltk
    #   open-clip-torch
    #   peft
    #   pqdm
+    #   pretrainedmodels
+    #   pytorch-lightning
+    #   segmentation-models-pytorch
    #   sentence-transformers
    #   tqdm-multiprocess
    #   transformers
@@ -829,7 +1152,9 @@ transformers==4.53.2
 transformers-stream-generator==0.0.5
    # via -r requirements/test.in
 triton==3.3.1
-    # via torch
+    # via
+    #   mamba-ssm
+    #   torch
 tritonclient==2.51.0
    # via
    #   -r requirements/test.in
@@ -843,18 +1168,34 @@ typer==0.15.2
    # via fastsafetensors
 types-python-dateutil==2.9.0.20241206
    # via arrow
+typeshed-client==2.8.2
+    # via jsonargparse
 typing-extensions==4.12.2
    # via
+    #   albumentations
+    #   alembic
+    #   fastapi
+    #   graphene
    #   huggingface-hub
    #   librosa
+    #   lightning
+    #   lightning-utilities
    #   mistral-common
+    #   mlflow-skinny
    #   mteb
+    #   opentelemetry-api
+    #   opentelemetry-sdk
+    #   opentelemetry-semantic-conventions
    #   pqdm
    #   pydantic
    #   pydantic-core
    #   pydantic-extra-types
+    #   pytorch-lightning
+    #   sqlalchemy
    #   torch
+    #   torchgeo
    #   typer
+    #   typeshed-client
    #   typing-inspection
 typing-inspection==0.4.1
    # via pydantic
@@ -866,9 +1207,13 @@ urllib3==2.2.3
    # via
    #   blobfile
    #   botocore
+    #   docker
+    #   lightly
    #   requests
    #   responses
    #   tritonclient
+uvicorn==0.35.0
+    # via mlflow-skinny
 vector-quantize-pytorch==1.21.2
    # via -r requirements/test.in
 virtualenv==20.31.2
@@ -880,11 +1225,15 @@ wcwidth==0.2.13
 webcolors==24.11.1
    # via jsonschema
 werkzeug==3.1.3
-    # via schemathesis
+    # via
+    #   flask
+    #   schemathesis
 word2number==1.1
    # via lm-eval
 wrapt==1.17.2
    # via smart-open
+xarray==2025.7.1
+    # via rioxarray
 xxhash==3.5.0
    # via
    #   datasets
@@ -893,5 +1242,7 @@ yarl==1.17.1
    # via
    #   aiohttp
    #   schemathesis
+zipp==3.23.0
+    # via importlib-metadata
 zstandard==0.23.0
    # via lm-eval
@@ -10,6 +10,7 @@ jinja2>=3.1.6
 ray[default]
 ray[data]
 setuptools==78.1.0
+nixl==0.3.0

 # Install torch_xla
 --pre
@@ -0,0 +1,350 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Test (piecewise) compilation with a simple model where multiple submodules
+are compiled and graph captured separately.
+"""
+import torch
+from torch import nn
+from torch.library import Library
+
+from vllm.compilation.backends import set_model_tag
+from vllm.compilation.counter import compilation_counter
+from vllm.compilation.decorators import (ignore_torch_compile,
+                                         support_torch_compile)
+from vllm.config import (CompilationConfig, CompilationLevel, VllmConfig,
+                         set_current_vllm_config)
+from vllm.envs import VLLM_USE_V1
+from vllm.forward_context import set_forward_context
+from vllm.utils import direct_register_custom_op
+
+# create a library to hold the custom op
+silly_lib = Library("silly", "FRAGMENT")  # noqa
+
+BATCH_SIZE = 32
+MLP_SIZE = 128
+HIDDEN_SIZE = 1024
+RANDOM_SEED = 0
+
+
+def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                    out: torch.Tensor) -> None:
+    out.copy_(q)
+    out += k
+    out += v
+
+
+def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                         out: torch.Tensor) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="attention",
+    op_func=silly_attention,
+    mutates_args=["out"],
+    fake_impl=silly_attention_fake,
+    target_lib=silly_lib,
+)
+
+
+@support_torch_compile
+class ParentModel(nn.Module):
+
+    def __init__(self,
+                 *,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+
+class Attention(nn.Module):
+
+    def __init__(self, mlp_size: int, hidden_size: int) -> None:
+        super().__init__()
+        self.pre_attn = nn.Linear(mlp_size, hidden_size, bias=False)
+        self.post_attn = nn.Linear(hidden_size, mlp_size, bias=False)
+        self.rms_norm_weight = nn.Parameter(torch.ones(hidden_size))
+
+        # Initialize to same weights for testing
+        nn.init.xavier_normal_(
+            self.pre_attn.weight.data,
+            generator=torch.Generator().manual_seed(RANDOM_SEED),
+            gain=0.001)
+        nn.init.xavier_normal_(
+            self.post_attn.weight.data,
+            generator=torch.Generator().manual_seed(RANDOM_SEED),
+            gain=0.001)
+
+    def rms_norm_ref(self, x: torch.Tensor) -> torch.Tensor:
+        x_f32 = x.float()
+        return (x_f32 * torch.rsqrt(
+            torch.mean(x_f32.square(), dim=-1, keepdim=True) + 1e-6) *
+                self.rms_norm_weight).to(x.dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.pre_attn(x)
+        x = self.rms_norm_ref(x)
+        attn_output = torch.empty_like(x)
+        torch.ops.silly.attention(x, x, x, attn_output)
+        x = attn_output
+        x = self.rms_norm_ref(x)
+        x = self.post_attn(x)
+        return x
+
+
+@support_torch_compile
+class CompiledAttention(nn.Module):
+
+    def __init__(self,
+                 *,
+                 mlp_size: int,
+                 hidden_size: int,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
+        super().__init__()
+        self.attn = Attention(mlp_size, hidden_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.attn(x)
+
+
+@support_torch_compile
+class CompiledAttentionTwo(CompiledAttention):
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.attn(x) + x
+
+
+@ignore_torch_compile
+class SimpleModelWithTwoGraphs(ParentModel):
+
+    def __init__(self,
+                 *,
+                 mlp_size: int,
+                 hidden_size: int,
+                 vllm_config: VllmConfig,
+                 prefix: str = '',
+                 **kwargs) -> None:
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        # Test will fail without set_model_tag here with error:
+        # "ValueError: too many values to unpack (expected 3)"
+        # This is because CompiledAttention and CompiledAttentionTwo
+        # have different implmentations but the same torch.compile
+        # cache dir will be used as default prefix is 'model_tag'
+        with set_model_tag("attn_one"):
+            self.attn_one = CompiledAttention(
+                mlp_size=mlp_size,
+                hidden_size=hidden_size,
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.attn_one",
+            )
+        with set_model_tag("attn_two"):
+            self.attn_two = CompiledAttentionTwo(
+                mlp_size=mlp_size,
+                hidden_size=hidden_size,
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.attn_two",
+            )
+
+        self.hidden_states = torch.zeros((BATCH_SIZE, MLP_SIZE)).cuda()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        bsz = x.shape[0]
+        # CUDAGraph expects same tensor addresses for each run
+        self.hidden_states[:bsz].copy_(x)
+        x = self.attn_one(self.hidden_states[:bsz])
+        self.hidden_states[:bsz].copy_(x)
+        x = self.attn_two(self.hidden_states[:bsz])
+        return x
+
+
+def test_ignore_torch_compile_decorator():
+    assert VLLM_USE_V1
+
+    # piecewise
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        use_cudagraph=True,
+        splitting_ops=["silly.attention"],
+        cudagraph_capture_sizes=[1, 2],
+    ))
+
+    @support_torch_compile
+    class A(nn.Module):
+
+        def __init__(self,
+                     *,
+                     vllm_config: VllmConfig,
+                     prefix: str = '',
+                     **kwargs) -> None:
+            super().__init__()
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            x = x + x
+            attn_output = torch.empty_like(x)
+            torch.ops.silly.attention(x, x, x, attn_output)
+            x = attn_output
+            x = x * 3
+            return x
+
+    @ignore_torch_compile
+    class B(A):
+        ...
+
+    @support_torch_compile
+    class C(B):
+        ...
+
+    with set_current_vllm_config(vllm_config):
+        mod_A = A(vllm_config=vllm_config, prefix='').eval().cuda()
+
+    # A has support_torch_compile
+    with compilation_counter.expect(
+            num_graphs_seen=1,
+            num_piecewise_graphs_seen=3,
+            num_piecewise_capturable_graphs_seen=2,
+            num_backend_compilations=2,
+            num_cudagraph_captured=4,
+            # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+    ), set_forward_context({}, vllm_config=vllm_config):
+        # first run is for compile
+        mod_A(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
+        # run cudagraph captured sizes
+        mod_A(torch.randn(2, MLP_SIZE).cuda())
+        mod_A(torch.randn(1, MLP_SIZE).cuda())
+
+    with set_current_vllm_config(vllm_config):
+        mod_B = B(vllm_config=vllm_config, prefix='').eval().cuda()
+
+    # B's ignore_torch_compile should override A's support_torch_compile
+    with compilation_counter.expect(
+            num_graphs_seen=0,
+            num_piecewise_graphs_seen=0,
+            num_piecewise_capturable_graphs_seen=0,
+            num_backend_compilations=0,
+            num_cudagraph_captured=0,
+    ), set_forward_context({}, vllm_config=vllm_config):
+        mod_B(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
+        mod_B(torch.randn(2, MLP_SIZE).cuda())
+        mod_B(torch.randn(1, MLP_SIZE).cuda())
+
+    with set_current_vllm_config(vllm_config):
+        mod_C = C(vllm_config=vllm_config, prefix='').eval().cuda()
+
+    # C's support_torch_compile should override B's ignore_torch_compile
+    with compilation_counter.expect(
+            num_graphs_seen=1,
+            num_piecewise_graphs_seen=3,
+            num_piecewise_capturable_graphs_seen=2,
+            num_backend_compilations=2,
+            num_cudagraph_captured=4,
+            # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+    ), set_forward_context({}, vllm_config=vllm_config):
+        mod_C(torch.randn(BATCH_SIZE, MLP_SIZE).cuda())
+        mod_C(torch.randn(2, MLP_SIZE).cuda())
+        mod_C(torch.randn(1, MLP_SIZE).cuda())
+
+
+@torch.inference_mode
+def run_model(vllm_config, model: nn.Module, inputs: torch.Tensor):
+    with set_forward_context({}, vllm_config=vllm_config):
+        # First run is for compile
+        model(inputs)
+
+        # Run CUDAGraph captured sizes
+        model(inputs[:2])
+        model(inputs[:1])
+
+        output = model(inputs[:2])
+
+        output = output.cpu()
+        return output.cpu()
+
+
+def test_multi_graph_piecewise_compile_outputs_equal():
+    outputs = []
+
+    # piecewise compile
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        use_cudagraph=True,
+        splitting_ops=["silly.attention"],
+        cudagraph_capture_sizes=[1, 2],
+    ))
+
+    with set_current_vllm_config(vllm_config):
+        model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
+                                         hidden_size=HIDDEN_SIZE,
+                                         vllm_config=vllm_config,
+                                         prefix='').eval().cuda()
+
+    # Pre-allocate memory for CUDAGraph which expects
+    # static tensor addresses
+    inputs = torch.randn(BATCH_SIZE, MLP_SIZE).cuda()
+
+    with compilation_counter.expect(
+            num_graphs_seen=2,  # two graphs for the model
+            num_piecewise_graphs_seen=6,
+            # attn_one, attn_two each has 3 piecewise graphs
+            # (pre attn, post attn, silly_attention) each
+            num_piecewise_capturable_graphs_seen=4,
+            # attn_one, attn_two has pre attn and post attn each, total=4
+            num_backend_compilations=4,  # num_piecewise_capturable_graphs_seen
+            num_cudagraph_captured=8,
+            # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
+    ):
+        outputs.append(run_model(vllm_config, model, inputs))
+
+    # no compile or cudagraph
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.NO_COMPILATION, ))
+
+    with set_current_vllm_config(vllm_config):
+        model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
+                                         hidden_size=HIDDEN_SIZE,
+                                         vllm_config=vllm_config,
+                                         prefix='').eval().cuda()
+
+    with compilation_counter.expect(
+            num_graphs_seen=0,
+            num_piecewise_graphs_seen=0,
+            num_piecewise_capturable_graphs_seen=0,
+            num_backend_compilations=0,
+            num_cudagraph_captured=0,
+    ):
+        outputs.append(run_model(vllm_config, model, inputs))
+
+    # piecewise compile without CUDA graph
+    vllm_config = VllmConfig(compilation_config=CompilationConfig(
+        level=CompilationLevel.PIECEWISE,
+        use_cudagraph=False,
+        splitting_ops=["silly.attention"],
+    ))
+
+    with set_current_vllm_config(vllm_config):
+        model = SimpleModelWithTwoGraphs(mlp_size=MLP_SIZE,
+                                         hidden_size=HIDDEN_SIZE,
+                                         vllm_config=vllm_config,
+                                         prefix='').eval().cuda()
+
+    with compilation_counter.expect(
+            num_graphs_seen=2,
+            num_piecewise_graphs_seen=6,
+            num_piecewise_capturable_graphs_seen=4,
+            num_backend_compilations=4,
+            num_cudagraph_captured=0,  # no cudagraph captured
+    ):
+        outputs.append(run_model(vllm_config, model, inputs))
+
+    # Generally don't expect outputs with and without inductor
+    # to be bitwise equivalent
+    assert torch.allclose(outputs[0], outputs[1])
+
+    # Expect bitwise equivalence using inductor w/ and w/o cudagraph
+    assert torch.equal(outputs[0], outputs[2])
@@ -73,9 +73,6 @@ def test_lm_eval_accuracy_v1_engine(model, monkeypatch: pytest.MonkeyPatch):
        if current_platform.is_tpu():
            # Limit compilation time for TPU V1

-            # xet doesn't work well for both Qwen/Qwen3-1.7B and
-            # google/gemma-3-1b-it
-            m.setenv("HF_HUB_DISABLE_XET", "1")
            more_args = "max_model_len=2048,max_num_seqs=64"

            # Add TP test (if provided)
@@ -69,8 +69,9 @@ def run_test(more_args):


@pytest.mark.skipif(not current_platform.is_cuda()
-                    and not current_platform.is_tpu(),
-                    reason="V1 currently only supported on CUDA and TPU")
+                    and not current_platform.is_tpu()
+                    and not current_platform.is_xpu(),
+                    reason="V1 currently only supported on CUDA, XPU and TPU")
 def test_lm_eval_accuracy_v1_engine(monkeypatch: pytest.MonkeyPatch):
    """Run with the V1 Engine."""

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # imports for guided decoding tests
 import json
+import os
 import shutil
 from tempfile import TemporaryDirectory
 from typing import Optional
@@ -26,10 +27,6 @@ MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
 # technically these adapters use a different base model,
 # but we're not testing generation quality here
 LORA_NAME = "typeof/zephyr-7b-beta-lora"
-PA_NAME = "swapnilbp/llama_tweet_ptune"
-# if PA_NAME changes, PA_NUM_VIRTUAL_TOKENS might also
-# need to change to match the prompt adapter
-PA_NUM_VIRTUAL_TOKENS = 8

 GUIDED_DECODING_BACKENDS = ["outlines", "lm-format-enforcer", "xgrammar"]

@@ -56,13 +53,7 @@ def zephyr_lora_added_tokens_files(zephyr_lora_files):


@pytest.fixture(scope="module")
-def zephyr_pa_files():
-    return snapshot_download(repo_id=PA_NAME)
-
-
-@pytest.fixture(scope="module")
-def default_server_args(zephyr_lora_files, zephyr_lora_added_tokens_files,
-                        zephyr_pa_files):
+def default_server_args(zephyr_lora_files, zephyr_lora_added_tokens_files):
    return [
        # use half precision for speed and memory savings in CI environment
        "--dtype",
@@ -81,15 +72,6 @@ def default_server_args(zephyr_lora_files, zephyr_lora_added_tokens_files,
        "64",
        "--max-cpu-loras",
        "2",
-        # pa config
-        "--enable-prompt-adapter",
-        "--prompt-adapters",
-        f"zephyr-pa={zephyr_pa_files}",
-        f"zephyr-pa2={zephyr_pa_files}",
-        "--max-prompt-adapters",
-        "2",
-        "--max-prompt-adapter-token",
-        "128",
    ]


@@ -98,8 +80,19 @@ def default_server_args(zephyr_lora_files, zephyr_lora_added_tokens_files,
 def server(default_server_args, request):
    if request.param:
        default_server_args.append(request.param)
-    with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
-        yield remote_server
+
+    original_value = os.environ.get('VLLM_USE_V1')
+    os.environ['VLLM_USE_V1'] = '0'
+    try:
+        with RemoteOpenAIServer(MODEL_NAME,
+                                default_server_args) as remote_server:
+            yield remote_server
+    finally:
+        # Restore original env value
+        if original_value is None:
+            os.environ.pop('VLLM_USE_V1', None)
+        else:
+            os.environ['VLLM_USE_V1'] = original_value


@pytest_asyncio.fixture
@@ -110,14 +103,11 @@ async def client(server):

@pytest.mark.asyncio
@pytest.mark.parametrize(
-    # first test base model, then test loras, then test prompt adapters
-    "model_name,num_virtual_tokens",
-    [(MODEL_NAME, 0), ("zephyr-lora", 0), ("zephyr-lora2", 0),
-     ("zephyr-pa", PA_NUM_VIRTUAL_TOKENS),
-     ("zephyr-pa2", PA_NUM_VIRTUAL_TOKENS)],
+    # first test base model, then test loras
+    "model_name",
+    [MODEL_NAME, "zephyr-lora", "zephyr-lora2"],
 )
-async def test_single_completion(client: openai.AsyncOpenAI, model_name: str,
-                                 num_virtual_tokens: int):
+async def test_single_completion(client: openai.AsyncOpenAI, model_name: str):
    completion = await client.completions.create(model=model_name,
                                                 prompt="Hello, my name is",
                                                 max_tokens=5,
@@ -130,9 +120,7 @@ async def test_single_completion(client: openai.AsyncOpenAI, model_name: str,
    assert len(choice.text) >= 5
    assert choice.finish_reason == "length"
    assert completion.usage == openai.types.CompletionUsage(
-        completion_tokens=5,
-        prompt_tokens=6 + num_virtual_tokens,
-        total_tokens=11 + num_virtual_tokens)
+        completion_tokens=5, prompt_tokens=6, total_tokens=11)

    # test using token IDs
    completion = await client.completions.create(
@@ -175,9 +163,9 @@ async def test_added_lora_tokens_base_model(client: openai.AsyncOpenAI):

@pytest.mark.asyncio
@pytest.mark.parametrize(
-    # first test base model, then test loras, then test prompt adapters
+    # first test base model, then test loras
    "model_name",
-    [MODEL_NAME, "zephyr-lora", "zephyr-lora2", "zephyr-pa", "zephyr-pa2"],
+    [MODEL_NAME, "zephyr-lora", "zephyr-lora2"],
 )
 async def test_no_logprobs(client: openai.AsyncOpenAI, model_name: str):
    # test using token IDs
@@ -194,9 +182,9 @@ async def test_no_logprobs(client: openai.AsyncOpenAI, model_name: str):

@pytest.mark.asyncio
@pytest.mark.parametrize(
-    # just test 1 lora and 1 pa hereafter
+    # just test 1 lora
    "model_name",
-    [MODEL_NAME, "zephyr-lora", "zephyr-pa"],
+    [MODEL_NAME, "zephyr-lora"],
 )
 async def test_zero_logprobs(client: openai.AsyncOpenAI, model_name: str):
    # test using token IDs
@@ -217,7 +205,7 @@ async def test_zero_logprobs(client: openai.AsyncOpenAI, model_name: str):
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "model_name",
-    [MODEL_NAME, "zephyr-lora", "zephyr-pa"],
+    [MODEL_NAME, "zephyr-lora"],
 )
 async def test_some_logprobs(client: openai.AsyncOpenAI, model_name: str):
    # test using token IDs
@@ -238,7 +226,7 @@ async def test_some_logprobs(client: openai.AsyncOpenAI, model_name: str):
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "model_name",
-    [MODEL_NAME, "zephyr-lora", "zephyr-pa"],
+    [MODEL_NAME, "zephyr-lora"],
 )
 async def test_too_many_completion_logprobs(client: openai.AsyncOpenAI,
                                            model_name: str):
@@ -314,7 +302,7 @@ async def test_prompt_logprobs_completion(client: openai.AsyncOpenAI,
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "model_name",
-    [MODEL_NAME, "zephyr-lora", "zephyr-pa"],
+    [MODEL_NAME, "zephyr-lora"],
 )
 async def test_completion_streaming(client: openai.AsyncOpenAI,
                                    model_name: str):
@@ -348,7 +336,7 @@ async def test_completion_streaming(client: openai.AsyncOpenAI,
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "model_name",
-    [MODEL_NAME, "zephyr-lora", "zephyr-pa"],
+    [MODEL_NAME, "zephyr-lora"],
 )
 async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):
    """Streaming for parallel sampling.
@@ -382,7 +370,7 @@ async def test_parallel_streaming(client: openai.AsyncOpenAI, model_name: str):
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "model_name",
-    [MODEL_NAME, "zephyr-lora", "zephyr-pa"],
+    [MODEL_NAME, "zephyr-lora"],
 )
 async def test_completion_stream_options(client: openai.AsyncOpenAI,
                                         model_name: str):
@@ -519,7 +507,7 @@ async def test_completion_stream_options(client: openai.AsyncOpenAI,
@pytest.mark.asyncio
@pytest.mark.parametrize(
    "model_name",
-    [MODEL_NAME, "zephyr-lora", "zephyr-pa"],
+    [MODEL_NAME, "zephyr-lora"],
 )
 async def test_batch_completions(client: openai.AsyncOpenAI, model_name: str):
    # test both text and token IDs
@@ -295,8 +295,6 @@ async def test_metrics_exist(server: RemoteOpenAIServer,


 def test_metrics_exist_run_batch(use_v1: bool):
-    if use_v1:
-        pytest.skip("Skipping test on vllm V1")
    input_batch = """{"custom_id": "request-0", "method": "POST", "url": "/v1/embeddings", "body": {"model": "intfloat/multilingual-e5-small", "input": "You are a helpful assistant."}}"""  # noqa: E501

    base_url = "0.0.0.0"
@@ -323,7 +321,8 @@ def test_metrics_exist_run_batch(use_v1: bool):
            base_url,
            "--port",
            port,
-        ], )
+        ],
+                                env={"VLLM_USE_V1": "1" if use_v1 else "0"})

        def is_server_up(url):
            try:
@@ -13,7 +13,6 @@ from ...utils import RemoteOpenAIServer
 from .test_completion import default_server_args  # noqa: F401
 from .test_completion import zephyr_lora_added_tokens_files  # noqa: F401
 from .test_completion import zephyr_lora_files  # noqa: F401
-from .test_completion import zephyr_pa_files  # noqa: F401
 from .test_completion import MODEL_NAME


@@ -32,8 +32,7 @@ async def _async_serving_models_init() -> OpenAIServingModels:
    serving_models = OpenAIServingModels(engine_client=mock_engine_client,
                                         base_model_paths=BASE_MODEL_PATHS,
                                         model_config=mock_model_config,
-                                         lora_modules=None,
-                                         prompt_adapters=None)
+                                         lora_modules=None)
    await serving_models.init_static_loras()

    return serving_models
@@ -6,6 +6,10 @@ from collections.abc import Mapping
 from typing import Literal, Optional

 import pytest
+from mistral_common.tokens.tokenizers.base import (SpecialTokenPolicy,
+                                                   SpecialTokens)
+from mistral_common.tokens.tokenizers.tekken import (SpecialTokenInfo,
+                                                     Tekkenizer)

 from vllm.assets.audio import AudioAsset
 from vllm.assets.image import ImageAsset
@@ -21,6 +25,7 @@ from vllm.multimodal import MultiModalDataDict
 from vllm.multimodal.utils import (encode_audio_base64, encode_image_base64,
                                   encode_video_base64)
 from vllm.transformers_utils.tokenizer_group import TokenizerGroup
+from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer

 from ..models.registry import HF_EXAMPLE_MODELS
 from ..utils import VLLM_PATH
@@ -1374,3 +1379,165 @@ def test_resolve_content_format_examples(template_path, expected_format):
    )

    assert resolved_format == expected_format
+
+
+def test_parse_chat_messages_include_thinking_chunk(mistral_model_config,
+                                                    mistral_tokenizer):
+    messages = [{
+        "role":
+        "system",
+        "content": [{
+            "type": "text",
+            "text": "You are a helpful assistant."
+        }, {
+            "type":
+            "thinking",
+            "closed":
+            True,
+            "thinking":
+            "Only return the answer when you are confident."
+        }]
+    }, {
+        "role": "user",
+        "content": "What is 2+2?"
+    }, {
+        "role":
+        "assistant",
+        "content": [{
+            "type": "text",
+            "text": "Let me think about it."
+        }, {
+            "type": "thinking",
+            "closed": True,
+            "thinking": "2+2 = 4"
+        }, {
+            "type": "text",
+            "text": "The answer is 4.",
+        }],
+    }]
+
+    conversation_with_thinking, _ = parse_chat_messages(
+        messages,
+        mistral_model_config,
+        mistral_tokenizer,
+        content_format="openai",
+    )
+
+    expected_conversation = [{
+        "role":
+        "system",
+        "content": [{
+            "type": "text",
+            "text": "You are a helpful assistant."
+        }, {
+            "type": "text",
+            "text": "Only return the answer when you are confident."
+        }],
+    }, {
+        "role":
+        "user",
+        "content": [{
+            "type": "text",
+            "text": "What is 2+2?"
+        }],
+    }, {
+        "role":
+        "assistant",
+        "content": [
+            {
+                "type": "text",
+                "text": "Let me think about it."
+            },
+            {
+                "type": "text",
+                "text": "2+2 = 4"
+            },
+            {
+                "type": "text",
+                "text": "The answer is 4."
+            },
+        ]
+    }]
+
+    assert conversation_with_thinking == expected_conversation
+
+
+def test_apply_mistral_chat_template_thinking_chunk():
+    # Moved import here to avoid yapf and isort conflicts
+    from vllm.entrypoints.chat_utils import apply_mistral_chat_template
+    messages = [{
+        "role":
+        "system",
+        "content": [{
+            "type": "text",
+            "text": "You are a helpful assistant."
+        }, {
+            "type":
+            "thinking",
+            "closed":
+            True,
+            "thinking":
+            "Only return the answer when you are confident."
+        }]
+    }, {
+        "role": "user",
+        "content": "What is 2+2?"
+    }, {
+        "role":
+        "assistant",
+        "content": [{
+            "type": "text",
+            "text": "Let me think about it."
+        }, {
+            "type": "thinking",
+            "closed": True,
+            "thinking": "2+2 = 4"
+        }, {
+            "type": "text",
+            "text": "The answer is 4.",
+        }],
+    }, {
+        "role": "user",
+        "content": "Thanks, what is 3+3?"
+    }]
+
+    # TODO(Julien): upon model release change to a tokenizer already configured.
+    # =================================================================
+    mistral_tokenizer = MistralTokenizer.from_pretrained(
+        "mistralai/Devstral-Small-2507")
+    assert isinstance(mistral_tokenizer.tokenizer, Tekkenizer)
+    # Add think special tokens to the tokenizer
+    mistral_tokenizer.tokenizer._all_special_tokens[35] = SpecialTokenInfo(
+        rank=35, is_control=True, token_str=SpecialTokens.begin_think.value)
+    mistral_tokenizer.tokenizer._all_special_tokens[36] = SpecialTokenInfo(
+        rank=36, is_control=True, token_str=SpecialTokens.end_think.value)
+    mistral_tokenizer.tokenizer._special_tokens_reverse_vocab = {
+        k: v
+        for k, v in
+        mistral_tokenizer.tokenizer._special_tokens_reverse_vocab.items()
+        if v not in {35, 36}
+    }
+    mistral_tokenizer.tokenizer._special_tokens_reverse_vocab[
+        SpecialTokens.begin_think.value] = 35
+    mistral_tokenizer.tokenizer._special_tokens_reverse_vocab[
+        SpecialTokens.end_think.value] = 36
+    mistral_tokenizer.instruct.BEGIN_THINK = 35
+    mistral_tokenizer.instruct.END_THINK = 36
+    # =================================================================
+
+    tokens_ids = apply_mistral_chat_template(mistral_tokenizer,
+                                             messages,
+                                             chat_template=None,
+                                             tools=None)
+
+    string_tokens = mistral_tokenizer.mistral.decode(
+        tokens_ids, special_token_policy=SpecialTokenPolicy.KEEP)
+
+    expected_tokens = (
+        r"<s>[SYSTEM_PROMPT]You are a helpful assistant.[THINK]Only return the"
+        r" answer when you are confident.[/THINK][/SYSTEM_PROMPT]"
+        r"[INST]What is 2+2?[/INST]"
+        r"Let me think about it.[THINK]2+2 = 4[/THINK]The answer is 4.</s>"
+        r"[INST]Thanks, what is 3+3?[/INST]")
+
+    assert string_tokens == expected_tokens
@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 from vllm import SamplingParams
-from vllm.config import LoadFormat

 test_model = "openai-community/gpt2"

@@ -17,7 +16,6 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)


 def test_model_loader_download_files(vllm_runner):
-    with vllm_runner(test_model,
-                     load_format=LoadFormat.FASTSAFETENSORS) as llm:
+    with vllm_runner(test_model, load_format="fastsafetensors") as llm:
        deserialized_outputs = llm.generate(prompts, sampling_params)
        assert deserialized_outputs
@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the shuffle_rows function
+
+Run `pytest tests/kernels/test_shuffle_rows.py`.
+"""
+
+import pytest
+import torch
+
+from vllm._custom_ops import shuffle_rows
+from vllm.platforms import current_platform
+
+
+@pytest.mark.parametrize("num_tokens", [1, 16, 64, 128, 256, 512, 1024])
+@pytest.mark.parametrize("hidden_size", [128, 256, 512, 1024, 2048, 4096])
+@pytest.mark.parametrize("dtype",
+                         [torch.float16, torch.bfloat16, torch.float32])
+def test_shuffle_rows_basic(num_tokens: int, hidden_size: int,
+                            dtype: torch.dtype):
+    """Test basic functionality of shuffle_rows with various tensor sizes and
+    dtypes."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    # Create input tensor
+    input_tensor = torch.randn(num_tokens,
+                               hidden_size,
+                               device="cuda",
+                               dtype=dtype)
+
+    # Create a simple permutation map (identity mapping)
+    dst2src_map = torch.arange(num_tokens, device="cuda", dtype=torch.int32)
+
+    # Test shuffle_rows
+    output = shuffle_rows(input_tensor, dst2src_map)
+
+    # With identity mapping, output should be identical to input
+    torch.testing.assert_close(output, input_tensor, atol=0, rtol=0)
+
+    # Check output shape
+    assert output.shape == (num_tokens, hidden_size)
+    assert output.dtype == dtype
+    assert output.device == input_tensor.device
+
+
+@pytest.mark.parametrize("num_tokens", [16, 64, 128])
+@pytest.mark.parametrize("hidden_size", [128, 512, 1024])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_shuffle_rows_permutation(num_tokens: int, hidden_size: int,
+                                  dtype: torch.dtype):
+    """Test shuffle_rows with actual permutation."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    # Create input tensor
+    input_tensor = torch.randn(num_tokens,
+                               hidden_size,
+                               device="cuda",
+                               dtype=dtype)
+
+    # Create a reverse permutation map
+    dst2src_map = torch.arange(num_tokens - 1,
+                               -1,
+                               -1,
+                               device="cuda",
+                               dtype=torch.int32)
+
+    # Test shuffle_rows
+    output = shuffle_rows(input_tensor, dst2src_map)
+
+    # Check that the output is the reverse of the input
+    expected_output = torch.flip(input_tensor, dims=[0])
+    torch.testing.assert_close(output, expected_output, atol=1e-6, rtol=1e-5)
+
+    # Check output shape and properties
+    assert output.shape == (num_tokens, hidden_size)
+    assert output.dtype == dtype
+    assert output.device == input_tensor.device
+
+
+@pytest.mark.parametrize("num_tokens", [32, 64])
+@pytest.mark.parametrize("hidden_size", [256, 512])
+def test_shuffle_rows_expansion(num_tokens: int, hidden_size: int):
+    """Test shuffle_rows with expansion (more output tokens than input
+    tokens)."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    dtype = torch.float16
+
+    # Create input tensor
+    input_tensor = torch.randn(num_tokens,
+                               hidden_size,
+                               device="cuda",
+                               dtype=dtype)
+
+    # Create a mapping that duplicates some tokens (expansion)
+    expanded_size = num_tokens * 2
+    dst2src_map = torch.randint(0,
+                                num_tokens, (expanded_size, ),
+                                device="cuda",
+                                dtype=torch.int32)
+
+    # Test shuffle_rows
+    output = shuffle_rows(input_tensor, dst2src_map)
+
+    # Check output shape
+    assert output.shape == (expanded_size, hidden_size)
+    assert output.dtype == dtype
+    assert output.device == input_tensor.device
+
+    # Verify that each output row matches the corresponding input row
+    for i in range(expanded_size):
+        src_idx = dst2src_map[i].item()
+        torch.testing.assert_close(output[i],
+                                   input_tensor[src_idx],
+                                   atol=1e-6,
+                                   rtol=1e-5)
+
+
+@pytest.mark.parametrize("num_tokens", [16, 64])
+@pytest.mark.parametrize("hidden_size", [128, 512])
+def test_shuffle_rows_random_permutation(num_tokens: int, hidden_size: int):
+    """Test shuffle_rows with random permutation."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    dtype = torch.float16
+
+    # Set seed for reproducibility
+    torch.manual_seed(42)
+
+    # Create input tensor
+    input_tensor = torch.randn(num_tokens,
+                               hidden_size,
+                               device="cuda",
+                               dtype=dtype)
+
+    # Create a random permutation map
+    dst2src_map = torch.randperm(num_tokens, device="cuda", dtype=torch.int32)
+
+    # Test shuffle_rows
+    output = shuffle_rows(input_tensor, dst2src_map)
+
+    # Check output shape and properties
+    assert output.shape == (num_tokens, hidden_size)
+    assert output.dtype == dtype
+    assert output.device == input_tensor.device
+
+    # Verify that each output row matches the corresponding input row
+    for i in range(num_tokens):
+        src_idx = dst2src_map[i].item()
+        torch.testing.assert_close(output[i],
+                                   input_tensor[src_idx],
+                                   atol=1e-6,
+                                   rtol=1e-5)
+
+
+def test_shuffle_rows_edge_cases():
+    """Test shuffle_rows with edge cases."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    dtype = torch.float16
+
+    # Test with single token
+    input_tensor = torch.randn(1, 128, device="cuda", dtype=dtype)
+    dst2src_map = torch.tensor([0], device="cuda", dtype=torch.int32)
+    output = shuffle_rows(input_tensor, dst2src_map)
+    torch.testing.assert_close(output, input_tensor, atol=0, rtol=0)
+
+    # Test with single feature dimension
+    input_tensor = torch.randn(16, 1, device="cuda", dtype=dtype)
+    dst2src_map = torch.arange(16, device="cuda", dtype=torch.int32)
+    output = shuffle_rows(input_tensor, dst2src_map)
+    torch.testing.assert_close(output, input_tensor, atol=0, rtol=0)
+
+
+def test_shuffle_rows_moe_like_scenario():
+    """Test shuffle_rows in a scenario similar to MoE usage."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    dtype = torch.float16
+    batch_size = 32
+    hidden_size = 1024
+    topk = 2
+
+    # Simulate input tokens
+    input_tensor = torch.randn(batch_size,
+                               hidden_size,
+                               device="cuda",
+                               dtype=dtype)
+
+    # Simulate expert assignment (each token goes to topk experts)
+    # This creates a mapping where tokens are duplicated for multiple experts
+    total_tokens = batch_size * topk
+    dst2src_map = torch.zeros(total_tokens, device="cuda", dtype=torch.int32)
+
+    # Fill the mapping to simulate MoE token distribution
+    for i in range(batch_size):
+        for k in range(topk):
+            dst2src_map[i * topk + k] = i
+
+    # Test shuffle_rows
+    output = shuffle_rows(input_tensor, dst2src_map)
+
+    # Check output shape
+    assert output.shape == (total_tokens, hidden_size)
+    assert output.dtype == dtype
+    assert output.device == input_tensor.device
+
+    # Verify that tokens are correctly duplicated
+    for i in range(batch_size):
+        for k in range(topk):
+            output_idx = i * topk + k
+            torch.testing.assert_close(output[output_idx],
+                                       input_tensor[i],
+                                       atol=1e-6,
+                                       rtol=1e-5)
+
+
+@pytest.mark.parametrize("dtype",
+                         [torch.float16, torch.bfloat16, torch.float32])
+def test_shuffle_rows_dtype_consistency(dtype: torch.dtype):
+    """Test that shuffle_rows preserves dtype correctly."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    num_tokens = 64
+    hidden_size = 512
+
+    # Create input tensor with specific dtype
+    input_tensor = torch.randn(num_tokens,
+                               hidden_size,
+                               device="cuda",
+                               dtype=dtype)
+    dst2src_map = torch.arange(num_tokens, device="cuda", dtype=torch.int32)
+
+    # Test shuffle_rows
+    output = shuffle_rows(input_tensor, dst2src_map)
+
+    # Verify dtype is preserved
+    assert output.dtype == dtype
+    assert output.device == input_tensor.device
+    torch.testing.assert_close(output, input_tensor, atol=1e-6, rtol=1e-5)
+
+
+def test_shuffle_rows_device_consistency():
+    """Test that shuffle_rows maintains device consistency."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    num_tokens = 32
+    hidden_size = 256
+    dtype = torch.float16
+
+    # Create input tensor on CUDA
+    input_tensor = torch.randn(num_tokens,
+                               hidden_size,
+                               device="cuda",
+                               dtype=dtype)
+    dst2src_map = torch.arange(num_tokens, device="cuda", dtype=torch.int32)
+
+    # Test shuffle_rows
+    output = shuffle_rows(input_tensor, dst2src_map)
+
+    # Verify device is maintained
+    assert output.device == input_tensor.device
+    assert output.device.type == "cuda"
+
+
+def test_shuffle_rows_contiguous_output():
+    """Test that shuffle_rows produces contiguous output."""
+    if not current_platform.is_cuda():
+        pytest.skip("shuffle_rows requires CUDA")
+
+    num_tokens = 64
+    hidden_size = 512
+    dtype = torch.float16
+
+    # Create input tensor
+    input_tensor = torch.randn(num_tokens,
+                               hidden_size,
+                               device="cuda",
+                               dtype=dtype)
+    dst2src_map = torch.arange(num_tokens, device="cuda", dtype=torch.int32)
+
+    # Test shuffle_rows
+    output = shuffle_rows(input_tensor, dst2src_map)
+
+    # Verify output is contiguous
+    assert output.is_contiguous()
@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from torch import nn
+
+from vllm.config import LoadConfig, ModelConfig
+from vllm.model_executor.model_loader import (get_model_loader,
+                                              register_model_loader)
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
+
+
+@register_model_loader("custom_load_format")
+class CustomModelLoader(BaseModelLoader):
+
+    def __init__(self, load_config: LoadConfig) -> None:
+        super().__init__(load_config)
+
+    def download_model(self, model_config: ModelConfig) -> None:
+        pass
+
+    def load_weights(self, model: nn.Module,
+                     model_config: ModelConfig) -> None:
+        pass
+
+
+def test_register_model_loader():
+    load_config = LoadConfig(load_format="custom_load_format")
+    assert isinstance(get_model_loader(load_config), CustomModelLoader)
+
+
+def test_invalid_model_loader():
+    with pytest.raises(ValueError):
+
+        @register_model_loader("invalid_load_format")
+        class InValidModelLoader:
+            pass
@@ -23,6 +23,8 @@ from transformers import (AutoConfig, AutoProcessor, AutoTokenizer,

 from vllm import LLM, SamplingParams

+from ....utils import multi_gpu_test
+
 # Sample prompts for testing
 PROMPTS: list[str] = [
    "Hello, my name is",
@@ -541,6 +543,7 @@ def run_reduced_model(model_path: str,
        print("-" * 40)


+@multi_gpu_test(num_gpus=2)
@pytest.mark.parametrize(
    "original_model_name,text_layers,num_experts,vision_layers,",
    [("meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 4, 4, 2)])
@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.utils import set_default_torch_num_threads
+
+from ....conftest import VllmRunner
+
+
+def generate_test_mm_data():
+    mm_data = {
+        "pixel_values": torch.full((6, 512, 512), 1.0, dtype=torch.float16),
+        "location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
+    }
+    return mm_data
+
+
+def _run_test(
+    vllm_runner: type[VllmRunner],
+    model: str,
+) -> None:
+
+    prompt = [
+        {
+            # This model deals with no text input
+            "prompt_token_ids": [1],
+            "multi_modal_data": generate_test_mm_data(),
+        } for _ in range(10)
+    ]
+
+    with (
+            set_default_torch_num_threads(1),
+            vllm_runner(
+                model,
+                task="embed",
+                dtype=torch.float16,
+                enforce_eager=True,
+                skip_tokenizer_init=True,
+                # Limit the maximum number of sequences to avoid the
+                # test going OOM during the warmup run
+                max_num_seqs=32,
+            ) as vllm_model,
+    ):
+        vllm_model.encode(prompt)
+
+
+MODELS = ["christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM"]
+
+
+@pytest.mark.core_model
+@pytest.mark.parametrize("model", MODELS)
+def test_models_image(
+    hf_runner,
+    vllm_runner,
+    image_assets,
+    model: str,
+) -> None:
+    _run_test(
+        vllm_runner,
+        model,
+    )
@@ -23,18 +23,14 @@ def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]:
        return ((name, torch.empty(0)) for name in weight_names)


-def create_model_dummy_weights(
-    repo: str,
-    model_arch: str,
-) -> Iterable[tuple[str, torch.Tensor]]:
+def create_dummy_model(repo: str, model_arch: str) -> PreTrainedModel:
    """
    Create weights from a dummy meta deserialized hf model with name conversion
    """
    model_cls: PreTrainedModel = getattr(transformers, model_arch)
    config = AutoConfig.from_pretrained(repo)
    with torch.device("meta"):
-        model: PreTrainedModel = model_cls._from_config(config)
-    return model.named_parameters()
+        return model_cls._from_config(config)


 def model_architectures_for_test() -> list[str]:
@@ -70,14 +66,21 @@ def test_hf_model_weights_mapper(model_arch: str):
    model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config)

    original_weights = create_repo_dummy_weights(model_id)
-    hf_converted_weights = create_model_dummy_weights(model_id, model_arch)
+    hf_dummy_model = create_dummy_model(model_id, model_arch)
+    hf_converted_weights = hf_dummy_model.named_parameters()
+    hf_converted_buffers = hf_dummy_model.named_buffers()
    mapper: WeightsMapper = model_cls.hf_to_vllm_mapper

    mapped_original_weights = mapper.apply(original_weights)
    mapped_hf_converted_weights = mapper.apply(hf_converted_weights)
+    mapped_hf_converted_buffers = mapper.apply(hf_converted_buffers)

    ref_weight_names = set(map(lambda x: x[0], mapped_original_weights))
    weight_names = set(map(lambda x: x[0], mapped_hf_converted_weights))
+    buffer_names = set(map(lambda x: x[0], mapped_hf_converted_buffers))
+
+    # Some checkpoints may have buffers, we ignore them for this test
+    ref_weight_names -= buffer_names

    weights_missing = ref_weight_names - weight_names
    weights_unmapped = weight_names - ref_weight_names
@@ -357,6 +357,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
                                                max_transformers_version="4.48",  # noqa: E501
                                                transformers_version_reason="HF model is not compatible.",  # noqa: E501
                                                hf_overrides={"architectures": ["DeepseekVLV2ForCausalLM"]}),  # noqa: E501
+    "Emu3ForConditionalGeneration": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),
    "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),
    "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),
    "GraniteSpeechForConditionalGeneration": _HfExamplesInfo("ibm-granite/granite-speech-3.3-2b"),  # noqa: E501
@@ -501,7 +502,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
                                    speculative_model="XiaomiMiMo/MiMo-7B-RL")
 }

-_TRANSFORMERS_MODELS = {
+_TRANSFORMERS_BACKEND_MODELS = {
    "TransformersForCausalLM": _HfExamplesInfo("hmellor/Ilama-3.2-1B", trust_remote_code=True),  # noqa: E501
    "TransformersForMultimodalLM": _HfExamplesInfo("OpenGVLab/InternVL3-1B-hf"),
 }
@@ -512,7 +513,7 @@ _EXAMPLE_MODELS = {
    **_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS,
    **_MULTIMODAL_EXAMPLE_MODELS,
    **_SPECULATIVE_DECODING_EXAMPLE_MODELS,
-    **_TRANSFORMERS_MODELS,
+    **_TRANSFORMERS_BACKEND_MODELS,
 }


@@ -1,48 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-
-import vllm
-from vllm.prompt_adapter.request import PromptAdapterRequest
-
-MODEL_PATH = "bigscience/bloomz-560m"
-PA_PATH = 'stevhliu/bloomz-560m_PROMPT_TUNING_CAUSAL_LM'
-
-
-def do_sample(llm, pa_name: str, pa_id: int):
-
-    prompts = [
-        "Tweet text : @nationalgridus I have no water and the bill is \
-        current and paid. Can you do something about this? Label : ",
-        "Tweet text : @nationalgridus Looks good thanks! Label : "
-    ]
-    sampling_params = vllm.SamplingParams(temperature=0.0,
-                                          max_tokens=3,
-                                          stop_token_ids=[3])
-
-    outputs = llm.generate(prompts,
-                           sampling_params,
-                           prompt_adapter_request=PromptAdapterRequest(
-                               pa_name, pa_id, PA_PATH, 8) if pa_id else None)
-
-    # Print the outputs.
-    generated_texts = []
-    for output in outputs:
-        prompt = output.prompt
-        generated_text = output.outputs[0].text.strip()
-        generated_texts.append(generated_text)
-        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
-    return generated_texts
-
-
-@pytest.mark.parametrize("enforce_eager", [True, False])
-def test_twitter_prompt_adapter(enforce_eager: bool):
-    llm = vllm.LLM(MODEL_PATH,
-                   enforce_eager=enforce_eager,
-                   enable_prompt_adapter=True,
-                   max_prompt_adapter_token=8)
-
-    expected_output = ['complaint', 'no complaint']
-
-    assert do_sample(llm, "twitter_pa", pa_id=1) == expected_output
@@ -1,56 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from vllm import EngineArgs, LLMEngine, SamplingParams
-from vllm.prompt_adapter.request import PromptAdapterRequest
-
-MODEL_PATH = "bigscience/bloomz-560m"
-pa_path = 'stevhliu/bloomz-560m_PROMPT_TUNING_CAUSAL_LM'
-pa_path2 = 'swapnilbp/angry_tweet_ptune'
-
-
-def do_sample(engine):
-
-    prompts = [
-        ("Tweet text: I have complaints! Label: ",
-         SamplingParams(temperature=0.0, max_tokens=3, stop_token_ids=[3]),
-         PromptAdapterRequest("hate_speech", 1, pa_path2, 8)),
-        ("Tweet text: I have no problems Label: ",
-         SamplingParams(temperature=0.0, max_tokens=3, stop_token_ids=[3]),
-         PromptAdapterRequest("hate_speech2", 2, pa_path2, 8)),
-        ("Tweet text: I have complaints! Label: ",
-         SamplingParams(temperature=0.0, max_tokens=3), None),
-        ("Tweet text: I have no problems Label: ",
-         SamplingParams(temperature=0.0, max_tokens=3, stop_token_ids=[3]),
-         PromptAdapterRequest("complain", 3, pa_path, 8)),
-    ]
-
-    request_id = 0
-    results = set()
-    while prompts or engine.has_unfinished_requests():
-        if prompts:
-            prompt, sampling_params, pa_request = prompts.pop(0)
-            engine.add_request(str(request_id),
-                               prompt,
-                               sampling_params,
-                               prompt_adapter_request=pa_request)
-            request_id += 1
-
-        request_outputs = engine.step()
-
-        for request_output in request_outputs:
-            if request_output.finished:
-                results.add(request_output.outputs[0].text)
-    return results
-
-
-def test_multi_prompt_adapters():
-    engine_args = EngineArgs(model=MODEL_PATH,
-                             max_prompt_adapters=3,
-                             enable_prompt_adapter=True,
-                             max_prompt_adapter_token=8)
-    engine = LLMEngine.from_engine_args(engine_args)
-    expected_output = {
-        ' quot;I', 'hate speech', 'no complaint', 'not hate speech'
-    }
-    assert do_sample(engine) == expected_output
@@ -1,64 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from huggingface_hub import snapshot_download
-
-from vllm import EngineArgs, LLMEngine, SamplingParams
-from vllm.lora.request import LoRARequest
-from vllm.prompt_adapter.request import PromptAdapterRequest
-
-MODEL_PATH = "meta-llama/Llama-2-7b-hf"
-pa_path = snapshot_download(repo_id="swapnilbp/llama_tweet_ptune")
-lora_path = snapshot_download(repo_id="yard1/llama-2-7b-sql-lora-test")
-
-
-def do_sample(engine):
-
-    prompt_text = "[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]"  # noqa: E501
-
-    # first prompt with a prompt adapter and second without adapter
-    prompts = [
-        (prompt_text,
-         SamplingParams(temperature=0.0, max_tokens=100,
-                        stop=["[/assistant]"]),
-         PromptAdapterRequest("hate_speech", 1, pa_path,
-                              8), LoRARequest("sql_test", 1, lora_path)),
-        (prompt_text,
-         SamplingParams(temperature=0.0, max_tokens=100,
-                        stop=["[/assistant]"]), None,
-         LoRARequest("sql_test", 1, lora_path)),
-    ]
-
-    request_id = 0
-    results = set()
-    while prompts or engine.has_unfinished_requests():
-        if prompts:
-            prompt, sampling_params, pa_request, lora_request = prompts.pop(0)
-            engine.add_request(str(request_id),
-                               prompt,
-                               sampling_params,
-                               prompt_adapter_request=pa_request,
-                               lora_request=lora_request)
-            request_id += 1
-
-        request_outputs = engine.step()
-
-        for request_output in request_outputs:
-            if request_output.finished:
-                results.add(request_output.outputs[0].text)
-    return results
-
-
-def test_lora_prompt_adapter():
-    engine_args = EngineArgs(model=MODEL_PATH,
-                             enable_prompt_adapter=True,
-                             enable_lora=True,
-                             max_num_seqs=60,
-                             max_prompt_adapter_token=8)
-    engine = LLMEngine.from_engine_args(engine_args)
-    result = do_sample(engine)
-
-    expected_output = {
-        "  SELECT icao FROM table_name_74 WHERE airport = 'lilongwe international airport' "  # noqa: E501
-    }
-    assert result == expected_output
@@ -0,0 +1,341 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from mistral_common.tokens.tokenizers.base import SpecialTokens
+from mistral_common.tokens.tokenizers.tekken import (SpecialTokenInfo,
+                                                     Tekkenizer)
+
+from tests.reasoning.utils import run_reasoning_extraction_mistral
+from vllm.reasoning import ReasoningParser, ReasoningParserManager
+from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer
+
+parser_name = "mistral"
+
+
+@pytest.fixture(scope="module")
+def mistral_tokenizer():
+    # TODO(Julien): upon model release change to a tokenizer already configured.
+    # =================================================================
+    mistral_tokenizer = MistralTokenizer.from_pretrained(
+        "mistralai/Devstral-Small-2507")
+    assert isinstance(mistral_tokenizer.tokenizer, Tekkenizer)
+    # Add think special tokens to the tokenizer
+    mistral_tokenizer.tokenizer._all_special_tokens[35] = SpecialTokenInfo(
+        rank=35, is_control=True, token_str=SpecialTokens.begin_think.value)
+    mistral_tokenizer.tokenizer._all_special_tokens[36] = SpecialTokenInfo(
+        rank=36, is_control=True, token_str=SpecialTokens.end_think.value)
+    mistral_tokenizer.tokenizer._special_tokens_reverse_vocab = {
+        k: v
+        for k, v in
+        mistral_tokenizer.tokenizer._special_tokens_reverse_vocab.items()
+        if v not in {35, 36}
+    }
+    mistral_tokenizer.tokenizer._special_tokens_reverse_vocab[
+        SpecialTokens.begin_think.value] = 35
+    mistral_tokenizer.tokenizer._special_tokens_reverse_vocab[
+        SpecialTokens.end_think.value] = 36
+    mistral_tokenizer.instruct.BEGIN_THINK = 35
+    mistral_tokenizer.instruct.END_THINK = 36
+    # =================================================================
+    return mistral_tokenizer
+
+
+SIMPLE_REASONING = {
+    "output": "This is a reasoning section[/THINK]This is the rest",
+    "reasoning_content": "This is a reasoning section",
+    "content": "This is the rest",
+    "is_reasoning_end": True,
+}
+COMPLETE_REASONING = {
+    "output": "This is a reasoning section[/THINK]",
+    "reasoning_content": "This is a reasoning section",
+    "content": None,
+    "is_reasoning_end": True,
+}
+NO_CONTENT = {
+    "output": "This is content",
+    "reasoning_content": "This is content",
+    "content": None,
+    "is_reasoning_end": False,
+}
+NO_REASONING_STREAMING = {
+    "output": "This is a reasoning section",
+    "reasoning_content": "This is a reasoning section",
+    "content": None,
+    "is_reasoning_end": False,
+}
+MULTIPLE_LINES = {
+    "output": "This\nThat[/THINK]This is the rest\nThat",
+    "reasoning_content": "This\nThat",
+    "content": "This is the rest\nThat",
+    "is_reasoning_end": True,
+}
+SHORTEST_REASONING_NO_STREAMING = {
+    "output": "[/THINK]This is the rest",
+    "reasoning_content": "",
+    "content": "This is the rest",
+    "is_reasoning_end": True,
+}
+SHORTEST_REASONING = {
+    "output": "[/THINK]This is the rest",
+    "reasoning_content": None,
+    "content": "This is the rest",
+    "is_reasoning_end": True,
+}
+REASONING_WITH_THINK = {
+    "output": "[THINK]This is a reasoning section[/THINK]This is the rest",
+    "reasoning_content": "This is a reasoning section",
+    "content": "This is the rest",
+    "is_reasoning_end": True,
+}
+COMPLETE_REASONING_WITH_THINK = {
+    "output": "[THINK]This is a reasoning section[/THINK]",
+    "reasoning_content": "This is a reasoning section",
+    "content": None,
+    "is_reasoning_end": True,
+}
+MULTIPLE_LINES_WITH_THINK = {
+    "output": "[THINK]This\nThat[/THINK]This is the rest\nThat",
+    "reasoning_content": "This\nThat",
+    "content": "This is the rest\nThat",
+    "is_reasoning_end": True,
+}
+SHORTEST_REASONING_NO_STREAMING_WITH_THINK = {
+    "output": "[/THINK]This is the rest",
+    "reasoning_content": "",
+    "content": "This is the rest",
+    "is_reasoning_end": True,
+}
+SHORTEST_REASONING_WITH_THINK = {
+    "output": "[/THINK]This is the rest",
+    "reasoning_content": None,
+    "content": "This is the rest",
+    "is_reasoning_end": True,
+}
+THINK_NO_END = {
+    "output": "[THINK]This is a reasoning section",
+    "reasoning_content": "This is a reasoning section",
+    "content": None,
+    "is_reasoning_end": False,
+}
+EMPTY = {
+    "output": "",
+    "reasoning_content": "",
+    "content": None,
+    "is_reasoning_end": False,
+}
+EMPTY_STREAMING = {
+    "output": "",
+    "reasoning_content": None,
+    "content": None,
+    "is_reasoning_end": False,
+}
+NEW_LINE = {
+    "output": "\n[THINK]This is a reasoning section[/THINK]\nThis is the rest",
+    "reasoning_content": "This is a reasoning section",
+    "content": "\nThis is the rest",
+    "is_reasoning_end": True,
+}
+# Streaming cannot handle new lines at the beginning of the output
+# because we need to support [THINK]...[/THINK] and [/THINK]...
+# We cannot know if the text before [THINK] is reasoning content
+# or not.
+NEW_LINE_STREAMING = {
+    "output": "\n[THINK]This is a reasoning section[/THINK]\nThis is the rest",
+    "reasoning_content": "\nThis is a reasoning section",
+    "content": "\nThis is the rest",
+    "is_reasoning_end": True,
+}
+
+TEST_CASES = [
+    pytest.param(
+        False,
+        SIMPLE_REASONING,
+        id="simple_reasoning",
+    ),
+    pytest.param(
+        True,
+        SIMPLE_REASONING,
+        id="simple_reasoning_streaming",
+    ),
+    pytest.param(
+        False,
+        COMPLETE_REASONING,
+        id="complete_reasoning",
+    ),
+    pytest.param(
+        True,
+        COMPLETE_REASONING,
+        id="complete_reasoning_streaming",
+    ),
+    pytest.param(
+        False,
+        NO_CONTENT,
+        id="no_content_token",
+    ),
+    pytest.param(
+        True,
+        NO_REASONING_STREAMING,
+        id="no_reasoning_token_streaming",
+    ),
+    pytest.param(
+        False,
+        MULTIPLE_LINES,
+        id="multiple_lines",
+    ),
+    pytest.param(
+        True,
+        MULTIPLE_LINES,
+        id="multiple_lines_streaming",
+    ),
+    pytest.param(
+        True,
+        SHORTEST_REASONING,
+        id="shortest",
+    ),
+    pytest.param(
+        False,
+        SHORTEST_REASONING_NO_STREAMING,
+        id="shortest_streaming",
+    ),
+    pytest.param(
+        False,
+        REASONING_WITH_THINK,
+        id="reasoning_with_think",
+    ),
+    pytest.param(
+        True,
+        REASONING_WITH_THINK,
+        id="reasoning_with_think_streaming",
+    ),
+    pytest.param(
+        False,
+        COMPLETE_REASONING_WITH_THINK,
+        id="complete_reasoning_with_think",
+    ),
+    pytest.param(
+        True,
+        COMPLETE_REASONING_WITH_THINK,
+        id="complete_reasoning_with_think_streaming",
+    ),
+    pytest.param(
+        False,
+        MULTIPLE_LINES_WITH_THINK,
+        id="multiple_lines_with_think",
+    ),
+    pytest.param(
+        True,
+        MULTIPLE_LINES_WITH_THINK,
+        id="multiple_lines_with_think_streaming",
+    ),
+    pytest.param(
+        False,
+        SHORTEST_REASONING_NO_STREAMING_WITH_THINK,
+        id="shortest_with_think",
+    ),
+    pytest.param(
+        True,
+        SHORTEST_REASONING_WITH_THINK,
+        id="shortest_with_think_streaming",
+    ),
+    pytest.param(
+        False,
+        THINK_NO_END,
+        id="think_no_end",
+    ),
+    pytest.param(
+        True,
+        THINK_NO_END,
+        id="think_no_end_streaming",
+    ),
+    pytest.param(
+        False,
+        EMPTY,
+        id="empty",
+    ),
+    pytest.param(
+        True,
+        EMPTY_STREAMING,
+        id="empty_streaming",
+    ),
+    pytest.param(
+        False,
+        NEW_LINE,
+        id="new_line",
+    ),
+    pytest.param(
+        True,
+        NEW_LINE_STREAMING,
+        id="new_line_streaming",
+    ),
+]
+
+
+@pytest.mark.parametrize("streaming, param_dict", TEST_CASES)
+def test_mistral_reasoning(
+    streaming: bool,
+    param_dict: dict,
+    mistral_tokenizer: MistralTokenizer,
+):
+    output = param_dict["output"]
+
+    index_think = output.find("[THINK]")
+    len_think = len("[THINK]")
+    index_end_think = output.find("[/THINK]")
+    len_end_think = len("[/THINK]")
+
+    # encode everything to tokens ids
+    output_tokens = []
+    if index_think != -1:
+        output_before_think = output[:index_think]
+        output_tokens += mistral_tokenizer.tokenizer.encode(
+            output_before_think, False, False)
+        output_tokens += [mistral_tokenizer.instruct.BEGIN_THINK]
+
+        if index_end_think != -1:
+            output_middle = output[index_think + len_think:index_end_think]
+            output_after_think = output[index_end_think + len_end_think:]
+            output_tokens += mistral_tokenizer.tokenizer.encode(
+                output_middle, False, False)
+            output_tokens += [mistral_tokenizer.instruct.END_THINK]
+            output_tokens += mistral_tokenizer.tokenizer.encode(
+                output_after_think, False, False)
+        else:
+            output_middle = output[index_think + len_think:]
+            output_tokens += mistral_tokenizer.tokenizer.encode(
+                output_middle, False, False)
+    elif index_end_think != -1:
+        output_before_think = output[:index_end_think]
+        output_after_think = output[index_end_think + len_end_think:]
+        output_tokens += mistral_tokenizer.tokenizer.encode(
+            output_before_think, False, False)
+        output_tokens += [mistral_tokenizer.instruct.END_THINK]
+        output_tokens += mistral_tokenizer.tokenizer.encode(
+            output_after_think, False, False)
+    else:
+        output_tokens += mistral_tokenizer.tokenizer.encode(
+            output, False, False)
+
+    parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(
+        parser_name)(mistral_tokenizer)
+
+    reasoning, content = run_reasoning_extraction_mistral(parser,
+                                                          output_tokens,
+                                                          streaming=streaming)
+
+    assert reasoning == param_dict["reasoning_content"]
+    assert content == param_dict["content"]
+
+    # Test is_reasoning_end
+    is_reasoning_end = parser.is_reasoning_end(output_tokens)
+    assert is_reasoning_end == param_dict["is_reasoning_end"]
+
+    # Test extract_content
+    if param_dict["content"] is not None:
+        content = parser.extract_content_ids(output_tokens)
+        assert content == mistral_tokenizer.tokenizer.encode(
+            param_dict["content"], bos=False, eos=False)
+    else:
+        content = parser.extract_content_ids(output_tokens)
+        assert content == []
@@ -6,6 +6,7 @@ from typing import Optional, Union
 from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
                                              DeltaMessage)
 from vllm.reasoning import ReasoningParser
+from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer


 class StreamingReasoningReconstructor:
@@ -54,6 +55,32 @@ def run_reasoning_extraction(
        return reasoning, content


+def run_reasoning_extraction_mistral(
+    reasoning_parser: ReasoningParser,
+    model_output: list[int],
+    request: Union[ChatCompletionRequest, None] = None,
+    streaming: bool = False,
+) -> tuple[Optional[str], Optional[str]]:
+    assert isinstance(reasoning_parser.model_tokenizer,
+                      MistralTokenizer), type(reasoning_parser.model_tokenizer)
+    if streaming:
+        reconstructor = run_reasoning_extraction_streaming_mistral(
+            reasoning_parser,
+            model_output,
+            request,
+        )
+        return (
+            reconstructor.reasoning_content,
+            reconstructor.other_content or None,
+        )
+    else:
+        str_output = reasoning_parser.model_tokenizer.convert_ids_to_tokens(
+            model_output)
+        reasoning, content = run_reasoning_extraction_nonstreaming(
+            reasoning_parser, str_output, request)
+        return reasoning, content
+
+
 def run_reasoning_extraction_nonstreaming(
    reasoning_parser: ReasoningParser,
    model_output: list[str],
@@ -94,3 +121,35 @@ def run_reasoning_extraction_streaming(
        previous_text = current_text
        previous_tokens = current_tokens
    return reconstructor
+
+
+def run_reasoning_extraction_streaming_mistral(
+    reasoning_parser: ReasoningParser,
+    model_deltas: list[int],
+    request: Union[ChatCompletionRequest, None] = None,
+) -> StreamingReasoningReconstructor:
+    assert isinstance(reasoning_parser.model_tokenizer,
+                      MistralTokenizer), type(reasoning_parser.model_tokenizer)
+    request = request or ChatCompletionRequest(messages=[], model="test-model")
+    reconstructor = StreamingReasoningReconstructor()
+    previous_text = ""
+    previous_tokens: list[int] = []
+    for model_delta in model_deltas:
+        token_delta = [model_delta]
+        delta = reasoning_parser.model_tokenizer.convert_ids_to_tokens(
+            [model_delta])[0]
+        current_text = previous_text + delta
+        current_tokens = previous_tokens + token_delta
+        delta_message = reasoning_parser.extract_reasoning_content_streaming(
+            previous_text,
+            current_text,
+            delta,
+            previous_tokens,
+            current_tokens,
+            token_delta,
+        )
+        if delta_message is not None:
+            reconstructor.append_delta(delta_message)
+        previous_text = current_text
+        previous_tokens = current_tokens
+    return reconstructor
@@ -2,9 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 from vllm import SamplingParams
-from vllm.config import LoadConfig, LoadFormat
+from vllm.config import LoadConfig
 from vllm.model_executor.model_loader import get_model_loader

+load_format = "runai_streamer"
 test_model = "openai-community/gpt2"

 prompts = [
@@ -18,7 +19,7 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)


 def get_runai_model_loader():
-    load_config = LoadConfig(load_format=LoadFormat.RUNAI_STREAMER)
+    load_config = LoadConfig(load_format=load_format)
    return get_model_loader(load_config)


@@ -28,6 +29,6 @@ def test_get_model_loader_with_runai_flag():


 def test_runai_model_loader_download_files(vllm_runner):
-    with vllm_runner(test_model, load_format=LoadFormat.RUNAI_STREAMER) as llm:
+    with vllm_runner(test_model, load_format=load_format) as llm:
        deserialized_outputs = llm.generate(prompts, sampling_params)
        assert deserialized_outputs
@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.model_executor.models.phi3v import Phi3VImagePixelInputs
+
+
+def test_tensor_schema_valid_tensor():
+    Phi3VImagePixelInputs(
+        data=torch.randn(16, 64, 3, 32, 32),
+        image_sizes=torch.randint(0, 256, (16, 2)),
+    )
+
+
+def test_tensor_schema_optional_fields():
+    Phi3VImagePixelInputs(
+        data=torch.randn(16, 64, 3, 32, 32),
+        image_sizes=None,
+    )
+
+    Phi3VImagePixelInputs(data=torch.randn(16, 64, 3, 32, 32), )
+
+
+def test_tensor_schema_constant_dim_failure():
+    with pytest.raises(ValueError, match="dim\\[2\\] expected 3, got 4"):
+        Phi3VImagePixelInputs(
+            data=torch.randn(16, 64, 4, 32, 32),  # dim[2] = 4
+            image_sizes=torch.randint(0, 256, (16, 2)),
+        )
+
+
+def test_tensor_schema_symbolic_dim_mismatch():
+    with pytest.raises(ValueError, match="expected 'bn'=12, got 16"):
+        Phi3VImagePixelInputs(
+            data=torch.randn(12, 64, 3, 32, 32),
+            image_sizes=torch.randint(0, 256, (16, 2)),
+        )
+
+
+def test_tensor_schema_list_tensor_valid():
+    Phi3VImagePixelInputs(
+        data=[torch.randn(64, 3, 32, 32) for _ in range(16)],
+        image_sizes=torch.randint(0, 256, (16, 2)),
+    )
+
+
+def test_tensor_schema_variable_patch_counts_valid():
+    # Each image has a different number of patches (p)
+    # Each tensor has shape (p, 3, 32, 32)
+    data = [
+        torch.randn(16, 3, 32, 32),  # p = 16
+        torch.randn(32, 3, 32, 32),  # p = 32
+        torch.randn(64, 3, 32, 32),  # p = 64
+    ]
+    image_sizes = torch.randint(0, 256, (3, 2))  # bn = 3
+    Phi3VImagePixelInputs(
+        data=data,
+        image_sizes=image_sizes,
+    )
+
+
+def test_tensor_schema_tuple_tensor_valid():
+    Phi3VImagePixelInputs(
+        data=tuple(torch.randn(64, 3, 32, 32) for _ in range(16)),
+        image_sizes=torch.randint(0, 256, (16, 2)),
+    )
+
+
+def test_tensor_schema_inconsistent_shapes_in_list():
+    with pytest.raises(ValueError, match="contains inconsistent shapes"):
+        Phi3VImagePixelInputs(
+            data=[torch.randn(64, 3, 32, 32),
+                  torch.randn(64, 3, 16, 16)] +
+            [torch.randn(64, 3, 32, 32) for _ in range(14)],
+            image_sizes=torch.randint(0, 256, (16, 2)),
+        )
+
+
+def test_tensor_schema_empty_list():
+    with pytest.raises(ValueError, match="is an empty list"):
+        Phi3VImagePixelInputs(
+            data=[],
+            image_sizes=torch.randint(0, 256, (0, 2)),
+        )
+
+
+def test_tensor_schema_validation_disabled_skips_shape_check():
+    # This should NOT raise, because validation is turned off
+    # This would normally fail (dim[2] should be 3, not 4)
+    Phi3VImagePixelInputs(
+        data=torch.randn(16, 64, 4, 32, 32),
+        image_sizes=torch.randint(0, 256, (16, 2)),
+        validate=False,
+    )
+
+
+def test_tensor_schema_with_valid_resolve_binding_dims():
+    data = torch.randn(16, 64, 3, 336, 336)  # h=336, w=336
+    image_sizes = torch.randint(0, 256, (16, 2))
+
+    Phi3VImagePixelInputs(
+        data=data,
+        image_sizes=image_sizes,
+        resolve_bindings={
+            "h": 336,
+            "w": 336
+        },
+    )
+
+
+def test_tensor_schema_with_invalid_resolve_binding_dims():
+    data = torch.randn(16, 64, 3, 36, 36)  # h=36, w=36
+    image_sizes = torch.randint(0, 256, (16, 2))
+
+    # Should raise because 'h' and 'w' don't match resolve bindings
+    with pytest.raises(ValueError, match="dim\\[3\\] expected 336, got 36"):
+        Phi3VImagePixelInputs(
+            data=data,
+            image_sizes=image_sizes,
+            resolve_bindings={
+                "h": 336,
+                "w": 336
+            },
+        )
@@ -11,7 +11,8 @@ from tests.v1.attention.utils import (BatchSpec, _Backend,
                                      create_vllm_config,
                                      get_attention_backend)
 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
+                                              set_kv_cache_layout)
 from vllm.v1.kv_cache_interface import FullAttentionSpec

 BACKENDS_TO_TEST = [
@@ -212,7 +213,7 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,

        from vllm.v1.attention.backends.flashinfer import PerLayerParameters

-        def mock_get_per_layer_parameters(vllm_config):
+        def mock_get_per_layer_parameters(vllm_config, impl_cls):
            # Return mock parameters for a single layer
            head_size = vllm_config.model_config.get_head_size()
            return {
@@ -297,7 +298,8 @@ def test_backend_correctness(batch_spec_name: str, model: str):
    5. Comparing the vLLM backend's output to the ground-truth SDPA output.
    """
    batch_spec = BATCH_SPECS[batch_spec_name]
-    vllm_config = create_vllm_config(model_name=model)
+    vllm_config = create_vllm_config(model_name=model,
+                                     max_model_len=max(batch_spec.seq_lens))
    device = torch.device("cuda:0")

    kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
@@ -419,6 +421,11 @@ def test_backend_correctness(batch_spec_name: str, model: str):
        if backend_name == _Backend.FLASHINFER_VLLM_V1:
            kv_cache_for_backend = kv_cache.transpose(0, 1)

+            # For FlashInfer default to HND layout and
+            kv_cache_for_backend = kv_cache_for_backend.transpose(
+                2, 3).contiguous().transpose(2, 3)
+            set_kv_cache_layout("HND")
+
        backend_output = run_attention_backend(backend_name, kv_cache_spec,
                                               vllm_config, device,
                                               common_attn_metadata,
@@ -66,7 +66,7 @@ def create_common_attn_metadata(
    num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)

    # Create block table (random for testing)
-    max_blocks = max(batch_spec.seq_lens) // block_size + 1
+    max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size
    block_table_tensor = torch.randint(0,
                                       max_block_idx,
                                       (batch_spec.batch_size, max_blocks),
@@ -565,8 +565,8 @@ def test_engine_core_proc_instantiation_cuda_empty(

        from vllm.v1.engine.utils import EngineZmqAddresses

-        def mock_startup_handshake(self, handshake_socket, on_head_node,
-                                   parallel_config):
+        def mock_startup_handshake(self, handshake_socket, local_client,
+                                   headless, parallel_config):
            return EngineZmqAddresses(inputs=["tcp://127.0.0.1:5555"],
                                      outputs=["tcp://127.0.0.1:5556"],
                                      coordinator_input=None,
@@ -0,0 +1,162 @@
+#!/bin/bash
+set -xe
+
+# Hosts / ports
+PREFILL_HOST=${PREFILL_HOST:-"localhost"}
+PREFILL_PORT=${PREFILL_PORT:-8100}
+PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577}
+DECODE_HOST=${DECODE_HOST:-"localhost"}
+DECODE_PORT=${DECODE_PORT:-8200}
+PROXY_HOST=${PROXY_HOST:-"localhost"}
+PROXY_PORT=${PROXY_PORT:-8192}
+BASELINE_HOST=${BASELINE_HOST:-"localhost"}
+BASELINE_PORT=${BASELINE_PORT:-9290}
+
+
+# Model to run.
+MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B-Instruct"}
+MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024}
+BLOCK_SIZE=${BLOCK_SIZE:-32}
+
+
+# execution env
+GIT_ROOT=$(git rev-parse --show-toplevel)
+EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration"
+CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"}
+CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"}
+
+OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.tpu_accuracy_test_outputs.txt"}
+
+# Trap the SIGINT signal (triggered by Ctrl+C)
+trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
+
+
+# Waits for vLLM server to start.
+wait_for_server() {
+  local host=$1
+  local port=$2
+  timeout 1200 bash -c "
+    until curl -s ${host}:${port}/v1/completions > /dev/null; do
+      sleep 1
+    done" && return 0 || return 1
+}
+
+# Cleanup function
+cleanup() {
+    echo "Caught Ctrl+C, cleaning up..."
+    # Cleanup commands
+    pgrep python | xargs kill -9 || true
+    # pkill -f python || true
+    echo "Cleanup complete. Exiting."
+}
+
+launch_baseline() {
+  BASELINE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
+  VLLM_LOGGING_LEVEL=DEBUG \
+  VLLM_USE_V1=1 \
+  PJRT_DEVICE=TPU \
+  VLLM_WORKER_MULTIPROC_METHOD=spawn \
+  VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
+      --host ${BASELINE_HOST} \
+      --port ${BASELINE_PORT} \
+      --max-model-len ${MAX_MODEL_LEN}\
+      --seed 42 \
+      --block-size ${BLOCK_SIZE} \
+      --gpu-memory-utilization 0.5 \
+      --disable-log-requests \
+      --enforce-eager"
+  echo ${BASELINE_BASE_CMD}
+  ssh -tt ${BASELINE_HOST} "${BASELINE_BASE_CMD}" &
+}
+
+launch_pd() {
+  PREFILL_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
+  UCX_TLS=tcp \
+  VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
+  VLLM_LOGGING_LEVEL=DEBUG \
+  VLLM_USE_V1=1 \
+  VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \
+  VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \
+  PJRT_DEVICE=TPU \
+  VLLM_WORKER_MULTIPROC_METHOD=spawn \
+  VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
+      --host ${PREFILL_HOST} \
+      --port ${PREFILL_PORT} \
+      --max-model-len ${MAX_MODEL_LEN}\
+      --seed 42 \
+      --block-size ${BLOCK_SIZE} \
+      --enforce-eager \
+      --gpu-memory-utilization 0.5 \
+      --disable-log-requests \
+      --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'"
+
+
+  DECODE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
+  UCX_TLS=tcp \
+  VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
+  VLLM_LOGGING_LEVEL=DEBUG \
+  VLLM_USE_V1=1 \
+  PJRT_DEVICE=TPU \
+  VLLM_WORKER_MULTIPROC_METHOD=spawn \
+  VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
+      --host ${DECODE_HOST} \
+      --port ${DECODE_PORT} \
+      --max-model-len ${MAX_MODEL_LEN}\
+      --seed 42 \
+      --block-size ${BLOCK_SIZE} \
+      --enforce-eager \
+      --gpu-memory-utilization 0.5 \
+      --disable-log-requests \
+      --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'"
+
+  echo ${PREFILL_BASE_CMD}
+  echo ${DECODE_BASE_CMD}
+  sleep 2
+
+  # execute on hosts
+  ssh -tt ${PREFILL_HOST} "${PREFILL_BASE_CMD}" &
+  ssh -tt ${DECODE_HOST} "${DECODE_BASE_CMD}" &
+  sleep 1
+  wait_for_server ${PREFILL_HOST} ${PREFILL_PORT}
+  sleep 1
+  wait_for_server ${DECODE_HOST} ${DECODE_PORT}
+  sleep 1
+}
+
+launch_pd_proxy(){
+  PROXY_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
+  python3 ${EXP_ROOT}/toy_proxy_server.py \
+  --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \
+  --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \
+  --host=${PROXY_HOST} --port ${PROXY_PORT}"
+  echo ${PROXY_BASE_CMD}
+  ssh -tt ${PROXY_HOST} "${PROXY_BASE_CMD}" &
+}
+
+run_tests(){
+  local service_url=$1
+  local mode=$2
+  python3 ${EXP_ROOT}/test_disagg_accuracy.py --service_url=${service_url} --model_name=${MODEL_NAME} --mode=${mode} --file_name=${OUTPUT_FILE}
+}
+
+
+# run non-disagg. baseline & save outputs
+launch_baseline
+sleep 2
+wait_for_server ${BASELINE_HOST} ${BASELINE_PORT}
+run_tests "http://${BASELINE_HOST}:${BASELINE_PORT}" "baseline"
+cleanup
+sleep 10
+
+
+# run disagg. & do exact-match with the outputs from baseline
+launch_pd
+launch_pd_proxy
+sleep 10
+run_tests "http://${PROXY_HOST}:${PROXY_PORT}" "disagg"
+echo "-----P/D success----"
+
+rm ${OUTPUT_FILE}
+cleanup
+
+exit 0
@@ -0,0 +1,128 @@
+#!/bin/bash
+set -xe
+
+# Hosts / ports
+PREFILL_HOST=${PREFILL_HOST:-"localhost"}
+PREFILL_PORT=${PREFILL_PORT:-8100}
+PREFILL_NIXL_SIDE_PORT=${PREFILL_NIXL_SIDE_PORT:-5577}
+DECODE_HOST=${DECODE_HOST:-"localhost"}
+DECODE_PORT=${DECODE_PORT:-8200}
+PROXY_HOST=${PROXY_HOST:-"localhost"}
+PROXY_PORT=${PROXY_PORT:-8192}
+BASELINE_HOST=${BASELINE_HOST:-"localhost"}
+BASELINE_PORT=${BASELINE_PORT:-9290}
+
+
+# Model to run.
+MODEL_NAME=${MODEL_NAME:-"meta-llama/Llama-3.2-3B-Instruct"}
+MAX_MODEL_LEN=${MAX_MODEL_LEN:-1024}
+BLOCK_SIZE=${BLOCK_SIZE:-32}
+
+
+# execution env
+GIT_ROOT=$(git rev-parse --show-toplevel)
+EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration"
+CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"}
+CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"}
+
+OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.tpu_accuracy_test_outputs.txt"}
+
+# Trap the SIGINT signal (triggered by Ctrl+C)
+trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
+
+# Waits for vLLM server to start.
+wait_for_server() {
+  local host=$1
+  local port=$2
+  timeout 1200 bash -c "
+    until curl -s ${host}:${port}/v1/completions > /dev/null; do
+      sleep 1
+    done" && return 0 || return 1
+}
+
+# Cleanup function
+cleanup() {
+    echo "Caught Ctrl+C, cleaning up..."
+    # Cleanup commands
+    pgrep python | xargs kill -9 || true
+    # pkill -f python || true
+    echo "Cleanup complete. Exiting."
+}
+
+
+launch_pd() {
+  PREFILL_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
+  UCX_TLS=tcp \
+  VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
+  VLLM_LOGGING_LEVEL=DEBUG \
+  VLLM_USE_V1=1 \
+  VLLM_NIXL_SIDE_CHANNEL_HOST=${PREFILL_HOST} \
+  VLLM_NIXL_SIDE_CHANNEL_PORT=${PREFILL_NIXL_SIDE_PORT} \
+  PJRT_DEVICE=TPU \
+  VLLM_WORKER_MULTIPROC_METHOD=spawn \
+  VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
+      --host ${PREFILL_HOST} \
+      --port ${PREFILL_PORT} \
+      --max-model-len ${MAX_MODEL_LEN}\
+      --seed 42 \
+      --block-size ${BLOCK_SIZE} \
+      --enforce-eager \
+      --gpu-memory-utilization 0.5 \
+      --disable-log-requests \
+      --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'"
+
+
+  DECODE_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
+  UCX_TLS=tcp \
+  VLLM_MULTIPROC_EXECUTE_MODEL_TIMEOUT_S=200 \
+  VLLM_LOGGING_LEVEL=DEBUG \
+  VLLM_USE_V1=1 \
+  PJRT_DEVICE=TPU \
+  VLLM_WORKER_MULTIPROC_METHOD=spawn \
+  VLLM_ENABLE_V1_MULTIPROCESSING=0 vllm serve $MODEL_NAME \
+      --host ${DECODE_HOST} \
+      --port ${DECODE_PORT} \
+      --max-model-len ${MAX_MODEL_LEN}\
+      --seed 42 \
+      --block-size ${BLOCK_SIZE} \
+      --enforce-eager \
+      --gpu-memory-utilization 0.5 \
+      --disable-log-requests \
+      --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"cpu\"}'"
+
+  echo ${PREFILL_BASE_CMD}
+  echo ${DECODE_BASE_CMD}
+  sleep 2
+
+  # execute on hosts
+  ssh -tt ${PREFILL_HOST} "${PREFILL_BASE_CMD}" &
+  ssh -tt ${DECODE_HOST} "${DECODE_BASE_CMD}" &
+  sleep 1
+  wait_for_server ${PREFILL_HOST} ${PREFILL_PORT}
+  sleep 1
+  wait_for_server ${DECODE_HOST} ${DECODE_PORT}
+  sleep 1
+}
+
+launch_pd_proxy(){
+  PROXY_BASE_CMD="source ${CONDA_PATH}/bin/activate ${CONDA_ENV_NAME};
+  python3 ${EXP_ROOT}/toy_proxy_server.py \
+  --prefiller-host ${PREFILL_HOST} --prefiller-port ${PREFILL_PORT} \
+  --decoder-host ${DECODE_HOST} --decoder-port ${DECODE_PORT} \
+  --host=${PROXY_HOST} --port ${PROXY_PORT}"
+  echo ${PROXY_BASE_CMD}
+  ssh -tt ${PROXY_HOST} "${PROXY_BASE_CMD}" &
+}
+
+
+# run disagg. & do exact-match with the outputs from baseline
+launch_pd
+launch_pd_proxy
+sleep 10
+
+PREFILL_HOST=${PREFILL_HOST} \
+PREFILL_PORT=${PREFILL_PORT} \
+DECODE_HOST=${DECODE_HOST} \
+DECODE_PORT=${DECODE_PORT} \
+PROXY_HOST=${PROXY_HOST} \
+PROXY_PORT=${PROXY_PORT} python -m pytest -s -v ${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_edge_cases.py
@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import argparse
+import json
+import os
+import time
+
+import openai
+import requests
+
+MAX_OUTPUT_LEN = 30
+
+SAMPLE_PROMPTS = (
+    "Red Hat is the best company in the world to work for because it works on "
+    "open source software, which means that all the contributions are "
+    "delivered to the community. As a result, when working on projects like "
+    "vLLM we are able to meet many amazing people from various organizations "
+    "like AMD, Google, NVIDIA, ",
+    "We hold these truths to be self-evident, that all men are created equal, "
+    "that they are endowed by their Creator with certain unalienable Rights, "
+    "that among these are Life, Liberty and the pursuit of Happiness.--That "
+    "to secure these rights, Governments are instituted among Men, deriving "
+    "their just powers from the consent of the governed, ",
+)
+
+
+def check_vllm_server(url: str, timeout=5, retries=3) -> bool:
+    """
+    Checks if the vLLM server is ready by sending a GET request to the
+    /health endpoint.
+
+    Args:
+        url (str): The base URL of the vLLM server.
+        timeout (int): Timeout in seconds for the request.
+        retries (int): Number of retries if the server is not ready.
+
+    Returns:
+        bool: True if the server is ready, False otherwise.
+    """
+    for attempt in range(retries):
+        try:
+            response = requests.get(url, timeout=timeout)
+            if response.status_code == 200:
+                return True
+            else:
+                print(f"Attempt {attempt + 1}: Server returned status code "
+                      "{response.status_code}")
+        except requests.exceptions.RequestException as e:
+            print(f"Attempt {attempt + 1}: Error connecting to server: {e}")
+        time.sleep(1)  # Wait before retrying
+    return False
+
+
+def run_simple_prompt(base_url: str, model_name: str,
+                      input_prompt: str) -> str:
+    client = openai.OpenAI(api_key="EMPTY", base_url=base_url)
+    completion = client.completions.create(model=model_name,
+                                           prompt=input_prompt,
+                                           max_tokens=MAX_OUTPUT_LEN,
+                                           temperature=0.0,
+                                           seed=42)
+
+    # print("-" * 50)
+    # print(f"Completion results for {model_name}:")
+    # print(completion)
+    # print("-" * 50)
+    return completion.choices[0].text
+
+
+def main():
+    """
+    This script demonstrates how to accept two optional string arguments
+    ("service_url" and "file_name") from the command line, each with a
+    default value of an empty string, using the argparse module.
+    """
+    parser = argparse.ArgumentParser(description="vLLM client script")
+
+    parser.add_argument(
+        "--service_url",  # Name of the first argument
+        type=str,
+        required=True,
+        help="The vLLM service URL.")
+
+    parser.add_argument(
+        "--model_name",  # Name of the first argument
+        type=str,
+        required=True,
+        help="model_name",
+    )
+
+    parser.add_argument(
+        "--mode",  # Name of the second argument
+        type=str,
+        default="baseline",
+        help="mode: baseline==non-disagg, or disagg",
+    )
+
+    parser.add_argument(
+        "--file_name",  # Name of the second argument
+        type=str,
+        default=".vllm_output.txt",
+        help="the file that saves the output tokens ",
+    )
+
+    args = parser.parse_args()
+
+    for arg in vars(args):
+        print(f"{arg}: {getattr(args, arg)}")
+
+    if args.mode == "baseline":
+        # non-disagg
+        health_check_url = f"{args.service_url}/health"
+    else:
+        # disagg proxy
+        health_check_url = f"{args.service_url}/healthcheck"
+        if not os.path.exists(args.file_name):
+            raise ValueError(
+                f"In disagg mode, the output file {args.file_name} from "
+                "non-disagg. baseline does not exist.")
+
+    service_url = f"{args.service_url}/v1"
+
+    if not check_vllm_server(health_check_url):
+        raise RuntimeError(
+            f"vllm server: {args.service_url} is not ready yet!")
+
+    output_strs = dict()
+    for prompt in SAMPLE_PROMPTS:
+        output_str = run_simple_prompt(base_url=service_url,
+                                       model_name=args.model_name,
+                                       input_prompt=prompt)
+        print(f"Prompt: {prompt}, output: {output_str}")
+        output_strs[prompt] = output_str
+
+    if args.mode == "baseline":
+        # baseline: save outputs
+        try:
+            with open(args.file_name, 'w') as json_file:
+                json.dump(output_strs, json_file, indent=4)
+        except OSError as e:
+            print(f"Error writing to file: {e}")
+            raise
+    else:
+        # disagg. verify outputs
+        baseline_outputs = None
+        try:
+            with open(args.file_name) as json_file:
+                baseline_outputs = json.load(json_file)
+        except OSError as e:
+            print(f"Error writing to file: {e}")
+            raise
+        assert isinstance(baseline_outputs, dict)
+        assert len(baseline_outputs) == len(output_strs)
+        for prompt, output in baseline_outputs.items():
+            assert prompt in output_strs, f"{prompt} not included"
+            assert output == output_strs[prompt], (
+                f"baseline_output: {output} != PD output: {output_strs[prompt]}"
+            )
+
+
+if __name__ == "__main__":
+    main()
@@ -4,8 +4,11 @@ import os

 import openai

+PREFILL_HOST = os.getenv("PREFILL_HOST", "localhost")
 PREFILL_PORT = os.getenv("PREFILL_PORT", None)
+DECODE_HOST = os.getenv("DECODE_HOST", "localhost")
 DECODE_PORT = os.getenv("DECODE_PORT", None)
+PROXY_HOST = os.getenv("PROXY_HOST", "localhost")
 PROXY_PORT = os.getenv("PROXY_PORT", None)

 if PREFILL_PORT is None or DECODE_PORT is None or PROXY_PORT is None:
@@ -21,15 +24,15 @@ def test_edge_cases():
    # Set the OpenAI API key and base URL
    decode_client = openai.OpenAI(
        api_key="MY_KEY",
-        base_url=f"http://localhost:{DECODE_PORT}/v1",
+        base_url=f"http://{DECODE_HOST}:{DECODE_PORT}/v1",
    )
    prefill_client = openai.OpenAI(
        api_key="MY_KEY",
-        base_url=f"http://localhost:{PREFILL_PORT}/v1",
+        base_url=f"http://{PREFILL_HOST}:{PREFILL_PORT}/v1",
    )
    proxy_client = openai.OpenAI(
        api_key="MY_KEY",
-        base_url=f"http://localhost:{PROXY_PORT}/v1",
+        base_url=f"http://{PROXY_HOST}:{PROXY_PORT}/v1",
    )

    # Get the list of models
@@ -3,6 +3,7 @@

 import argparse
 import itertools
+import logging
 import os
 import uuid
 from contextlib import asynccontextmanager
@@ -11,9 +12,8 @@ import httpx
 from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse

-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)


@asynccontextmanager
@@ -1,10 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import contextlib
+import inspect
 import os
 import tempfile
 import textwrap
 import time
+import uuid
+from collections import defaultdict
+from typing import Optional
 from unittest.mock import patch

 import pytest
@@ -16,30 +21,118 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
    KVConnectorRole, NixlAgentMetadata, NixlConnector, NixlConnectorMetadata,
    NixlConnectorWorker)
 from vllm.forward_context import ForwardContext
-from vllm.mocks.mock_nixl_connector import FakeNixlWrapper
 from vllm.sampling_params import SamplingParams

 from .utils import create_request, create_scheduler, create_vllm_config


-def _make_stub_pkg() -> str:
-    """Return a directory that makes
-       `from nixl._api import nixl_agent` resolve to our FakeNixlWrapper."""
-    td = tempfile.mkdtemp()
-    pkg_root = os.path.join(td, "nixl", "_api")
-    os.makedirs(pkg_root, exist_ok=True)
+class FakeNixlWrapper:
+    """Mock implementation of NixlWrapper for testing.

-    stub = textwrap.dedent("""\
-        # Forward the real FakeNixlWrapper that the driver already defined.
-        print("In fake package")
-        from vllm.mocks.mock_nixl_connector import FakeNixlWrapper as nixl_agent
-    """)
-    with open(os.path.join(pkg_root, "__init__.py"), "w") as f:
-        f.write(stub)
+    We don't inherit from nixl._api.nixl_agent because nixl may not be
+    installed.
+    
+    Note: The complete source of this class is also used in the
+    `_make_fake_nixl_pkg` function to create a fake nixl package
+    for Ray workers.
+    """

-    # touch parent package
-    open(os.path.join(td, "nixl", "__init__.py"), "w").close()
-    return td
+    AGENT_METADATA = b"fake_agent_metadata"
+    REMOTE_AGENT_NAME = "remote_agent"
+
+    def __init__(self, agent_name: str, *args, **kwargs):
+        self._cycles_before_xfer_done = 0
+        self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(
+            lambda: 0)
+
+    def get_reg_descs(self, caches_data, memory_type: str) -> list:
+        return [str(uuid.uuid4()) for _ in caches_data]
+
+    def register_memory(self, descs) -> None:
+        pass
+
+    def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
+        return [str(uuid.uuid4()) for _ in blocks_data]
+
+    def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
+        return uuid.uuid4().int
+
+    def get_agent_metadata(self) -> bytes:
+        return self.AGENT_METADATA
+
+    def add_remote_agent(self, agent_metadata: bytes) -> str:
+        return self.REMOTE_AGENT_NAME
+
+    def get_new_notifs(self) -> dict[str, list[bytes]]:
+        # Used to collect done_sending, which we don't test yet.
+        return {}
+
+    def check_xfer_state(self, handle: int) -> str:
+        if self._check_xfer_state_cycles[
+                handle] >= self._cycles_before_xfer_done:
+            return "DONE"
+        self._check_xfer_state_cycles[handle] += 1
+        return "PROC"
+
+    def release_xfer_handle(self, handle: int) -> None:
+        pass
+
+    def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
+        pass
+
+    def make_prepped_xfer(self,
+                          xfer_type: str,
+                          local_xfer_side_handle: int,
+                          local_block_descs_ids: list[int],
+                          remote_xfer_side_handle: int,
+                          remote_block_descs_ids: list[int],
+                          notif_msg: Optional[bytes] = None) -> int:
+        return uuid.uuid4().int
+
+    def transfer(self, handle: int) -> str:
+        return "PROC"
+
+    ############################################################
+    # Follow are for changing the behavior during testing.
+    ############################################################
+
+    def set_cycles_before_xfer_done(self, cycles: int):
+        """Set the number of cycles before a transfer is considered done."""
+        self._cycles_before_xfer_done = cycles
+
+
+@contextlib.contextmanager
+def _make_fake_nixl_pkg():
+    """Context manager that creates a temporary package making
+       `from nixl._api import nixl_agent` resolve to our FakeNixlWrapper.
+       
+    Automatically cleans up the temporary directory when done.
+    """
+    with tempfile.TemporaryDirectory() as td:
+        pkg_root = os.path.join(td, "nixl", "_api")
+        os.makedirs(pkg_root, exist_ok=True)
+
+        # Get the source code of FakeNixlWrapper class and dedent it
+        fake_nixl_source = inspect.getsource(FakeNixlWrapper)
+        fake_nixl_source = textwrap.dedent(fake_nixl_source)
+
+        stub = f"""\
+# Copy of FakeNixlWrapper implementation for Ray workers
+import uuid
+from collections import defaultdict
+from typing import Optional
+
+{fake_nixl_source}
+
+# Export as nixl_agent
+nixl_agent = FakeNixlWrapper
+"""
+        with open(os.path.join(pkg_root, "__init__.py"), "w") as f:
+            f.write(stub)
+
+        # touch parent package
+        open(os.path.join(td, "nixl", "__init__.py"), "w").close()
+        yield td


 def test_basic_interface():
@@ -351,27 +444,37 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
        kv_connector="NixlConnector",
        kv_role="kv_both",
    )
+    llm_kwargs = {
+        "model": model_name,
+        "enforce_eager": True,
+        "gpu_memory_utilization": 0.5,
+        "kv_transfer_config": kv_transfer_config,
+        "distributed_executor_backend": distributed_executor_backend,
+    }
+
    timeout = 6
    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
    monkeypatch.setenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", str(timeout))

-    # Build runtime_env only if we’re using Ray
+    # Build runtime_env only if we're using Ray
    if distributed_executor_backend == "ray":
-        runtime_env = {
-            "working_dir": _make_stub_pkg(),  # ship stub package
-            "env_vars": {
-                "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout),
-            },
-        }
-        ray.init(runtime_env=runtime_env)
+        with _make_fake_nixl_pkg() as working_dir:
+            runtime_env = {
+                "working_dir": working_dir,  # ship fake nixl package
+                "env_vars": {
+                    "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout),
+                },
+            }
+            ray.init(runtime_env=runtime_env)

-    llm = LLM(
-        model=model_name,
-        enforce_eager=True,
-        gpu_memory_utilization=0.5,
-        kv_transfer_config=kv_transfer_config,
-        distributed_executor_backend=distributed_executor_backend,
-    )
+            _run_abort_timeout_test(llm_kwargs, timeout)
+    else:
+        _run_abort_timeout_test(llm_kwargs, timeout)
+
+
+def _run_abort_timeout_test(llm_kwargs: dict, timeout: int):
+    """Helper function to run the abort timeout test logic."""
+    llm = LLM(**llm_kwargs)
    remote_prefill_opts = {
        "do_remote_decode": True,
        "do_remote_prefill": False,
@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+import os
+import threading
+import time
+from contextlib import AsyncExitStack
+
+import openai  # use the official client for correctness check
+import pytest
+import pytest_asyncio
+
+from tests.utils import RemoteOpenAIServer
+from tests.v1.test_utils import check_request_balancing
+from vllm.platforms import Platform
+
+MODEL_NAME = "ibm-research/PowerMoE-3b"
+
+# Number of data parallel ranks for hybrid LB testing (4 total)
+DP_SIZE = int(os.getenv("DP_SIZE", "4"))
+# Default tensor parallel size to use
+TP_SIZE = int(os.getenv("TP_SIZE", "1"))
+
+# Number of nodes (2 nodes, each with 2 DP ranks)
+NUM_NODES = 2
+DP_SIZE_LOCAL = DP_SIZE // NUM_NODES  # 2 ranks per node
+
+
+class HybridLBServerManager:
+    """Manages hybrid data parallel vLLM server instances where each node 
+    runs a single logical API server that balances requests only to the 
+    DP engines running on that same node."""
+
+    def __init__(self,
+                 model_name: str,
+                 dp_size: int,
+                 api_server_count: int,
+                 base_server_args: list,
+                 dp_size_local: int = DP_SIZE_LOCAL,
+                 tp_size: int = TP_SIZE):
+        self.model_name = model_name
+        self.dp_size = dp_size
+        self.dp_size_local = dp_size_local
+        self.tp_size = tp_size
+        self.api_server_count = api_server_count
+        self.base_server_args = base_server_args
+        self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
+        self.server_threads: list[threading.Thread] = []
+        self.num_nodes = dp_size // dp_size_local
+
+    def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
+        """Start all server instances for hybrid LB mode."""
+        for node_id in range(self.num_nodes):
+            # Create server args for this specific node
+            server_args = self.base_server_args.copy()
+
+            # Calculate start rank for this node
+            start_rank = node_id * self.dp_size_local
+
+            # Add hybrid LB specific arguments
+            server_args.extend([
+                "--data-parallel-size",
+                str(self.dp_size),
+                "--data-parallel-size-local",
+                str(self.dp_size_local),
+                "--data-parallel-start-rank",
+                str(start_rank),
+                "--data-parallel-hybrid-lb",  # Enable hybrid LB mode
+                "--tensor-parallel-size",
+                str(self.tp_size),
+                "--port",
+                str(8000 + node_id),  # Different port for each node
+                "--api-server-count",
+                str(self.api_server_count),
+                "--data-parallel-address",
+                "127.0.0.1",
+                "--data-parallel-rpc-port",
+                "13345",
+            ])
+
+            # Use a thread to start each server to allow parallel initialization
+            def start_server(node: int, sargs: list[str]):
+                try:
+                    # Calculate GPU devices for this node
+                    gpus_per_node = self.dp_size_local * self.tp_size
+                    gpu_start = node * gpus_per_node
+                    gpu_end = gpu_start + gpus_per_node
+
+                    # Start the server
+                    server = RemoteOpenAIServer(
+                        self.model_name,
+                        sargs,
+                        auto_port=False,
+                        env_dict={
+                            "CUDA_VISIBLE_DEVICES":
+                            ",".join(
+                                str(Platform.device_id_to_physical_device_id(
+                                    i)) for i in range(gpu_start, gpu_end))
+                        })
+                    server.__enter__()
+                    print(f"Hybrid LB node {node} started successfully with "
+                          f"{self.dp_size_local} local DP ranks and "
+                          f"{self.api_server_count} API servers")
+                    self.servers.append((server, sargs))
+                except Exception as e:
+                    print(f"Failed to start hybrid LB node {node}: {e}")
+                    raise
+
+            thread = threading.Thread(target=start_server,
+                                      args=(node_id, server_args))
+            thread.start()
+
+            self.server_threads.append(thread)
+
+        # Wait for all servers to start
+        for thread in self.server_threads:
+            thread.join()
+
+        # Give servers additional time to fully initialize and coordinate
+        time.sleep(3)
+
+        if len(self.servers) != self.num_nodes:
+            raise Exception("Servers failed to start")
+
+        return self.servers
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Stop all server instances."""
+        while self.servers:
+            try:
+                self.servers.pop()[0].__exit__(exc_type, exc_val, exc_tb)
+            except Exception as e:
+                print(f"Error stopping server: {e}")
+
+
+@pytest.fixture(scope="module")
+def default_server_args():
+    return [
+        # use half precision for speed and memory savings in CI environment
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        "2048",
+        "--max-num-seqs",
+        "128",
+        "--enforce-eager",
+    ]
+
+
+@pytest.fixture(scope="module", params=[1, 4])
+def servers(request, default_server_args):
+    api_server_count = request.param
+    with HybridLBServerManager(MODEL_NAME, DP_SIZE, api_server_count,
+                               default_server_args, DP_SIZE_LOCAL,
+                               TP_SIZE) as server_list:
+        yield server_list
+
+
+@pytest_asyncio.fixture
+async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
+    # Create a client for each node (each node has its own API endpoint)
+    async with AsyncExitStack() as stack:
+        yield [
+            await stack.enter_async_context(server.get_async_client())
+            for server, _ in servers
+        ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_hybrid_lb_completion(clients: list[openai.AsyncOpenAI],
+                                    servers: list[tuple[RemoteOpenAIServer,
+                                                        list[str]]],
+                                    model_name: str) -> None:
+
+    async def make_request(client: openai.AsyncOpenAI):
+        completion = await client.completions.create(
+            model=model_name,
+            prompt="Hello, my name is",
+            max_tokens=10,
+            temperature=1.0)
+
+        assert completion.id is not None
+        assert completion.choices is not None and len(completion.choices) == 1
+
+        choice = completion.choices[0]
+        # The exact number of tokens can vary slightly with temperature=1.0,
+        # so we check for a reasonable minimum length.
+        assert len(choice.text) >= 1
+        # Finish reason might not always be 'length' if the model finishes early
+        # or due to other reasons, especially with high temperature.
+        # So, we'll accept 'length' or 'stop'.
+        assert choice.finish_reason in ("length", "stop")
+
+        # Token counts can also vary, so we check they are positive.
+        assert completion.usage.completion_tokens > 0
+        assert completion.usage.prompt_tokens > 0
+        assert completion.usage.total_tokens > 0
+        return completion
+
+    # Test single request to each node
+    for i, client in enumerate(clients):
+        result = await make_request(client)
+        assert result is not None
+        print(
+            f"Hybrid LB node {i} handled single completion request successfully"
+        )
+
+    await asyncio.sleep(0.5)
+
+    # Send requests to all nodes - each should balance within its local DP ranks
+    num_requests_per_node = 25  # Total 50 requests across 2 nodes
+    all_tasks = []
+
+    for i, client in enumerate(clients):
+        tasks = [make_request(client) for _ in range(num_requests_per_node)]
+        all_tasks.extend(tasks)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests_per_node * len(clients)
+    assert all(completion is not None for completion in results)
+
+    await asyncio.sleep(0.5)
+
+    # Second burst of requests
+    all_tasks = []
+    for i, client in enumerate(clients):
+        tasks = [make_request(client) for _ in range(num_requests_per_node)]
+        all_tasks.extend(tasks)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests_per_node * len(clients)
+    assert all(completion is not None for completion in results)
+
+    _, server_args = servers[0]
+    api_server_count = (
+        server_args.count('--api-server-count')
+        and server_args[server_args.index('--api-server-count') + 1] or 1)
+    print(
+        f"Successfully completed hybrid LB test with {len(clients)} nodes "
+        f"({DP_SIZE_LOCAL} DP ranks each, API server count: {api_server_count})"
+    )
+
+    # Check request balancing within each node
+    for i, (server, _) in enumerate(servers):
+        print(f"Checking request balancing for node {i}")
+        check_request_balancing(server, DP_SIZE_LOCAL)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_hybrid_lb_completion_streaming(clients: list[
+    openai.AsyncOpenAI], servers: list[tuple[RemoteOpenAIServer, list[str]]],
+                                              model_name: str) -> None:
+    prompt = "What is an LLM?"
+
+    async def make_streaming_request(client: openai.AsyncOpenAI):
+        # Perform a non-streaming request to get the expected full output
+        single_completion = await client.completions.create(
+            model=model_name,
+            prompt=prompt,
+            max_tokens=5,
+            temperature=0.0,
+        )
+        single_output = single_completion.choices[0].text
+
+        # Perform the streaming request
+        stream = await client.completions.create(model=model_name,
+                                                 prompt=prompt,
+                                                 max_tokens=5,
+                                                 temperature=0.0,
+                                                 stream=True)
+        chunks: list[str] = []
+        finish_reason_count = 0
+        last_chunk = None
+        async for chunk in stream:
+            chunks.append(chunk.choices[0].text)
+            if chunk.choices[0].finish_reason is not None:
+                finish_reason_count += 1
+            last_chunk = chunk  # Keep track of the last chunk
+
+        # finish reason should only return in the last block for OpenAI API
+        assert finish_reason_count == 1, (
+            "Finish reason should appear exactly once.")
+        assert last_chunk is not None, (
+            "Stream should have yielded at least one chunk.")
+        assert last_chunk.choices[
+            0].finish_reason == "length", "Finish reason should be 'length'."
+        # Check that the combined text matches the non-streamed version.
+        assert "".join(
+            chunks
+        ) == single_output, "Streamed output should match non-streamed output."
+        return True  # Indicate success for this request
+
+    # Test single request to each node
+    for i, client in enumerate(clients):
+        result = await make_streaming_request(client)
+        assert result is not None
+        print(
+            f"Hybrid LB node {i} handled single streaming request successfully"
+        )
+
+    await asyncio.sleep(0.5)
+
+    # Send streaming requests to all nodes
+    num_requests_per_node = 25  # Total 50 requests across 2 nodes
+    all_tasks = []
+
+    for i, client in enumerate(clients):
+        tasks = [
+            make_streaming_request(client)
+            for _ in range(num_requests_per_node)
+        ]
+        all_tasks.extend(tasks)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests_per_node * len(clients)
+    assert all(results), "Not all streaming requests completed successfully."
+
+    await asyncio.sleep(0.5)
+
+    # Second burst of streaming requests
+    all_tasks = []
+    for i, client in enumerate(clients):
+        tasks = [
+            make_streaming_request(client)
+            for _ in range(num_requests_per_node)
+        ]
+        all_tasks.extend(tasks)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests_per_node * len(clients)
+    assert all(results), "Not all streaming requests completed successfully."
+
+    _, server_args = servers[0]
+    api_server_count = (
+        server_args.count('--api-server-count')
+        and server_args[server_args.index('--api-server-count') + 1] or 1)
+    print(f"Successfully completed hybrid LB streaming test with "
+          f"{len(clients)} nodes ({DP_SIZE_LOCAL} DP ranks each, "
+          f"API server count: {api_server_count})")
+
+    # Check request balancing within each node
+    for i, (server, _) in enumerate(servers):
+        print(f"Checking streaming request balancing for node {i}")
+        check_request_balancing(server, DP_SIZE_LOCAL)
@@ -18,6 +18,8 @@ if TYPE_CHECKING:

 MODELS = [
    "Qwen/Qwen2.5-1.5B-Instruct",
+    # TODO: Enable this model when fixed.
+    # "Qwen/Qwen1.5-MoE-A2.7B",
    # TODO: Enable this models with v6e
    # "Qwen/Qwen2-7B-Instruct",
    # "meta-llama/Llama-3.1-8B",
@@ -460,11 +460,16 @@ def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
        {"load_config": {
            "load_format": original_load_format
        }})
-    model_runner_2.load_model()  # Load real weights inplace
+    model_runner_2.reload_weights()  # Load real weights inplace
    assert str(model_runner.get_model().state_dict()) == str(
        model_runner_2.get_model().state_dict())


+def test_reload_weights_before_load_model(model_runner):
+    with pytest.raises(AssertionError):
+        model_runner.reload_weights()
+
+
 def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():
    torch.set_default_dtype(torch.float16)
    layer_0 = "model.layers.0.self_attn.attn"
@@ -31,6 +31,5 @@ run_mypy vllm/inputs
 run_mypy vllm/lora
 run_mypy vllm/model_executor
 run_mypy vllm/plugins
-run_mypy vllm/prompt_adapter
 run_mypy vllm/worker
 run_mypy vllm/v1
@@ -1169,16 +1169,12 @@ class FlashInferImpl(AttentionImpl):
                    query=decode_query,
                    kv_cache=kv_cache.permute(*stride_order),
                    workspace_buffer=workspace_buffer,
-                    num_heads=num_heads,
-                    num_kv_heads=num_kv_heads,
-                    scale=softmax_scale,
                    block_tables=attn_metadata.block_tables,
                    seq_lens=decode_meta.seq_lens_tensor,
-                    block_size=attn_metadata.page_size,
                    max_seq_len=attn_metadata.max_decode_seq_len,
-                    kv_cache_dtype=kv_cache_dtype,
-                    k_scale=layer._k_scale_float,
-                    v_scale=layer._v_scale_float)
+                    bmm1_scale=layer._k_scale_float * softmax_scale,
+                    bmm2_scale=layer._v_scale_float,
+                )

        if prefill_output is None and decode_output is not None:
            # Decode only batch.
@@ -143,6 +143,8 @@ class Attention(nn.Module):
        # the backends)
        if envs.VLLM_USE_V1:
            self.use_irope = extra_impl_args.pop("use_irope", False)
+        else:
+            self.use_irope = extra_impl_args.get("use_irope", False)

        quant_method = quant_config.get_quant_method(
            self, prefix=prefix) if quant_config else None
@@ -177,7 +179,6 @@ class Attention(nn.Module):
                             kv_sharing_target_layer_name, **extra_impl_args)
        self.backend = backend_name_to_enum(attn_backend.get_name())
        self.dtype = dtype
-        self.use_irope = extra_impl_args.get("use_irope", False)

        # For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
        # torch.compile works by registering the attention as one giant
@@ -148,7 +148,9 @@ async def run_vllm_async(
    from vllm import SamplingParams

    async with build_async_engine_client_from_engine_args(
-            engine_args, disable_frontend_multiprocessing) as llm:
+        engine_args,
+        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
+    ) as llm:
        model_config = await llm.get_model_config()
        assert all(
            model_config.max_model_len >= (request.prompt_len +
@@ -423,6 +423,12 @@ class InductorAdaptor(CompilerInterface):
            if is_torch_equal_or_newer("2.6"):
                stack.enter_context(
                    torch._inductor.config.patch(fx_graph_remote_cache=False))
+                # InductorAdaptor (unfortunately) requires AOTAutogradCache
+                # to be turned off to run. It will fail to acquire the hash_str
+                # and error if not.
+                # StandaloneInductorAdaptor (PyTorch 2.8+) fixes this problem.
+                stack.enter_context(
+                    torch._functorch.config.patch(enable_autograd_cache=False))
                stack.enter_context(
                    torch._functorch.config.patch(
                        enable_remote_autograd_cache=False))
@@ -20,9 +20,38 @@ from .monitor import start_monitoring_torch_compile

 logger = init_logger(__name__)

+IGNORE_COMPILE_KEY = "_ignore_compile_vllm"
+
 _T = TypeVar("_T", bound=type[nn.Module])


+def ignore_torch_compile(cls: _T) -> _T:
+    """
+    A decorator to ignore support_torch_compile decorator
+    on the class. This is useful when a parent class has
+    a support_torch_compile decorator, but we don't want to
+    compile the class `cls` that inherits the parent class.
+    This only ignores compiling the forward of the class the
+    decorator is applied to. 
+
+    If the parent has ignore_torch_compile but the child has
+    support_torch_compile, the child will still be compiled.
+    
+    If the class has one or more submodules
+    that have support_torch_compile decorator applied, compile will
+    not be ignored for those submodules.
+    """
+    setattr(cls, IGNORE_COMPILE_KEY, True)
+    return cls
+
+
+def _should_ignore_torch_compile(cls) -> bool:
+    """
+    Check if the class should be ignored for torch.compile.
+    """
+    return getattr(cls, IGNORE_COMPILE_KEY, False)
+
+
@overload
 def support_torch_compile(
    *,
@@ -148,6 +177,8 @@ def _support_torch_compile(

    old_init = cls.__init__

+    setattr(cls, IGNORE_COMPILE_KEY, False)
+
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = '', **kwargs):
        old_init(self, vllm_config=vllm_config, prefix=prefix, **kwargs)
        self.vllm_config = vllm_config
@@ -156,9 +187,11 @@ def _support_torch_compile(
        self.do_not_compile = \
            vllm_config.compilation_config.level in [
            CompilationLevel.NO_COMPILATION, CompilationLevel.DYNAMO_AS_IS
-        ] or not supports_dynamo()
+        ] or not supports_dynamo() or _should_ignore_torch_compile(
+            self.__class__)
        if self.do_not_compile:
            return
+
        compilation_counter.num_models_seen += 1
        TorchCompileWrapperWithCustomDispatcher.__init__(
            self, compilation_level=vllm_config.compilation_config.level)
@@ -5,12 +5,15 @@ from torch import fx as fx

 from vllm.config import VllmConfig
 from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+if current_platform.is_cuda_alike():
+    from .fusion import FusionPass
+    from .collective_fusion import AllReduceFusionPass, AsyncTPPass
+    from .fusion_attn import AttnFusionPass

 from .activation_quant_fusion import ActivationQuantFusionPass
-from .collective_fusion import AllReduceFusionPass, AsyncTPPass
 from .fix_functionalization import FixFunctionalizationPass
-from .fusion import FusionPass
-from .fusion_attn import AttnFusionPass
 from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context
 from .noop_elimination import NoOpEliminationPass
 from .sequence_parallelism import SequenceParallelismPass
@@ -65,7 +65,7 @@ if TYPE_CHECKING:
    from vllm.model_executor.layers.quantization import QuantizationMethods
    from vllm.model_executor.layers.quantization.base_config import (
        QuantizationConfig)
-    from vllm.model_executor.model_loader import BaseModelLoader
+    from vllm.model_executor.model_loader import LoadFormats
    from vllm.model_executor.model_loader.tensorizer import TensorizerConfig

    ConfigType = type[DataclassInstance]
@@ -78,6 +78,7 @@ else:
    QuantizationConfig = Any
    QuantizationMethods = Any
    BaseModelLoader = Any
+    LoadFormats = Any
    TensorizerConfig = Any
    ConfigType = type
    HfOverrides = Union[dict[str, Any], Callable[[type], type]]
@@ -651,6 +652,8 @@ class ModelConfig:
        self.original_max_model_len = self.max_model_len
        self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
        self.multimodal_config = self._init_multimodal_config()
+        self.model_supports_multimodal_raw_input = (
+            self.registry.supports_multimodal_raw_input(self.architectures))
        if not self.skip_tokenizer_init:
            self._verify_tokenizer_mode()

@@ -1243,10 +1246,10 @@ class ModelConfig:
        return self.get_hf_config_sliding_window()

    def get_vocab_size(self) -> int:
-        return self.hf_text_config.vocab_size
+        return getattr(self.hf_text_config, "vocab_size", 0)

    def get_hidden_size(self) -> int:
-        return self.hf_text_config.hidden_size
+        return getattr(self.hf_text_config, "hidden_size", 0)

    @property
    def is_deepseek_mla(self) -> bool:
@@ -1771,29 +1774,12 @@ class CacheConfig:
            logger.warning("Possibly too large swap space. %s", msg)


-class LoadFormat(str, enum.Enum):
-    AUTO = "auto"
-    PT = "pt"
-    SAFETENSORS = "safetensors"
-    NPCACHE = "npcache"
-    DUMMY = "dummy"
-    TENSORIZER = "tensorizer"
-    SHARDED_STATE = "sharded_state"
-    GGUF = "gguf"
-    BITSANDBYTES = "bitsandbytes"
-    MISTRAL = "mistral"
-    RUNAI_STREAMER = "runai_streamer"
-    RUNAI_STREAMER_SHARDED = "runai_streamer_sharded"
-    FASTSAFETENSORS = "fastsafetensors"
-
-
@config
@dataclass
 class LoadConfig:
    """Configuration for loading the model weights."""

-    load_format: Union[str, LoadFormat,
-                       "BaseModelLoader"] = LoadFormat.AUTO.value
+    load_format: Union[str, LoadFormats] = "auto"
    """The format of the model weights to load:\n
    - "auto" will try to load the weights in the safetensors format and fall
    back to the pytorch bin format if safetensors format is not available.\n
@@ -1814,7 +1800,8 @@ class LoadConfig:
    - "gguf" will load weights from GGUF format files (details specified in
    https://github.com/ggml-org/ggml/blob/master/docs/gguf.md).\n
    - "mistral" will load weights from consolidated safetensors files used by
-    Mistral models."""
+    Mistral models.
+    - Other custom values can be supported via plugins."""
    download_dir: Optional[str] = None
    """Directory to download and load the weights, default to the default
    cache directory of Hugging Face."""
@@ -1862,10 +1849,7 @@ class LoadConfig:
        return hash_str

    def __post_init__(self):
-        if isinstance(self.load_format, str):
-            load_format = self.load_format.lower()
-            self.load_format = LoadFormat(load_format)
-
+        self.load_format = self.load_format.lower()
        if self.ignore_patterns is not None and len(self.ignore_patterns) > 0:
            logger.info(
                "Ignoring the following patterns when downloading weights: %s",
@@ -1906,8 +1890,16 @@ class ParallelConfig:
    """Backend to use for data parallel, either "mp" or "ray"."""
    data_parallel_external_lb: bool = False
    """Whether to use "external" DP LB mode. Applies only to online serving
-    and when data_parallel_size > 0. Set implicitly when
-    data_parallel_rank is provided explicitly to vllm serve."""
+    and when data_parallel_size > 0. This is useful for a "one-pod-per-rank"
+    wide-EP setup in Kuberentes. Set implicitly when --data-parallel-rank
+    is provided explicitly to vllm serve."""
+    data_parallel_hybrid_lb: bool = False
+    """Whether to use "hybrid" DP LB mode. Applies only to online serving
+    and when data_parallel_size > 0. Enables running an AsyncLLM
+    and API server on a "per-node" basis where vLLM load balances
+    between local data parallel ranks, but an external LB balances
+    between vLLM nodes/replicas. Set explicitly in conjunction with 
+    --data-parallel-start-rank."""
    enable_expert_parallel: bool = False
    """Use expert parallelism instead of tensor parallelism for MoE layers."""
    enable_eplb: bool = False
@@ -3141,59 +3133,6 @@ class LoRAConfig:
            self.lora_dtype = getattr(torch, self.lora_dtype)


-@config
-@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
-class PromptAdapterConfig:
-    """Configuration for PromptAdapters."""
-
-    max_prompt_adapters: int = 1
-    """Max number of PromptAdapters in a batch."""
-    max_prompt_adapter_token: int = 0
-    """Max number of PromptAdapters tokens."""
-    max_cpu_prompt_adapters: Optional[int] = None
-    """Maximum number of PromptAdapters to store in CPU memory. Must be >= than
-    `max_prompt_adapters`."""
-    prompt_adapter_dtype: Union[torch.dtype, str] = "auto"
-    """Data type for PromptAdapter. If auto, will default to base model dtype.
-    """
-
-    def compute_hash(self) -> str:
-        """
-        WARNING: Whenever a new field is added to this config,
-        ensure that it is included in the factors list if
-        it affects the computation graph.
-
-        Provide a hash that uniquely identifies all the configs
-        that affect the structure of the computation
-        graph from input ids/embeddings to the final hidden states,
-        excluding anything before input ids/embeddings and after
-        the final hidden states.
-        """
-        # no factors to consider.
-        # this config will not affect the computation graph.
-        factors: list[Any] = []
-        hash_str = hashlib.md5(str(factors).encode(),
-                               usedforsecurity=False).hexdigest()
-        return hash_str
-
-    def __post_init__(self):
-
-        if self.max_prompt_adapters < 1:
-            raise ValueError(f"max_prompt_adapters "
-                             f"({self.max_prompt_adapters}) must be >= 1.")
-        if self.max_prompt_adapter_token == 0:
-            raise ValueError("max_prompt_adapter_token must be set.")
-        if self.max_cpu_prompt_adapters is None:
-            self.max_cpu_prompt_adapters = self.max_prompt_adapters
-
-    def verify_with_model_config(self, model_config: ModelConfig):
-        if self.prompt_adapter_dtype == "auto":
-            self.prompt_adapter_dtype = model_config.dtype
-        elif isinstance(self.prompt_adapter_dtype, str):
-            self.prompt_adapter_dtype = getattr(torch,
-                                                self.prompt_adapter_dtype)
-
-
@config
@dataclass
 class MultiModalConfig:
@@ -4400,8 +4339,6 @@ class VllmConfig:
    """Decoding configuration."""
    observability_config: Optional[ObservabilityConfig] = None
    """Observability configuration."""
-    prompt_adapter_config: Optional[PromptAdapterConfig] = None
-    """Prompt adapter configuration."""
    quant_config: Optional[QuantizationConfig] = None
    """Quantization configuration."""
    compilation_config: CompilationConfig = field(
@@ -4498,10 +4435,6 @@ class VllmConfig:
            vllm_factors.append(self.observability_config.compute_hash())
        else:
            vllm_factors.append("None")
-        if self.prompt_adapter_config:
-            vllm_factors.append(self.prompt_adapter_config.compute_hash())
-        else:
-            vllm_factors.append("None")
        if self.quant_config:
            pass  # should be captured by model_config.quantization
        if self.compilation_config:
@@ -4609,9 +4542,6 @@ class VllmConfig:
        if self.lora_config is not None:
            self.lora_config.verify_with_cache_config(self.cache_config)
            self.lora_config.verify_with_model_config(self.model_config)
-        if self.prompt_adapter_config is not None:
-            self.prompt_adapter_config.verify_with_model_config(
-                self.model_config)

        if self.quant_config is None and self.model_config is not None:
            self.quant_config = VllmConfig._get_quantization_config(
@@ -15,7 +15,6 @@ from vllm.config import CacheConfig, LoRAConfig, SchedulerConfig
 from vllm.core.interfaces import AllocStatus, BlockSpaceManager
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
-from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sequence import (Sequence, SequenceData, SequenceGroup,
                           SequenceGroupBase, SequenceGroupMetadata,
                           SequenceGroupMetadataDelta, SequenceStage,
@@ -165,8 +164,6 @@ class SchedulerOutputs:
        if self.num_loras > 0:
            self._sort_by_lora_ids()

-        self.num_prompt_adapters: int = len(self.prompt_adapter_requests)
-
    def is_empty(self) -> bool:
        # NOTE: We do not consider the ignored sequence groups.
        return (not self.scheduled_seq_groups and not self.blocks_to_swap_in
@@ -194,14 +191,6 @@ class SchedulerOutputs:
            if g.seq_group.lora_request is not None
        }

-    @property
-    def prompt_adapter_requests(self) -> Set[PromptAdapterRequest]:
-        return {
-            g.seq_group.prompt_adapter_request
-            for g in self.scheduled_seq_groups
-            if g.seq_group.prompt_adapter_request is not None
-        }
-

@dataclass
 class SchedulerRunningOutputs:
@@ -1648,7 +1637,6 @@ class Scheduler:
                    multi_modal_placeholders=(
                        seq_group.multi_modal_placeholders
                        if scheduler_outputs.num_prefill_groups > 0 else None),
-                    prompt_adapter_request=seq_group.prompt_adapter_request,
                )
            else:
                # When SPMD mode is enabled, we only send delta data except for
@@ -120,8 +120,8 @@ class KVOutputAggregator:
    output corresponding to Rank 0 for scheduler."""

    def __init__(self, world_size: int):
-        # Complete transfer tracker. Used by to track finished requests
-        # [req_id -> n_finished_workers]
+        # Complete transfer tracker. Used to track finished requests
+        # [req_id -> n_remaining_workers]
        self._recv_remaining_count = defaultdict[str, int](lambda: world_size)
        self._send_remaining_count = defaultdict[str, int](lambda: world_size)

@@ -134,12 +134,10 @@ class KVOutputAggregator:
                                remaining_count_dict: dict[str, int],
                                finished_set: set[str]) -> None:
            for req_id in req_ids or ():
-                new_count = remaining_count_dict[req_id] - 1
-                if new_count == 0:
+                remaining_count_dict[req_id] -= 1
+                if remaining_count_dict[req_id] == 0:
                    finished_set.add(req_id)
                    del remaining_count_dict[req_id]
-                else:
-                    remaining_count_dict[req_id] = new_count

        finished_sending = set[str]()
        finished_recving = set[str]()
@@ -32,7 +32,7 @@ The class provides the following primitives:

 import enum
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Callable, Literal, Optional

 import torch

@@ -46,6 +46,12 @@ if TYPE_CHECKING:
    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
    from vllm.v1.request import Request

+# s_tensor_list, d_tensor_list, s_indices, d_indices, direction
+CopyBlocksOp = Callable[[
+    dict[str, torch.Tensor], dict[
+        str, torch.Tensor], list[int], list[int], Literal["h2d", "d2h"]
+], None]
+
 logger = init_logger(__name__)


@@ -127,6 +133,13 @@ class KVConnectorBase_V1(ABC):
        """
        return

+    def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
+        """
+        Set the xPU-specific ops for copying KV between host and device.
+        Needed when host buffer is used for kv transfer (e.g., in NixlConnector)
+        """
+        return
+
    @abstractmethod
    def start_load_kv(self, forward_context: "ForwardContext",
                      **kwargs) -> None:
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import contextlib
+import logging
 import math
 import queue
 import threading
@@ -20,14 +21,14 @@ from vllm import envs
 from vllm.attention.selector import backend_name_to_enum, get_attn_backend
 from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
-    KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
+    CopyBlocksOp, KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)
 from vllm.distributed.parallel_state import (
    get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,
    get_tp_group)
 from vllm.distributed.utils import divide
 from vllm.forward_context import ForwardContext
 from vllm.logger import init_logger
-from vllm.platforms import _Backend
+from vllm.platforms import _Backend, current_platform
 from vllm.utils import make_zmq_path, make_zmq_socket, round_down
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.request import RequestStatus
@@ -40,6 +41,7 @@ if TYPE_CHECKING:
 Transfer = tuple[int, float]  # (xfer_handle, start_time)
 EngineId = str
 ReqId = str
+
 GET_META_MSG = b"get_meta_msg"

 logger = init_logger(__name__)
@@ -52,6 +54,13 @@ except ImportError:
    logger.warning("NIXL is not available")
    NixlWrapper = None

+# Supported xPUs and types of kv transfer buffer.
+# {xPU: tuple of supported kv buffer types}
+_NIXL_SUPPORTED_XPUS = {
+    "cuda": ("cuda", ),
+    "tpu": ("cpu", ),
+}
+

 class NixlAgentMetadata(
        msgspec.Struct,
@@ -80,6 +89,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):

    def __init__(self):
        self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
+        self.reqs_to_save: dict[ReqId, ReqMeta] = {}
        self.reqs_to_send: dict[ReqId, float] = {}

    def add_new_req(
@@ -87,8 +97,12 @@ class NixlConnectorMetadata(KVConnectorMetadata):
        request_id: ReqId,
        local_block_ids: list[int],
        kv_transfer_params: dict[str, Any],
+        load_remote_cache: bool = True,
+        save_to_host: bool = False,
    ):
-        self.reqs_to_recv[request_id] = ReqMeta(
+        # save and load are mutually exclusive
+        assert load_remote_cache ^ save_to_host
+        _req = ReqMeta(
            local_block_ids=local_block_ids,
            remote_block_ids=kv_transfer_params["remote_block_ids"],
            remote_engine_id=kv_transfer_params["remote_engine_id"],
@@ -97,6 +111,10 @@ class NixlConnectorMetadata(KVConnectorMetadata):
            # P workers don't need to receive tp_size from proxy here.
            tp_size=kv_transfer_params.get("tp_size", 1),
        )
+        if save_to_host:
+            self.reqs_to_save[request_id] = _req
+        if load_remote_cache:
+            self.reqs_to_recv[request_id] = _req


 class NixlConnector(KVConnectorBase_V1):
@@ -155,6 +173,10 @@ class NixlConnector(KVConnectorBase_V1):
        assert self.connector_worker is not None
        self.connector_worker.register_kv_caches(kv_caches)

+    def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
+        assert self.connector_worker is not None
+        self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
+
    def get_finished(self,
                     finished_req_ids: set[str]) -> tuple[set[str], set[str]]:
        """Get the finished recving and sending requests."""
@@ -177,8 +199,11 @@ class NixlConnector(KVConnectorBase_V1):
        pass

    def wait_for_save(self):
-        """NixlConnector does not save explicitly."""
-        pass
+        assert self.connector_worker is not None
+        assert isinstance(self._connector_metadata, NixlConnectorMetadata)
+        if self.connector_worker.use_host_buffer and \
+           self.connector_worker.copy_blocks:
+            self.connector_worker.save_kv_to_host(self._connector_metadata)


 class NixlConnectorScheduler:
@@ -193,12 +218,15 @@ class NixlConnectorScheduler:
            envs.VLLM_NIXL_SIDE_CHANNEL_PORT +
            vllm_config.parallel_config.data_parallel_rank *
            vllm_config.parallel_config.tensor_parallel_size)
+        self.use_host_buffer = \
+            vllm_config.kv_transfer_config.kv_buffer_device == "cpu"
        logger.info("Initializing NIXL Scheduler %s", engine_id)

        # Requests that need to start recv/send.
        # New requests are added by update_state_after_alloc in
        # the scheduler. Used to make metadata passed to Worker.
        self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
+        self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
        # Reqs to send and their expiration time
        self._reqs_need_send: dict[ReqId, float] = {}

@@ -248,7 +276,25 @@ class NixlConnectorScheduler:
            "num_external_tokens=%s, kv_transfer_params=%s",
            num_external_tokens, params)

-        if params is not None and params.get("do_remote_prefill"):
+        if not params:
+            return
+        if self.use_host_buffer and params.get("do_remote_decode"):
+            # NOTE: when accelerator is not directly supported by Nixl,
+            # prefilled blocks need to be saved to host memory before transfer.
+
+            # figure out full computed blocks to save
+            block_ids = blocks.get_block_ids()[0]
+            all_full = request.num_tokens % self.block_size == 0
+            full_block_ids = (block_ids if all_full else block_ids[:-1])
+            # TODO: skip the blocks that are already in the host xfer buffer.
+            # Currently, the host xfer buffer block is 1-to-1 mapped to device
+            # kv blocks, so host blocks won't be flushed as long as its device
+            # block is not overwritten; and it will be safe to skip saving them
+            # to host xfer buffer.
+            if full_block_ids:
+                self._reqs_need_save[request.request_id] = \
+                    (request, full_block_ids)
+        elif params.get("do_remote_prefill"):
            if params.get("remote_block_ids"):
                if all(p in params for p in ("remote_engine_id", "remote_host",
                                             "remote_port")):
@@ -260,6 +306,7 @@ class NixlConnectorScheduler:
                    # Get unhashed blocks to pull from remote.
                    self._reqs_need_recv[request.request_id] = (
                        request, local_block_ids)
+
                else:
                    logger.warning(
                        "Got invalid KVTransferParams: %s. This "
@@ -284,10 +331,21 @@ class NixlConnectorScheduler:
                kv_transfer_params=req.kv_transfer_params,
            )

-        # Clear the list once workers start the transfers
-        self._reqs_need_recv.clear()
+        for req_id, (req, block_ids) in self._reqs_need_save.items():
+            assert req.kv_transfer_params is not None
+            meta.add_new_req(
+                request_id=req_id,
+                local_block_ids=block_ids,
+                kv_transfer_params=req.kv_transfer_params,
+                load_remote_cache=False,
+                save_to_host=True,
+            )

        meta.reqs_to_send = self._reqs_need_send
+
+        # Clear the list once workers start the transfers
+        self._reqs_need_recv.clear()
+        self._reqs_need_save.clear()
        self._reqs_need_send = {}

        return meta
@@ -379,9 +437,36 @@ class NixlConnectorWorker:
        self.tp_rank = get_tensor_model_parallel_rank()
        self.world_size = get_tensor_model_parallel_world_size()
        self.tp_group = get_tp_group()
+        self.num_blocks = 0

        # KV Caches and nixl tracking data.
-        self.kv_caches: dict[str, torch.Tensor] = {}
+        self.device_type = current_platform.device_type
+        self.kv_buffer_device: str = \
+            vllm_config.kv_transfer_config.kv_buffer_device
+        if self.device_type not in _NIXL_SUPPORTED_XPUS:
+            raise RuntimeError(f"{self.device_type} is not supported.")
+        elif self.kv_buffer_device not in _NIXL_SUPPORTED_XPUS[
+                self.device_type]:
+            raise RuntimeError(
+                f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
+                "is not supported.")
+        self.device_kv_caches: dict[str, torch.Tensor] = {}
+
+        # cpu kv buffer for xfer
+        # used when xPU memory can not be registered under nixl
+        self.host_xfer_buffers: dict[str, torch.Tensor] = {}
+        self.use_host_buffer = self.kv_buffer_device == "cpu"
+        if self.kv_buffer_device == "cuda":
+            self.nixl_memory_type = "VRAM"
+        elif self.kv_buffer_device == "cpu":
+            self.nixl_memory_type = "DRAM"
+        else:
+            raise RuntimeError(
+                f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
+                "is not supported.")
+
+        # Note: host xfer buffer ops when use_host_buffer is True
+        self.copy_blocks: Optional[CopyBlocksOp] = None

        # Map of engine_id -> kv_caches_base_addr. For TP case, each local
        # rank will still only pull from a single remote TP worker.
@@ -404,6 +489,7 @@ class NixlConnectorWorker:

        # In progress transfers.
        # [req_id -> list[handle]]
+        self._recving_metadata: dict[ReqId, ReqMeta] = {}
        self._recving_transfers = defaultdict[ReqId, list[Transfer]](list)
        # Track the expiration time of requests that are waiting to be sent.
        self._reqs_to_send: dict[ReqId, float] = {}
@@ -440,6 +526,7 @@ class NixlConnectorWorker:
        self.backend_name = backend.get_name()
        attn_backend = backend_name_to_enum(self.backend_name)
        self._use_flashinfer = attn_backend == _Backend.FLASHINFER_VLLM_V1
+        self._use_pallas_v1 = attn_backend == _Backend.PALLAS_VLLM_V1
        logger.debug("Detected attention backend %s", self.backend_name)

        self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
@@ -529,6 +616,31 @@ class NixlConnectorWorker:
        # Remote rank -> agent name.
        return {p_remote_rank: remote_agent_name}

+    def initialize_host_xfer_buffer(
+            self, kv_caches: dict[str, torch.Tensor]) -> None:
+        """
+        Initialize transfer buffer in CPU mem for accelerators
+        NOT directly supported by NIXL (e.g., tpu)
+        """
+        xfer_buffers: dict[str, torch.Tensor] = {}
+        try:
+            for layer_name, kv_cache in kv_caches.items():
+                kv_shape = kv_cache.shape
+                kv_dtype = kv_cache.dtype
+                xfer_buffers[layer_name] = torch.empty(kv_shape,
+                                                       dtype=kv_dtype,
+                                                       device="cpu")
+        except MemoryError as e:
+            logger.error("NIXLConnectorWorker gets %s.", e)
+            raise
+
+        self.host_xfer_buffers = xfer_buffers
+
+    def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
+        """Assign copy (d2h, h2d) operations when host buffer is used."""
+        assert self.use_host_buffer
+        self.copy_blocks = copy_operation
+
    def _background_nixl_handshake(self, req_id: str,
                                   remote_engine_id: EngineId, meta: ReqMeta):
        # Do NIXL handshake in background and add to _ready_requests when done.
@@ -562,47 +674,76 @@ class NixlConnectorWorker:
        _, first_kv_cache = next(iter(kv_caches.items()))
        kv_elem_size = first_kv_cache.element_size()

+        if self.use_host_buffer:
+            self.initialize_host_xfer_buffer(kv_caches=kv_caches)
+            assert len(self.host_xfer_buffers) == len(kv_caches), (
+                f"host_buffer: {len(self.host_xfer_buffers)}, "
+                f"kv_caches: {len(kv_caches)}")
+            xfer_buffers = self.host_xfer_buffers
+        else:
+            xfer_buffers = kv_caches
+            assert not self.host_xfer_buffers, (
+                "host_xfer_buffer should not be initialized when "
+                f"kv_buffer_device is {self.kv_buffer_device}")
+
        # TODO(tms): Find a more robust way to detect and handle MLA
        # NOTE (NickLucche) To move blocks efficiently with NIXL, the expected
        # KV memory layout is HND, as opposed to the default NHD. Note that it
        # will only affects the strides. For MLA instead, we make require no
        # such thing and resort to the standard layout.
        use_mla = len(first_kv_cache.shape) == 3
-        assert use_mla == self.use_mla
-
-        # TODO (NickLucche) not compatible with hybrid allocator. Enforce check
-        # once it goes live, as a single kv layout is expected for xfers.
-        if use_mla:
-            # MLA case.
+        if self.device_type == "tpu":
+            assert not use_mla, f"{self.kv_buffer_device} does not support MLA."
+            assert self._use_pallas_v1, f"attn backend: {self.backend_name}"
+            # tpu (v1) kv shape per layer:
+            # (num_blocks, block_size, num_kv_heads * 2, head_size)
            self.num_blocks = first_kv_cache.shape[0]
-            block_rank = 2  # [block_size, latent_dim]
+            block_rank = 3  # [block_size, kv_heads, head_dim]
            block_shape = first_kv_cache.shape[-block_rank:]
-            block_size, kv_latent_dim = block_shape
-            self.slot_size_bytes = kv_elem_size * kv_latent_dim
-        else:
-            # [2 (k and v), num_blocks, ...]
-            if self._use_flashinfer:
-                # FlashInfer swaps 2<->num_blocks dimensions.
+            block_size, n_kv_heads_x_2, head_dim = block_shape
+            self.slot_size_bytes = kv_elem_size * n_kv_heads_x_2 * head_dim
+        elif self.device_type == "cuda":
+            assert use_mla == self.use_mla
+            # TODO (NickLucche) not compatible with hybrid allocator.
+            # Enforce check once it goes live, as a single kv layout
+            # is expected for xfers.
+            if use_mla:
+                # MLA case.
                self.num_blocks = first_kv_cache.shape[0]
-                block_rank = 4  # [2, block_size, kv_heads, head_dim]
+                block_rank = 2  # [block_size, latent_dim]
+                block_shape = first_kv_cache.shape[-block_rank:]
+                block_size, kv_latent_dim = block_shape
+                self.slot_size_bytes = kv_elem_size * kv_latent_dim
            else:
-                self.num_blocks = first_kv_cache.shape[1]
-                block_rank = 3  # [block_size, kv_heads, head_dim]
-            block_shape = first_kv_cache.shape[-block_rank:]
-            block_size, n_kv_heads, head_dim = block_shape[-3:]
-            # head size in bytes.
-            self.slot_size_bytes = kv_elem_size * n_kv_heads * head_dim
-        assert block_size == self.block_size
+                # [2 (k and v), num_blocks, ...]
+                if self._use_flashinfer:
+                    # FlashInfer swaps 2<->num_blocks dimensions.
+                    self.num_blocks = first_kv_cache.shape[0]
+                    block_rank = 4  # [2, block_size, kv_heads, head_dim]
+                else:
+                    self.num_blocks = first_kv_cache.shape[1]
+                    block_rank = 3  # [block_size, kv_heads, head_dim]
+                block_shape = first_kv_cache.shape[-block_rank:]
+                block_size, n_kv_heads, head_dim = block_shape[-3:]
+                # head size in bytes.
+                self.slot_size_bytes = kv_elem_size * n_kv_heads * head_dim
+            assert block_size == self.block_size
+        else:
+            raise RuntimeError(
+                f"{self.device_type} ({self.backend_name}) is not supported.")
+
        # TODO(tms): self.block_len needs to be per-layer for sliding window,
        # hybrid attn, etc
        # block size in bytes
        self.block_len = kv_elem_size * math.prod(block_shape)
        logger.info(
-            "Registering KV_Caches: use_mla: %s, num_blocks: %s, "
-            "block_shape: %s, per_layer_kv_cache_shape: %s", use_mla,
-            self.num_blocks, block_shape, first_kv_cache.shape)
+            "Registering KV_Caches. use_mla: %s, kv_buffer_device: %s, "
+            "use_host_buffer: %s, num_blocks: %s, block_shape: %s, "
+            "per_layer_kv_cache_shape: %s", use_mla, self.kv_buffer_device,
+            self.use_host_buffer, self.num_blocks, block_shape,
+            first_kv_cache.shape)
        self.dst_num_blocks[self.engine_id] = self.num_blocks
-        self.kv_caches = kv_caches
+        self.device_kv_caches = kv_caches
        kv_caches_base_addr = []
        caches_data = []

@@ -614,19 +755,21 @@ class NixlConnectorWorker:
        # (roughly 8KB vs 5KB).
        # Conversely for FlashInfer, K and V are transferred in the same tensor
        # to better exploit the memory layout (ie num_blocks is the first dim).
-        for cache_or_caches in kv_caches.values():
+        for cache_or_caches in xfer_buffers.values():
            # Normalize to always be a list of caches
-            cache_list = [cache_or_caches] if use_mla or self._use_flashinfer \
-                else cache_or_caches
+            cache_list = [cache_or_caches] if use_mla \
+                         or self._use_pallas_v1 or self._use_flashinfer \
+                         else cache_or_caches
            for cache in cache_list:
                base_addr = cache.data_ptr()
                region_len = self.num_blocks * self.block_len
-                caches_data.append(
-                    (base_addr, region_len, cache.device.index, ""))
+                # NOTE: use tp_rank for device_id since multi-node TP
+                # is rarely used.
+                caches_data.append((base_addr, region_len, self.tp_rank, ""))
                kv_caches_base_addr.append(base_addr)
        self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
        self.num_regions = len(caches_data)
-        self.num_layers = len(self.kv_caches.keys())
+        self.num_layers = len(xfer_buffers.keys())

        # TODO(mgoin): remove this once we have hybrid memory allocator
        # Optimization for models with local attention (Llama 4)
@@ -648,7 +791,8 @@ class NixlConnectorWorker:
                         self.block_window_per_layer)
            assert len(self.block_window_per_layer) == self.num_layers

-        descs = self.nixl_wrapper.get_reg_descs(caches_data, "VRAM")
+        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)
        logger.debug("Done registering descs")
@@ -666,11 +810,13 @@ class NixlConnectorWorker:
                block_offset = block_id * self.block_len
                addr = base_addr + block_offset
                # (addr, len, device id)
+                # TODO: does device_id matter to DRAM?
                blocks_data.append((addr, self.block_len, self.tp_rank))
        logger.debug("Created %s blocks for src engine %s and rank %s",
                     len(blocks_data), self.engine_id, self.tp_rank)

-        descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
+        descs = self.nixl_wrapper.get_xfer_descs(blocks_data,
+                                                 self.nixl_memory_type)
        # NIXL_INIT_AGENT to be used for preparations of local descs.
        self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(
            "NIXL_INIT_AGENT", descs)
@@ -755,6 +901,8 @@ class NixlConnectorWorker:
        tp_ratio = divide(self._tp_size[self.engine_id],
                          self._tp_size[engine_id])
        assert tp_ratio > 0, "Decode TP cannot be smaller than prefill TP"
+        assert not self._use_pallas_v1 or tp_ratio == 1, \
+               "TPU (pallas_v1) DOES NOT support heterogeneous TP yet."

        # Handle tp_size>num_kv_heads: replicate KV cache.
        total_num_kv_heads = self.model_config.get_total_num_kv_heads()
@@ -813,13 +961,43 @@ class NixlConnectorWorker:
            self.tp_rank)

        # Register with NIXL.
-        descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")
+        descs = self.nixl_wrapper.get_xfer_descs(blocks_data,
+                                                 self.nixl_memory_type)
        self.dst_xfer_side_handles[
            engine_id] = self.nixl_wrapper.prep_xfer_dlist(
                remote_agent_name, descs)

        return remote_agent_name

+    def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
+        """copy recved kv from host buffer to device."""
+        assert self.use_host_buffer
+        assert self.copy_blocks is not None
+
+        local_block_ids = meta.local_block_ids
+        self.copy_blocks(self.host_xfer_buffers, self.device_kv_caches,
+                         local_block_ids, local_block_ids, "h2d")
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(
+                "synced recved kv of request[%s] to device kv buffer,"
+                "local_block_ids: %s. ", req_id,
+                ",".join(map(str, meta.local_block_ids)))
+
+    def save_kv_to_host(self, metadata: NixlConnectorMetadata):
+        """copy kv from device to host buffer."""
+        assert self.use_host_buffer
+        assert self.copy_blocks is not None
+
+        for req_id, meta in metadata.reqs_to_save.items():
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.debug(
+                    "save_load_kv for request[%s] to host xfer buffer."
+                    "local_block_ids: %s. ", req_id,
+                    ",".join(map(str, meta.local_block_ids)))
+            # blocking
+            self.copy_blocks(self.device_kv_caches, self.host_xfer_buffers,
+                             meta.local_block_ids, meta.local_block_ids, "d2h")
+
    def get_finished(self) -> tuple[set[str], set[str]]:
        """
        Get requests that are done sending or recving on this specific worker.
@@ -834,6 +1012,12 @@ class NixlConnectorWorker:
                "and %s requests done recving", self.tp_rank,
                len(done_sending), len(done_recving))

+        if self.use_host_buffer:
+            for req_id in done_recving:
+                meta = self._recving_metadata.pop(req_id)
+                assert meta, f"{req_id} not found in recving_metadata list"
+                self.sync_recved_kv_to_device(req_id, meta)
+
        # Handle timeout to avoid stranding blocks on remote.
        now = time.perf_counter()
        while self._reqs_to_send:
@@ -904,6 +1088,8 @@ class NixlConnectorWorker:
                "Num local_block_ids: %s. Num remote_block_ids: %s. ", req_id,
                remote_engine_id, len(meta.local_block_ids),
                len(meta.remote_block_ids))
+            if self.use_host_buffer:
+                self._recving_metadata[req_id] = meta
            if remote_engine_id not in self._remote_agents:
                # Initiate handshake with remote engine to exchange metadata.
                with self._handshake_lock:
@@ -26,13 +26,12 @@ from vllm.config import (BlockSize, CacheConfig, CacheDType, CompilationConfig,
                         DetailedTraceModules, Device, DeviceConfig,
                         DistributedExecutorBackend, GuidedDecodingBackend,
                         GuidedDecodingBackendV1, HfOverrides, KVEventsConfig,
-                         KVTransferConfig, LoadConfig, LoadFormat,
-                         LogprobsMode, LoRAConfig, ModelConfig, ModelDType,
-                         ModelImpl, MultiModalConfig, ObservabilityConfig,
-                         ParallelConfig, PoolerConfig, PrefixCachingHashAlgo,
-                         PromptAdapterConfig, SchedulerConfig, SchedulerPolicy,
-                         SpeculativeConfig, TaskOption, TokenizerMode,
-                         VllmConfig, get_attr_docs, get_field)
+                         KVTransferConfig, LoadConfig, LogprobsMode,
+                         LoRAConfig, ModelConfig, ModelDType, ModelImpl,
+                         MultiModalConfig, ObservabilityConfig, ParallelConfig,
+                         PoolerConfig, PrefixCachingHashAlgo, SchedulerConfig,
+                         SchedulerPolicy, SpeculativeConfig, TaskOption,
+                         TokenizerMode, VllmConfig, get_attr_docs, get_field)
 from vllm.logger import init_logger
 from vllm.platforms import CpuArchEnum, current_platform
 from vllm.plugins import load_general_plugins
@@ -47,10 +46,12 @@ from vllm.utils import (STR_DUAL_CHUNK_FLASH_ATTN_VAL, FlexibleArgumentParser,
 if TYPE_CHECKING:
    from vllm.executor.executor_base import ExecutorBase
    from vllm.model_executor.layers.quantization import QuantizationMethods
+    from vllm.model_executor.model_loader import LoadFormats
    from vllm.usage.usage_lib import UsageContext
 else:
    ExecutorBase = Any
    QuantizationMethods = Any
+    LoadFormats = Any
    UsageContext = Any

 logger = init_logger(__name__)
@@ -276,7 +277,7 @@ class EngineArgs:
    trust_remote_code: bool = ModelConfig.trust_remote_code
    allowed_local_media_path: str = ModelConfig.allowed_local_media_path
    download_dir: Optional[str] = LoadConfig.download_dir
-    load_format: str = LoadConfig.load_format
+    load_format: Union[str, LoadFormats] = LoadConfig.load_format
    config_format: str = ModelConfig.config_format
    dtype: ModelDType = ModelConfig.dtype
    kv_cache_dtype: CacheDType = CacheConfig.cache_dtype
@@ -295,9 +296,11 @@ class EngineArgs:
    tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
    data_parallel_size: int = ParallelConfig.data_parallel_size
    data_parallel_rank: Optional[int] = None
+    data_parallel_start_rank: Optional[int] = None
    data_parallel_size_local: Optional[int] = None
    data_parallel_address: Optional[str] = None
    data_parallel_rpc_port: Optional[int] = None
+    data_parallel_hybrid_lb: bool = False
    data_parallel_backend: str = ParallelConfig.data_parallel_backend
    enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
    enable_eplb: bool = ParallelConfig.enable_eplb
@@ -358,11 +361,6 @@ class EngineArgs:
    max_cpu_loras: Optional[int] = LoRAConfig.max_cpu_loras
    lora_dtype: Optional[Union[str, torch.dtype]] = LoRAConfig.lora_dtype
    lora_extra_vocab_size: int = LoRAConfig.lora_extra_vocab_size
-    # PromptAdapter fields
-    enable_prompt_adapter: bool = False
-    max_prompt_adapters: int = PromptAdapterConfig.max_prompt_adapters
-    max_prompt_adapter_token: int = \
-        PromptAdapterConfig.max_prompt_adapter_token

    num_scheduler_steps: int = SchedulerConfig.num_scheduler_steps
    multi_step_stream_outputs: bool = SchedulerConfig.multi_step_stream_outputs
@@ -437,6 +435,8 @@ class EngineArgs:
        ParallelConfig.enable_multimodal_encoder_data_parallel

    async_scheduling: bool = SchedulerConfig.async_scheduling
+    # DEPRECATED
+    enable_prompt_adapter: bool = False

    def __post_init__(self):
        # support `EngineArgs(compilation_config={...})`
@@ -548,9 +548,7 @@ class EngineArgs:
            title="LoadConfig",
            description=LoadConfig.__doc__,
        )
-        load_group.add_argument("--load-format",
-                                choices=[f.value for f in LoadFormat],
-                                **load_kwargs["load_format"])
+        load_group.add_argument("--load-format", **load_kwargs["load_format"])
        load_group.add_argument("--download-dir",
                                **load_kwargs["download_dir"])
        load_group.add_argument("--model-loader-extra-config",
@@ -607,6 +605,11 @@ class EngineArgs:
            type=int,
            help='Data parallel rank of this instance. '
            'When set, enables external load balancer mode.')
+        parallel_group.add_argument('--data-parallel-start-rank',
+                                    '-dpr',
+                                    type=int,
+                                    help='Starting data parallel rank '
+                                    'for secondary nodes.')
        parallel_group.add_argument('--data-parallel-size-local',
                                    '-dpl',
                                    type=int,
@@ -628,6 +631,9 @@ class EngineArgs:
                                    default='mp',
                                    help='Backend for data parallel, either '
                                    '"mp" or "ray".')
+        parallel_group.add_argument(
+            "--data-parallel-hybrid-lb",
+            **parallel_kwargs["data_parallel_hybrid_lb"])
        parallel_group.add_argument(
            "--enable-expert-parallel",
            **parallel_kwargs["enable_expert_parallel"])
@@ -729,23 +735,6 @@ class EngineArgs:
        lora_group.add_argument("--default-mm-loras",
                                **lora_kwargs["default_mm_loras"])

-        # PromptAdapter related configs
-        prompt_adapter_kwargs = get_kwargs(PromptAdapterConfig)
-        prompt_adapter_group = parser.add_argument_group(
-            title="PromptAdapterConfig",
-            description=PromptAdapterConfig.__doc__,
-        )
-        prompt_adapter_group.add_argument(
-            "--enable-prompt-adapter",
-            action=argparse.BooleanOptionalAction,
-            help="If True, enable handling of PromptAdapters.")
-        prompt_adapter_group.add_argument(
-            "--max-prompt-adapters",
-            **prompt_adapter_kwargs["max_prompt_adapters"])
-        prompt_adapter_group.add_argument(
-            "--max-prompt-adapter-token",
-            **prompt_adapter_kwargs["max_prompt_adapter_token"])
-
        # Speculative arguments
        speculative_group = parser.add_argument_group(
            title="SpeculativeConfig",
@@ -850,6 +839,12 @@ class EngineArgs:
        parser.add_argument('--disable-log-stats',
                            action='store_true',
                            help='Disable logging statistics.')
+        parser.add_argument('--enable-prompt-adapter',
+                            action='store_true',
+                            deprecated=True,
+                            help='[DEPRECATED] Prompt adapter has been '
+                            'removed. Setting this flag to True or False'
+                            ' has no effect on vLLM behavior.')

        return parser

@@ -868,10 +863,9 @@ class EngineArgs:

        # NOTE: This is to allow model loading from S3 in CI
        if (not isinstance(self, AsyncEngineArgs) and envs.VLLM_CI_USE_S3
-                and self.model in MODELS_ON_S3
-                and self.load_format == LoadFormat.AUTO):  # noqa: E501
+                and self.model in MODELS_ON_S3 and self.load_format == "auto"):
            self.model = f"{MODEL_WEIGHTS_S3_BUCKET}/{self.model}"
-            self.load_format = LoadFormat.RUNAI_STREAMER
+            self.load_format = "runai_streamer"

        return ModelConfig(
            model=self.model,
@@ -986,6 +980,7 @@ class EngineArgs:
    def create_engine_config(
        self,
        usage_context: Optional[UsageContext] = None,
+        headless: bool = False,
    ) -> VllmConfig:
        """
        Create the VllmConfig.
@@ -1074,15 +1069,41 @@ class EngineArgs:
            # but we should not do this here.
            placement_group = ray.util.get_current_placement_group()

+        assert not headless or not self.data_parallel_hybrid_lb, (
+            "data_parallel_hybrid_lb is not applicable in "
+            "headless mode")
+
        data_parallel_external_lb = self.data_parallel_rank is not None
+        # Local DP rank = 1, use pure-external LB.
        if data_parallel_external_lb:
            assert self.data_parallel_size_local in (1, None), (
                "data_parallel_size_local must be 1 when data_parallel_rank "
                "is set")
            data_parallel_size_local = 1
+            # Use full external lb if we have local_size of 1.
+            self.data_parallel_hybrid_lb = False
        elif self.data_parallel_size_local is not None:
            data_parallel_size_local = self.data_parallel_size_local
+
+            if self.data_parallel_start_rank and not headless:
+                # Infer hybrid LB mode.
+                self.data_parallel_hybrid_lb = True
+
+            if self.data_parallel_hybrid_lb and data_parallel_size_local == 1:
+                # Use full external lb if we have local_size of 1.
+                data_parallel_external_lb = True
+                self.data_parallel_hybrid_lb = False
+
+            if data_parallel_size_local == self.data_parallel_size:
+                # Disable hybrid LB mode if set for a single node
+                self.data_parallel_hybrid_lb = False
+
+            self.data_parallel_rank = self.data_parallel_start_rank or 0
        else:
+            assert not self.data_parallel_hybrid_lb, (
+                "data_parallel_size_local must be set to use "
+                "data_parallel_hybrid_lb.")
+
            # Local DP size defaults to global DP size if not set.
            data_parallel_size_local = self.data_parallel_size

@@ -1139,6 +1160,7 @@ class EngineArgs:
            data_parallel_master_ip=data_parallel_address,
            data_parallel_rpc_port=data_parallel_rpc_port,
            data_parallel_backend=self.data_parallel_backend,
+            data_parallel_hybrid_lb=self.data_parallel_hybrid_lb,
            enable_expert_parallel=self.enable_expert_parallel,
            enable_eplb=self.enable_eplb,
            num_redundant_experts=self.num_redundant_experts,
@@ -1234,11 +1256,6 @@ class EngineArgs:

        load_config = self.create_load_config()

-        prompt_adapter_config = PromptAdapterConfig(
-            max_prompt_adapters=self.max_prompt_adapters,
-            max_prompt_adapter_token=self.max_prompt_adapter_token) \
-                                        if self.enable_prompt_adapter else None
-
        decoding_config = DecodingConfig(
            backend=self.guided_decoding_backend,
            disable_fallback=self.guided_decoding_disable_fallback,
@@ -1266,7 +1283,6 @@ class EngineArgs:
            load_config=load_config,
            decoding_config=decoding_config,
            observability_config=observability_config,
-            prompt_adapter_config=prompt_adapter_config,
            compilation_config=self.compilation_config,
            kv_transfer_config=self.kv_transfer_config,
            kv_events_config=self.kv_events_config,
@@ -1281,7 +1297,7 @@ class EngineArgs:
        #############################################################
        # Unsupported Feature Flags on V1.

-        if self.load_format == LoadFormat.SHARDED_STATE.value:
+        if self.load_format == "sharded_state":
            _raise_or_fallback(
                feature_name=f"--load_format {self.load_format}",
                recommend_to_remove=False)
@@ -1342,12 +1358,6 @@ class EngineArgs:
                                   recommend_to_remove=False)
                return False

-        # No Prompt Adapter so far.
-        if self.enable_prompt_adapter:
-            _raise_or_fallback(feature_name="--enable-prompt-adapter",
-                               recommend_to_remove=False)
-            return False
-
        # No text embedding inputs so far.
        if self.enable_prompt_embeds:
            _raise_or_fallback(feature_name="--enable-prompt-embeds",
@@ -1469,7 +1479,6 @@ class EngineArgs:

                if (is_gpu and not use_sliding_window and not use_spec_decode
                        and not self.enable_lora
-                        and not self.enable_prompt_adapter
                        and model_config.runner_type != "pooling"):
                    self.enable_chunked_prefill = True
                    logger.warning(
@@ -29,7 +29,6 @@ from vllm.model_executor.guided_decoding import (
 from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.outputs import PoolingRequestOutput, RequestOutput
 from vllm.pooling_params import PoolingParams
-from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
 from vllm.sequence import ExecuteModelRequest
 from vllm.transformers_utils.tokenizer import AnyTokenizer
@@ -435,7 +434,6 @@ class _AsyncLLMEngine(LLMEngine):
        arrival_time: Optional[float] = None,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
        data_parallel_rank: Optional[int] = None,
        tokenization_kwargs: Optional[dict[str, Any]] = None,
@@ -468,7 +466,6 @@ class _AsyncLLMEngine(LLMEngine):
        processed_inputs = await self.input_preprocessor.preprocess_async(
            prompt,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
            tokenization_kwargs=tokenization_kwargs,
        )

@@ -491,7 +488,6 @@ class _AsyncLLMEngine(LLMEngine):
            params=params,
            arrival_time=arrival_time,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
            trace_headers=trace_headers,
            priority=priority,
        )
@@ -861,7 +857,6 @@ class AsyncLLMEngine(EngineClient):
        arrival_time: Optional[float] = None,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
        data_parallel_rank: Optional[int] = None,
        tokenization_kwargs: Optional[dict[str, Any]] = None,
@@ -889,7 +884,6 @@ class AsyncLLMEngine(EngineClient):
            arrival_time=arrival_time or time.time(),
            lora_request=lora_request,
            trace_headers=trace_headers,
-            prompt_adapter_request=prompt_adapter_request,
            priority=priority,
            data_parallel_rank=data_parallel_rank,
            tokenization_kwargs=tokenization_kwargs,
@@ -904,7 +898,6 @@ class AsyncLLMEngine(EngineClient):
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
        data_parallel_rank: Optional[int] = None,
    ) -> AsyncGenerator[RequestOutput, None]:
@@ -922,8 +915,6 @@ class AsyncLLMEngine(EngineClient):
            request_id: The unique id of the request.
            lora_request: LoRA request to use for generation, if any.
            trace_headers: OpenTelemetry trace headers.
-            prompt_adapter_request: Prompt Adapter request to use
-                                            for generation, if any.
            priority: The priority of the request.
                Only applicable with priority scheduling.
            data_parallel_rank: The (global) data parallel rank that must
@@ -983,7 +974,6 @@ class AsyncLLMEngine(EngineClient):
                    sampling_params,
                    lora_request=lora_request,
                    trace_headers=trace_headers,
-                    prompt_adapter_request=prompt_adapter_request,
                    priority=priority,
                    data_parallel_rank=data_parallel_rank,
            ):
@@ -44,7 +44,6 @@ from vllm.multimodal.processing import EncDecMultiModalProcessor
 from vllm.outputs import (PoolingRequestOutput, RequestOutput,
                          RequestOutputFactory)
 from vllm.pooling_params import PoolingParams
-from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.sequence import (ExecuteModelRequest, ParallelSampleSequenceGroup,
                           PoolingSequenceGroupOutput, Sequence, SequenceGroup,
@@ -223,7 +222,6 @@ class LLMEngine:
        self.load_config = vllm_config.load_config
        self.decoding_config = vllm_config.decoding_config or DecodingConfig(  # noqa
        )
-        self.prompt_adapter_config = vllm_config.prompt_adapter_config  # noqa
        self.observability_config = vllm_config.observability_config or ObservabilityConfig(  # noqa
        )

@@ -238,14 +236,14 @@ class LLMEngine:
        self.log_stats = log_stats
        self.use_cached_outputs = use_cached_outputs

-        if not self.model_config.skip_tokenizer_init:
-            self.tokenizer = self._init_tokenizer()
-            self.detokenizer = Detokenizer(self.tokenizer)
-            tokenizer_group = self.get_tokenizer_group()
-        else:
+        if self.model_config.skip_tokenizer_init:
            self.tokenizer = None
            self.detokenizer = None
            tokenizer_group = None
+        else:
+            self.tokenizer = self._init_tokenizer()
+            self.detokenizer = Detokenizer(self.tokenizer)
+            tokenizer_group = self.get_tokenizer_group()

        # Ensure that the function doesn't contain a reference to self,
        # to avoid engine GC issues
@@ -294,8 +292,6 @@ class LLMEngine:
                    # Feature flags
                    "enable_lora":
                    bool(self.lora_config),
-                    "enable_prompt_adapter":
-                    bool(self.prompt_adapter_config),
                    "enable_prefix_caching":
                    self.cache_config.enable_prefix_caching,
                    "enforce_eager":
@@ -542,9 +538,6 @@ class LLMEngine:
            self.lora_config.verify_with_model_config(self.model_config)
            self.lora_config.verify_with_scheduler_config(
                self.scheduler_config)
-        if self.prompt_adapter_config:
-            self.prompt_adapter_config.verify_with_model_config(
-                self.model_config)

    def _add_processed_request(
        self,
@@ -553,7 +546,6 @@ class LLMEngine:
        params: Union[SamplingParams, PoolingParams],
        arrival_time: float,
        lora_request: Optional[LoRARequest],
-        prompt_adapter_request: Optional[PromptAdapterRequest],
        trace_headers: Optional[Mapping[str, str]] = None,
        priority: int = 0,
    ) -> Optional[SequenceGroup]:
@@ -569,7 +561,6 @@ class LLMEngine:
                arrival_time=arrival_time,
                lora_request=lora_request,
                trace_headers=trace_headers,
-                prompt_adapter_request=prompt_adapter_request,
                priority=priority,
            )
            return None
@@ -583,11 +574,10 @@ class LLMEngine:
        encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)

        seq = Sequence(seq_id, decoder_inputs, block_size, eos_token_id,
-                       lora_request, prompt_adapter_request)
+                       lora_request)

        encoder_seq = (None if encoder_inputs is None else Sequence(
-            seq_id, encoder_inputs, block_size, eos_token_id, lora_request,
-            prompt_adapter_request))
+            seq_id, encoder_inputs, block_size, eos_token_id, lora_request))

        # Create a SequenceGroup based on SamplingParams or PoolingParams
        if isinstance(params, SamplingParams):
@@ -598,7 +588,6 @@ class LLMEngine:
                arrival_time=arrival_time,
                lora_request=lora_request,
                trace_headers=trace_headers,
-                prompt_adapter_request=prompt_adapter_request,
                encoder_seq=encoder_seq,
                priority=priority)
        elif isinstance(params, PoolingParams):
@@ -608,7 +597,6 @@ class LLMEngine:
                params,
                arrival_time=arrival_time,
                lora_request=lora_request,
-                prompt_adapter_request=prompt_adapter_request,
                encoder_seq=encoder_seq,
                priority=priority)
        else:
@@ -637,7 +625,6 @@ class LLMEngine:
        lora_request: Optional[LoRARequest] = None,
        tokenization_kwargs: Optional[dict[str, Any]] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
    ) -> None:
        """Add a request to the engine's request pool.
@@ -658,7 +645,6 @@ class LLMEngine:
                the current monotonic time.
            lora_request: The LoRA request to add.
            trace_headers: OpenTelemetry trace headers.
-            prompt_adapter_request: The prompt adapter request to add.
            priority: The priority of the request.
                Only applicable with priority scheduling.

@@ -719,7 +705,6 @@ class LLMEngine:
            prompt,
            tokenization_kwargs=tokenization_kwargs,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
        )

        self._add_processed_request(
@@ -728,7 +713,6 @@ class LLMEngine:
            params=params,
            arrival_time=arrival_time,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
            trace_headers=trace_headers,
            priority=priority,
        )
@@ -741,7 +725,6 @@ class LLMEngine:
        arrival_time: float,
        lora_request: Optional[LoRARequest],
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        encoder_seq: Optional[Sequence] = None,
        priority: int = 0,
    ) -> SequenceGroup:
@@ -769,17 +752,15 @@ class LLMEngine:
        if self.vllm_config.speculative_config is not None:
            draft_size = \
                self.vllm_config.speculative_config.num_speculative_tokens + 1
-        seq_group = SequenceGroup(
-            request_id=request_id,
-            seqs=[seq],
-            arrival_time=arrival_time,
-            sampling_params=sampling_params,
-            lora_request=lora_request,
-            trace_headers=trace_headers,
-            prompt_adapter_request=prompt_adapter_request,
-            encoder_seq=encoder_seq,
-            priority=priority,
-            draft_size=draft_size)
+        seq_group = SequenceGroup(request_id=request_id,
+                                  seqs=[seq],
+                                  arrival_time=arrival_time,
+                                  sampling_params=sampling_params,
+                                  lora_request=lora_request,
+                                  trace_headers=trace_headers,
+                                  encoder_seq=encoder_seq,
+                                  priority=priority,
+                                  draft_size=draft_size)

        return seq_group

@@ -790,7 +771,6 @@ class LLMEngine:
        pooling_params: PoolingParams,
        arrival_time: float,
        lora_request: Optional[LoRARequest],
-        prompt_adapter_request: Optional[PromptAdapterRequest],
        encoder_seq: Optional[Sequence] = None,
        priority: int = 0,
    ) -> SequenceGroup:
@@ -798,15 +778,13 @@ class LLMEngine:
        # Defensive copy of PoolingParams, which are used by the pooler
        pooling_params = pooling_params.clone()
        # Create the sequence group.
-        seq_group = SequenceGroup(
-            request_id=request_id,
-            seqs=[seq],
-            arrival_time=arrival_time,
-            lora_request=lora_request,
-            pooling_params=pooling_params,
-            prompt_adapter_request=prompt_adapter_request,
-            encoder_seq=encoder_seq,
-            priority=priority)
+        seq_group = SequenceGroup(request_id=request_id,
+                                  seqs=[seq],
+                                  arrival_time=arrival_time,
+                                  lora_request=lora_request,
+                                  pooling_params=pooling_params,
+                                  encoder_seq=encoder_seq,
+                                  priority=priority)
        return seq_group

    def abort_request(self, request_id: Union[str, Iterable[str]]) -> None:
@@ -1834,16 +1812,6 @@ class LLMEngine:
    def pin_lora(self, lora_id: int) -> bool:
        return self.model_executor.pin_lora(lora_id)

-    def add_prompt_adapter(
-            self, prompt_adapter_request: PromptAdapterRequest) -> bool:
-        return self.model_executor.add_prompt_adapter(prompt_adapter_request)
-
-    def remove_prompt_adapter(self, prompt_adapter_id: int) -> bool:
-        return self.model_executor.remove_prompt_adapter(prompt_adapter_id)
-
-    def list_prompt_adapters(self) -> List[int]:
-        return self.model_executor.list_prompt_adapters()
-
    def start_profile(self) -> None:
        self.model_executor.start_profile()

@@ -10,7 +10,6 @@ from vllm import PoolingParams
 from vllm.inputs import PromptType
 from vllm.lora.request import LoRARequest
 from vllm.outputs import RequestOutput
-from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
 from vllm.utils import Device

@@ -33,7 +32,6 @@ class RPCProcessRequest:
    request_id: str
    lora_request: Optional[LoRARequest] = None
    trace_headers: Optional[Mapping[str, str]] = None
-    prompt_adapter_request: Optional[PromptAdapterRequest] = None
    priority: int = 0

    def __init__(
@@ -43,7 +41,6 @@ class RPCProcessRequest:
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
    ) -> None:
        super().__init__()
@@ -53,7 +50,6 @@ class RPCProcessRequest:
        self.request_id = request_id
        self.lora_request = lora_request
        self.trace_headers = trace_headers
-        self.prompt_adapter_request = prompt_adapter_request
        self.priority = priority


@@ -45,7 +45,6 @@ from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.outputs import PoolingRequestOutput, RequestOutput
-from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import SamplingParams
 from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
 from vllm.utils import Device
@@ -448,7 +447,6 @@ class MQLLMEngineClient(EngineClient):
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
    ) -> AsyncGenerator[RequestOutput, None]:
        """Generate outputs for a request.
@@ -465,8 +463,6 @@ class MQLLMEngineClient(EngineClient):
            request_id: The unique id of the request.
            lora_request: LoRA request to use for generation, if any.
            trace_headers: OpenTelemetry trace headers.
-            prompt_adapter_request: Prompt Adapter request to use
-                                            for generation, if any.
            priority: Priority of the request (lower means earlier handling).
                Any priority other than 0 will lead to an error if the
                scheduling policy is not "priority".
@@ -474,8 +470,7 @@ class MQLLMEngineClient(EngineClient):
        return cast(
            AsyncGenerator[RequestOutput, None],
            self._process_request(prompt, sampling_params, request_id,
-                                  lora_request, trace_headers,
-                                  prompt_adapter_request, priority))
+                                  lora_request, trace_headers, priority))

    def encode(
        self,
@@ -521,7 +516,6 @@ class MQLLMEngineClient(EngineClient):
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
    ) -> Union[AsyncGenerator[RequestOutput, None], AsyncGenerator[
            PoolingRequestOutput, None]]:
@@ -575,7 +569,6 @@ class MQLLMEngineClient(EngineClient):
                    request_id=request_id,
                    lora_request=lora_request,
                    trace_headers=trace_headers,
-                    prompt_adapter_request=prompt_adapter_request,
                    priority=priority,
                ))

@@ -304,14 +304,12 @@ class MQLLMEngine:
            self._send_outputs(rpc_err)

        try:
-            self.engine.add_request(
-                request_id=request_id,
-                prompt=request.prompt,
-                params=request.params,
-                lora_request=request.lora_request,
-                trace_headers=request.trace_headers,
-                prompt_adapter_request=request.prompt_adapter_request,
-                priority=request.priority)
+            self.engine.add_request(request_id=request_id,
+                                    prompt=request.prompt,
+                                    params=request.params,
+                                    lora_request=request.lora_request,
+                                    trace_headers=request.trace_headers,
+                                    priority=request.priority)

            if self.log_requests:
                logger.info("Added request %s.", request.request_id)
@@ -16,7 +16,6 @@ from vllm.lora.request import LoRARequest
 from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput
 from vllm.pooling_params import PoolingParams
-from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import BeamSearchParams, SamplingParams
 from vllm.transformers_utils.tokenizer import AnyTokenizer
 from vllm.utils import Device, collect_from_async_generator, random_uuid
@@ -55,7 +54,6 @@ class EngineClient(ABC):
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
    ) -> AsyncGenerator[RequestOutput, None]:
        """Generate outputs for a request."""
@@ -151,6 +151,27 @@ class CustomChatCompletionContentSimpleVideoParam(TypedDict, total=False):
    video_url: Required[str]


+class CustomThinkCompletionContentParam(TypedDict, total=False):
+    """A Think Completion Content Param that accepts a plain text and a boolean.
+
+    Example:
+    {
+        "thinking": "I am thinking about the answer",
+        "closed": True,
+        "type": "thinking"
+    }
+    """
+
+    thinking: Required[str]
+    """The thinking content."""
+
+    closed: bool
+    """Whether the thinking is closed."""
+
+    type: Required[Literal["thinking"]]
+    """The thinking type."""
+
+
 ChatCompletionContentPartParam: TypeAlias = Union[
    OpenAIChatCompletionContentPartParam, ChatCompletionContentPartAudioParam,
    ChatCompletionContentPartInputAudioParam,
@@ -159,7 +180,8 @@ ChatCompletionContentPartParam: TypeAlias = Union[
    CustomChatCompletionContentSimpleImageParam,
    ChatCompletionContentPartImageEmbedsParam,
    CustomChatCompletionContentSimpleAudioParam,
-    CustomChatCompletionContentSimpleVideoParam, str]
+    CustomChatCompletionContentSimpleVideoParam, str,
+    CustomThinkCompletionContentParam]


 class CustomChatCompletionMessageParam(TypedDict, total=False):
@@ -938,6 +960,7 @@ _ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)
 _InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)
 _RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)
 _PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)
+_ThinkParser = partial(cast, CustomThinkCompletionContentParam)
 # Need to validate url objects
 _ImageParser = TypeAdapter(ChatCompletionContentPartImageParam).validate_python
 _AudioParser = TypeAdapter(ChatCompletionContentPartAudioParam).validate_python
@@ -954,6 +977,8 @@ MM_PARSER_MAP: dict[
 ] = {
    "text":
    lambda part: _TextParser(part).get("text", None),
+    "thinking":
+    lambda part: _ThinkParser(part).get("thinking", None),
    "input_text":
    lambda part: _TextParser(part).get("text", None),
    "input_image":
@@ -1100,7 +1125,7 @@ def _parse_chat_message_content_part(
            "with empty / unparsable content.", part, part_type)
        return None

-    if part_type in ("text", "input_text", "refusal"):
+    if part_type in ("text", "input_text", "refusal", "thinking"):
        str_content = cast(str, content)
        if wrap_dicts:
            return {'type': 'text', 'text': str_content}
@@ -21,7 +21,7 @@ from vllm.entrypoints.utils import (VLLM_SUBCMD_PARSER_EPILOG,
 from vllm.executor.multiproc_worker_utils import _add_prefix
 from vllm.logger import init_logger
 from vllm.usage.usage_lib import UsageContext
-from vllm.utils import FlexibleArgumentParser, get_tcp_uri
+from vllm.utils import FlexibleArgumentParser, bind_process_name, get_tcp_uri
 from vllm.v1.engine.core import EngineCoreProc
 from vllm.v1.engine.utils import CoreEngineProcManager, launch_core_engines
 from vllm.v1.executor.abstract import Executor
@@ -45,11 +45,6 @@ class ServeSubcommand(CLISubcommand):
        if args.headless or args.api_server_count < 1:
            run_headless(args)
        else:
-            if args.data_parallel_start_rank:
-                raise ValueError(
-                    "data_parallel_start_rank is only applicable "
-                    "in headless mode. "
-                    "Add --headless flag to enable headless mode.")
            if args.api_server_count > 1:
                run_multi_api_server(args)
            else:
@@ -82,17 +77,18 @@ def run_headless(args: argparse.Namespace):

    if args.api_server_count > 1:
        raise ValueError("api_server_count can't be set in headless mode")
-
+    bind_process_name("APIServer_Headless")
    # Create the EngineConfig.
    engine_args = vllm.AsyncEngineArgs.from_cli_args(args)
    usage_context = UsageContext.OPENAI_API_SERVER
-    vllm_config = engine_args.create_engine_config(usage_context=usage_context)
+    vllm_config = engine_args.create_engine_config(usage_context=usage_context,
+                                                   headless=True)

    if not envs.VLLM_USE_V1:
        raise ValueError("Headless mode is only supported for V1")

-    if engine_args.data_parallel_rank is not None:
-        raise ValueError("data_parallel_rank is not applicable in "
+    if engine_args.data_parallel_hybrid_lb:
+        raise ValueError("data_parallel_hybrid_lb is not applicable in "
                         "headless mode")

    parallel_config = vllm_config.parallel_config
@@ -122,7 +118,7 @@ def run_headless(args: argparse.Namespace):
    engine_manager = CoreEngineProcManager(
        target_fn=EngineCoreProc.run_engine_core,
        local_engine_count=local_engine_count,
-        start_index=args.data_parallel_start_rank,
+        start_index=vllm_config.parallel_config.data_parallel_rank,
        local_start_index=0,
        vllm_config=vllm_config,
        local_client=False,
@@ -175,7 +171,8 @@ def run_multi_api_server(args: argparse.Namespace):
    parallel_config = vllm_config.parallel_config
    dp_rank = parallel_config.data_parallel_rank
    external_dp_lb = parallel_config.data_parallel_external_lb
-    assert external_dp_lb or dp_rank == 0
+    hybrid_dp_lb = parallel_config.data_parallel_hybrid_lb
+    assert external_dp_lb or hybrid_dp_lb or dp_rank == 0

    api_server_manager: Optional[APIServerProcessManager] = None

@@ -195,12 +192,12 @@ def run_multi_api_server(args: argparse.Namespace):
            stats_update_address=coordinator.get_stats_publish_address()
            if coordinator else None)

-        # For dp ranks > 0 in external DP LB mode, we must delay the
+        # For dp ranks > 0 in external/hybrid DP LB modes, we must delay the
        # start of the API servers until the local engine is started
        # (after the launcher context manager exits),
        # since we get the front-end stats update address from the coordinator
        # via the handshake with the local engine.
-        if dp_rank == 0 or not external_dp_lb:
+        if dp_rank == 0 or not (external_dp_lb or hybrid_dp_lb):
            # Start API servers using the manager.
            api_server_manager = APIServerProcessManager(
                **api_server_manager_kwargs)
@@ -45,7 +45,6 @@ from vllm.outputs import (ClassificationRequestOutput, EmbeddingRequestOutput,
                          PoolingRequestOutput, RequestOutput,
                          ScoringRequestOutput)
 from vllm.pooling_params import PoolingParams, PoolingTask
-from vllm.prompt_adapter.request import PromptAdapterRequest
 from vllm.sampling_params import (BeamSearchParams, GuidedDecodingParams,
                                  RequestOutputKind, SamplingParams)
 from vllm.transformers_utils.tokenizer import (AnyTokenizer, MistralTokenizer,
@@ -314,7 +313,6 @@ class LLM:
        *,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        guided_options_request: Optional[Union[LLMGuidedOptions,
                                               GuidedDecodingRequest]] = None,
    ) -> list[RequestOutput]:
@@ -330,7 +328,6 @@ class LLM:
        prompt_token_ids: Optional[list[int]] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        guided_options_request: Optional[Union[LLMGuidedOptions,
                                               GuidedDecodingRequest]] = None,
    ) -> list[RequestOutput]:
@@ -346,7 +343,6 @@ class LLM:
        prompt_token_ids: Optional[list[list[int]]] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        guided_options_request: Optional[Union[LLMGuidedOptions,
                                               GuidedDecodingRequest]] = None,
    ) -> list[RequestOutput]:
@@ -363,7 +359,6 @@ class LLM:
        prompt_token_ids: list[int],
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        guided_options_request: Optional[Union[LLMGuidedOptions,
                                               GuidedDecodingRequest]] = None,
    ) -> list[RequestOutput]:
@@ -380,7 +375,6 @@ class LLM:
        prompt_token_ids: list[list[int]],
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        guided_options_request: Optional[Union[LLMGuidedOptions,
                                               GuidedDecodingRequest]] = None,
    ) -> list[RequestOutput]:
@@ -395,7 +389,6 @@ class LLM:
        prompt_token_ids: Union[list[int], list[list[int]]],
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        guided_options_request: Optional[Union[LLMGuidedOptions,
                                               GuidedDecodingRequest]] = None,
    ) -> list[RequestOutput]:
@@ -415,7 +408,6 @@ class LLM:
        prompt_token_ids: Optional[Union[list[int], list[list[int]]]] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        guided_options_request: Optional[Union[LLMGuidedOptions,
                                               GuidedDecodingRequest]] = None,
        priority: Optional[list[int]] = None,
@@ -440,8 +432,6 @@ class LLM:
                it is used to create the progress bar.
                If `False`, no progress bar is created.
            lora_request: LoRA request to use for generation, if any.
-            prompt_adapter_request: Prompt Adapter request to use for
-                generation, if any.
            priority: The priority of the requests, if any.
                Only applicable when priority scheduling policy is enabled.

@@ -507,7 +497,6 @@ class LLM:
            params=sampling_params,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
            guided_options=guided_options_request,
            tokenization_kwargs=tokenization_kwargs,
            priority=priority,
@@ -963,7 +952,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        pooling_task: PoolingTask = "encode",
        tokenization_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[PoolingRequestOutput]:
@@ -980,7 +968,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        pooling_task: PoolingTask = "encode",
        tokenization_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[PoolingRequestOutput]:
@@ -997,7 +984,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        pooling_task: PoolingTask = "encode",
        tokenization_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[PoolingRequestOutput]:
@@ -1015,7 +1001,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        pooling_task: PoolingTask = "encode",
        tokenization_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[PoolingRequestOutput]:
@@ -1033,7 +1018,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        pooling_task: PoolingTask = "encode",
        tokenization_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[PoolingRequestOutput]:
@@ -1049,7 +1033,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        pooling_task: PoolingTask = "encode",
        tokenization_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[PoolingRequestOutput]:
@@ -1070,7 +1053,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        pooling_task: PoolingTask = "encode",
        tokenization_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[PoolingRequestOutput]:
@@ -1092,8 +1074,6 @@ class LLM:
                it is used to create the progress bar.
                If `False`, no progress bar is created.
            lora_request: LoRA request to use for generation, if any.
-            prompt_adapter_request: Prompt Adapter request to use for
-                generation, if any.
            pooling_task: Override the pooling task to use.

        Returns:
@@ -1150,7 +1130,6 @@ class LLM:
            use_tqdm=use_tqdm,
            lora_request=lora_request,
            tokenization_kwargs=tokenization_kwargs,
-            prompt_adapter_request=prompt_adapter_request,
        )

        outputs = self._run_engine(use_tqdm=use_tqdm)
@@ -1167,7 +1146,6 @@ class LLM:
        pooling_params: Optional[Union[PoolingParams,
                                       Sequence[PoolingParams]]] = None,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
    ) -> list[EmbeddingRequestOutput]:
        """
        Generate an embedding vector for each prompt.
@@ -1187,8 +1165,6 @@ class LLM:
                it is used to create the progress bar.
                If `False`, no progress bar is created.
            lora_request: LoRA request to use for generation, if any.
-            prompt_adapter_request: Prompt Adapter request to use for
-                generation, if any.

        Returns:
            A list of `EmbeddingRequestOutput` objects containing the
@@ -1205,7 +1181,6 @@ class LLM:
            use_tqdm=use_tqdm,
            pooling_params=pooling_params,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
            pooling_task="embed",
        )

@@ -1218,7 +1193,6 @@ class LLM:
        *,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
    ) -> list[ClassificationRequestOutput]:
        """
        Generate class logits for each prompt.
@@ -1236,8 +1210,6 @@ class LLM:
                it is used to create the progress bar.
                If `False`, no progress bar is created.
            lora_request: LoRA request to use for generation, if any.
-            prompt_adapter_request: Prompt Adapter request to use for
-                generation, if any.

        Returns:
            A list of `ClassificationRequestOutput` objects containing the
@@ -1253,7 +1225,6 @@ class LLM:
            prompts,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
            pooling_task="classify",
        )

@@ -1267,7 +1238,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
    ) -> list[ScoringRequestOutput]:

        encoded_output: list[PoolingRequestOutput] = self.encode(
@@ -1275,7 +1245,6 @@ class LLM:
            truncate_prompt_tokens=truncate_prompt_tokens,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
            pooling_task="embed",
        )

@@ -1303,7 +1272,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
    ) -> list[ScoringRequestOutput]:

        if isinstance(tokenizer, MistralTokenizer):
@@ -1361,7 +1329,6 @@ class LLM:
            params=pooling_params,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
-            prompt_adapter_request=prompt_adapter_request,
        )

        outputs = self._run_engine(use_tqdm=use_tqdm)
@@ -1381,7 +1348,6 @@ class LLM:
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
    ) -> list[ScoringRequestOutput]:
        """Generate similarity scores for all pairs `<text,text_pair>` or
          `<multi-modal data, multi-modal data pair>`.
@@ -1412,8 +1378,6 @@ class LLM:
                it is used to create the progress bar.
                If `False`, no progress bar is created.
            lora_request: LoRA request to use for generation, if any.
-            prompt_adapter_request: Prompt Adapter request to use for
-                generation, if any.

        Returns:
            A list of `ScoringRequestOutput` objects containing the
@@ -1504,8 +1468,7 @@ class LLM:
                data_2,  # type: ignore[arg-type]
                truncate_prompt_tokens,
                use_tqdm,
-                lora_request,
-                prompt_adapter_request)
+                lora_request)
        else:
            return self._embedding_score(
                tokenizer,
@@ -1513,8 +1476,7 @@ class LLM:
                data_2,  # type: ignore[arg-type]
                truncate_prompt_tokens,
                use_tqdm,
-                lora_request,
-                prompt_adapter_request)
+                lora_request)

    def start_profile(self) -> None:
        self.llm_engine.start_profile()
@@ -1625,7 +1587,6 @@ class LLM:
        *,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[Union[Sequence[LoRARequest], LoRARequest]],
-        prompt_adapter_request: Optional[PromptAdapterRequest],
        tokenization_kwargs: Optional[dict[str, Any]] = None,
        guided_options: Optional[GuidedDecodingRequest] = None,
        priority: Optional[list[int]] = None,
@@ -1671,7 +1632,6 @@ class LLM:
                tokenization_kwargs=tokenization_kwargs,
                lora_request=lora_request[i] if isinstance(
                    lora_request, Sequence) else lora_request,
-                prompt_adapter_request=prompt_adapter_request,
                priority=priority[i] if priority else 0,
            )

@@ -1681,7 +1641,6 @@ class LLM:
        params: Union[SamplingParams, PoolingParams],
        tokenization_kwargs: Optional[dict[str, Any]] = None,
        lora_request: Optional[LoRARequest] = None,
-        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
    ) -> None:
        request_id = str(next(self.request_counter))
@@ -1691,7 +1650,6 @@ class LLM:
            params,
            lora_request=lora_request,
            tokenization_kwargs=tokenization_kwargs,
-            prompt_adapter_request=prompt_adapter_request,
            priority=priority,
        )

--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Chengji YaoandGitHub	c0a8db461c	Revert "[TPU][Bugfix] fix OOM issue in CI test (#21550 )" This reverts commit `40d86ee412`.	2025-07-24 23:04:37 -07:00
Chengji YaoandGitHub	40d86ee412	[TPU][Bugfix] fix OOM issue in CI test (#21550 ) Signed-off-by: Chengji Yao <chengjiyao@google.com>	2025-07-24 23:01:53 -07:00
Yang ChenandGitHub	85d051f026	[Misc] Removed undefined cmake variables MOE_PERMUTE_ARCHS (#21262 ) Signed-off-by: Yang Chen <yangche@fb.com>	2025-07-24 22:54:23 -07:00
Ignacio SicaandGitHub	5140f54b89	[CI/Build] fix cpu_extension for apple silicon (#21195 ) Signed-off-by: ignaciosica <mignacio.sica@gmail.com>	2025-07-24 22:53:59 -07:00
Chengji Yao GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	947edd099e	[Misc][Tools] make max-model-len a parameter in auto_tune script (#21321 ) Signed-off-by: Chengji Yao <chengjiyao@google.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2025-07-24 22:46:43 -07:00
hfanandGitHub	fde60ee775	[Model] Fix a check for None but the return value was empty list in Gemma3 MM vision_embeddings (#21479 ) Signed-off-by: Hongmin Fan <fanhongmin@google.com>	2025-07-25 13:46:06 +08:00
Jason GuandGitHub	b38bc652ac	[Model] Support tensor parallel for timm ViT in Deepseek_vl2 (#21494 ) Signed-off-by: wzqd <1057337859@qq.com>	2025-07-24 22:45:16 -07:00
Ning XieandGitHub	adaf2c6d4f	[Bugfix] fix modelscope snapshot_download serialization (#21536 ) Signed-off-by: Andy Xie <andy.xning@gmail.com>	2025-07-24 22:44:38 -07:00
Li, JiangandGitHub	42343f1f89	[CI] Update CODEOWNERS for CPU and Intel GPU (#21582 ) Signed-off-by: jiang1.li <jiang1.li@intel.com>	2025-07-24 21:58:03 -07:00
Benji BeckandGitHub	965bc71b04	Integrate TensorSchema with shape validation for Phi3VImagePixelInputs (#21232 ) Signed-off-by: Benji Beck <benjibeck@meta.com>	2025-07-24 21:43:52 -07:00
Zhou FangandGitHub	807a328bb6	[Docs] Add `requirements/common.txt` to run unit tests (#21572 ) Signed-off-by: Zhou Fang <fang.github@gmail.com>	2025-07-24 20:51:15 -07:00
QiliangCuiandGitHub	e0be2c4d09	[TPU][Test] Temporarily suspend this MoE model in test_basic.py. (#21560 ) Signed-off-by: Qiliang Cui <derrhein@gmail.com>	2025-07-24 20:44:50 -07:00
Nick HillandGitHub	9c8b2c2a8a	[DP] Support api-server-count > 0 in hybrid DP LB mode (#21510 ) Signed-off-by: Nick Hill <nhill@redhat.com>	2025-07-24 20:18:16 -07:00
Varun Sundar Rabindranath GitHub Varun Sundar Rabindranath	2212cd6cfb	[Bugfix] DeepGemm utils : Fix hardcoded type-cast (#21517 ) Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>	2025-07-24 20:17:29 -07:00
Burkhard Ringlein GitHub Thomas Parnell	ce3a9b1378	[Kernel] adding fused_moe configs for upcoming granite4 (#21332 ) Signed-off-by: Burkhard Ringlein <ngl@zurich.ibm.com> Co-authored-by: Thomas Parnell <tpa@zurich.ibm.com>	2025-07-24 20:16:59 -07:00
Yuxuan ZhangandGitHub	2ce90e5b01	Fix GLM-4 PP Missing Layer When using with PP. (#21531 ) Signed-off-by: zRzRzRzRzRzRzR <2448370773@qq.com>	2025-07-24 20:07:38 -07:00
Wentao YeandGitHub	633f6e804b	[Bug] Fix DeepGemm Init Error (#21554 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2025-07-24 20:07:22 -07:00
Harry MellorandGitHub	b57296bb9a	[Docs] Fix `site_url` for RunLLM (#21564 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2025-07-24 20:05:58 -07:00
Cyrus LeungandGitHub	34ddcf9ff4	[Frontend] `run-batch` supports V1 (#21541 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>	2025-07-24 20:05:55 -07:00
Woosuk KwonandGitHub	fe56180c7f	[MoE] More balanced expert sharding (#21497 ) Signed-off-by: Woosuk Kwon <woosuk@thinkingmachines.ai>	2025-07-24 15:56:08 -07:00
QiliangCuiandGitHub	07d80d7b0e	[TPU][TEST] HF_HUB_DISABLE_XET=1 the test 3. (#21539 ) Signed-off-by: Qiliang Cui <derrhein@gmail.com>	2025-07-24 15:33:04 -07:00
weiliangandGitHub	2dd72d23d9	update flashinfer to v0.2.9rc1 (#21485 ) Signed-off-by: Weiliang Liu <weiliangl@nvidia.com>	2025-07-24 14:06:11 -07:00
Simon Mo GitHub Harry Mellor	a6c7fb8cff	[Docs] Add Expert Parallelism Initial Documentation (#21373 ) Signed-off-by: simon-mo <simon.mo@hey.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2025-07-24 12:36:06 -07:00
Ricardo DecalandGitHub	a7272c23d0	[Docs][minor] Fix broken gh-file link in distributed serving docs (#21543 ) Signed-off-by: Ricardo Decal <rdecal@anyscale.com>	2025-07-24 10:36:56 -07:00
Juncheng Gu GitHub Richard Liu Richard Liu	6066284914	[P/D] Support CPU Transfer in NixlConnector (#18293 ) Signed-off-by: Juncheng Gu <juncgu@gmail.com> Signed-off-by: Richard Liu <ricliu@google.com> Co-authored-by: Richard Liu <39319471+richardsliu@users.noreply.github.com> Co-authored-by: Richard Liu <ricliu@google.com>	2025-07-24 17:58:42 +01:00
Rui QiaoandGitHub	1e9ea8e69d	[P/D] Move FakeNixlWrapper to test dir (#21328 ) Signed-off-by: Rui Qiao <ruisearch42@gmail.com>	2025-07-24 08:53:45 -07:00
Chaojun ZhangandGitHub	d9f9a3fd96	[XPU] Conditionally import CUDA-specific passes to avoid import errors on xpu platform (#21036 ) Signed-off-by: chzhang <chaojun.zhang@intel.com>	2025-07-24 23:23:36 +08:00
Shu WangandGitHub	1b25f1fe75	Update flashinfer CUTLASS MoE Kernel (#21408 ) Signed-off-by: Shu Wang. <shuw@nvidia.com>	2025-07-24 08:13:31 -07:00
Wentao YeandGitHub	e8cb0d0495	[Bug] Fix Compressed Tensor NVFP4 `cutlass_fp4_group_mm` illegal memory access (#21465 ) Signed-off-by: yewentao256 <zhyanwentao@126.com>	2025-07-24 08:13:24 -07:00
Ricardo Decal GitHub Simon Mo Harry Mellor	684174115d	[Docs] Rewrite Distributed Inference and Serving guide (#20593 ) Signed-off-by: Ricardo Decal <rdecal@anyscale.com> Co-authored-by: Simon Mo <simon.mo@hey.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2025-07-24 08:13:05 -07:00
Sanger Steel GitHub William Goldby	cdb79ee63d	[Docs] Update Tensorizer usage documentation (#21190 ) Signed-off-by: Sanger Steel <sangersteel@gmail.com> Signed-off-by: William Goldby <willgoldby@gmail.com> Co-authored-by: William Goldby <willgoldby@gmail.com>	2025-07-24 06:56:18 -07:00
elvischenvandGitHub	5a19a6c670	[Fix] Update mamba_ssm to 2.2.5 (#21421 ) Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>	2025-07-24 03:25:41 -07:00
Ming YangandGitHub	2ded067fd2	[Bugfix] Fix CUDA arch flags for MoE permute (#21426 ) Signed-off-by: Ming Yang <minos.future@gmail.com>	2025-07-24 03:23:59 -07:00
Harry MellorandGitHub	13abd0eaf9	[Model] Officially support Emu3 with Transformers backend (#21319 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2025-07-24 03:22:12 -07:00
Lucas WilkinsonandGitHub	61b8cea3b4	[Attention] Optimize FlashInfer MetadataBuilder Build call (#21137 ) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>	2025-07-24 03:21:46 -07:00
cjackalandGitHub	526078a96c	bump `flashinfer` to `v0.2.8` (#21385 ) Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>	2025-07-24 03:20:38 -07:00
ChaunceyandGitHub	6da0078523	[Feat] Allow custom naming of vLLM processes (#21445 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>	2025-07-24 03:15:23 -07:00
Rui QiaoandGitHub	73e3949d07	[Misc] Improve comment for DPEngineCoreActor._set_cuda_visible_devices() (#21501 ) Signed-off-by: Rui Qiao <ruisearch42@gmail.com>	2025-07-24 03:13:40 -07:00
Shintarou Okada GitHub Harry Mellor	6eca337ce0	Replace `--expand-tools-even-if-tool-choice-none` with `--exclude-tools-when-tool-choice-none` for v0.10.0 (#20544 ) Signed-off-by: okada <kokuzen@gmail.com> Signed-off-by: okada shintarou <okada@preferred.jp> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2025-07-24 02:56:36 -07:00
Yuxuan ZhangandGitHub	85bda9e7d0	remove GLM-4.5 quantization wrong Code (#21435 )	2025-07-24 01:52:43 -07:00
22quinnandGitHub	610852a423	[Core] Support model loader plugins (#21067 ) Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>	2025-07-24 01:49:44 -07:00
Nick HillandGitHub	f0f4de8f26	[Misc] Fix duplicate FusedMoEConfig debug messages (#21455 ) Signed-off-by: Nick Hill <nhill@redhat.com>	2025-07-24 01:27:30 -07:00
Zhou FangandGitHub	fc5f756db4	[v1][Core] Clean up usages of `SpecializedManager` (#21407 ) Signed-off-by: Zhou Fang <fang.github@gmail.com>	2025-07-24 00:40:11 -07:00
Chengji Yao GitHub Simon Mo	e74bfc70e4	[TPU][Bugfix] fix moe layer (#21340 ) Signed-off-by: Chengji Yao <chengjiyao@google.com> Co-authored-by: Simon Mo <simon.mo@hey.com>	2025-07-24 00:38:39 -07:00
Gregory ShtrasbergandGitHub	90eeea8f85	[Bugfix][ROCm] Fix for warp_size uses on host (#21205 ) Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>	2025-07-24 00:37:19 -07:00
Harry MellorandGitHub	dde295a934	Deduplicate Transformers backend code using inheritance (#21461 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2025-07-24 00:16:23 -07:00
Julien DenizeandGitHub	6d8d0a24c0	Add think chunk (#21333 ) Signed-off-by: Julien Denize <julien.denize@mistral.ai>	2025-07-23 21:51:32 -07:00
Yinghai Lu GitHub Nick Hill Rui Qiao	11ef7a611e	[BugFix] Set CUDA_VISIBLE_DEVICES before spawning the subprocesses (#21211 ) Signed-off-by: Yinghai Lu <yinghai@thinkingmachines.ai> Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: Rui Qiao <ruisearch42@gmail.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Rui Qiao <ruisearch42@gmail.com>	2025-07-23 21:44:04 -07:00
Woosuk KwonandGitHub	dc2f159f8a	Dump input metadata on crash for async scheduling (#21258 ) Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>	2025-07-23 21:10:30 -07:00
Robert Shaw GitHub Robert Shaw Nick Hill Tyler Michael Smith	d5b981f8b1	[DP] Internal Load Balancing Per Node [`one-pod-per-node`] (#21238 ) Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>	2025-07-23 20:57:32 -07:00
Nick HillandGitHub	eec6942014	[BugFix] Fix KVConnector TP worker aggregation (#21473 ) Signed-off-by: Nick Hill <nhill@redhat.com>	2025-07-23 20:56:49 -07:00
KazusatoOoko GitHub KazusatoOko	fd48d99ffd	[BugFix]: Batch generation from prompt_embeds fails for long prompts (#21390 ) Signed-off-by: KazusatoOko <kazusto.oko@sakana.ai> Co-authored-by: KazusatoOko <kazusto.oko@sakana.ai>	2025-07-23 20:43:17 -07:00
WeiQing ChenandGitHub	f8c15c4efb	[Bugfix] Fix example disagg_example_p2p_nccl_xpyd.sh zombie process (#21437 ) Signed-off-by: David Chen <530634352@qq.com>	2025-07-23 20:42:11 -07:00
Matthew BonanniandGitHub	aa08a954f9	[Bugfix] Fix casing warning (#21468 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>	2025-07-23 20:41:23 -07:00
Liangliang MaandGitHub	13e4ee1dc3	[XPU][UT] increase intel xpu CI test scope (#21492 ) Signed-off-by: Ma, Liangliang <liangliang.ma@intel.com>	2025-07-23 20:24:04 -07:00
Ming Yang GitHub Cyrus Leung	772ce5af97	[Misc] Add dummy maverick test to CI (#21324 ) Signed-off-by: Ming Yang <minos.future@gmail.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>	2025-07-23 20:22:42 -07:00
deven-labovitchandGitHub	63d92abb7c	[Frontend] Set MAX_AUDIO_CLIP_FILESIZE_MB via env var instead of hardcoding (#21374 ) Signed-off-by: Deven Labovitch <deven@videa.ai>	2025-07-23 20:22:19 -07:00
Hardik Gupta GitHub Isotr0pygemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	11599b0e1f	feat(gguf_loader): accept HF repo paths & URLs for GGUF (#20793 ) Signed-off-by: Hardik <hardikgupta1999@gmail.com> Signed-off-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2025-07-23 20:21:02 -07:00
Michael GoinandGitHub	f3137cdd81	[Core] Freeze gc during cuda graph capture to speed up init (#21146 ) Signed-off-by: Codex <codex@openai.com> Signed-off-by: mgoin <mgoin64@gmail.com>	2025-07-23 17:20:14 -07:00
Michael GoinandGitHub	82ec66f514	[V0 Deprecation] Remove Prompt Adapters (#20588 ) Signed-off-by: mgoin <mgoin64@gmail.com>	2025-07-23 16:36:48 -07:00
Yong Hoon ShinandGitHub	78c13e30e1	[V1] Fix local chunked attention always disabled (#21419 ) Signed-off-by: Yong Hoon Shin <yhshin@meta.com>	2025-07-23 15:59:30 -07:00
22quinnandGitHub	5c9b807b34	[Core] Add `reload_weights` RPC method (#20096 ) Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>	2025-07-23 14:24:52 -07:00
QiliangCuiandGitHub	14bf19e39f	[TPU][TEST] Fix the downloading issue in TPU v1 test 11. (#21418 ) Signed-off-by: Qiliang Cui <derrhein@gmail.com>	2025-07-23 11:29:36 -07:00
Yong Hoon ShinandGitHub	4ac7713e32	Add test case for compiling multiple graphs (#21044 ) Signed-off-by: Yong Hoon Shin <yhshin@meta.com>	2025-07-23 11:00:47 -07:00
Christian PintoandGitHub	8560a5b258	[Core][Model] PrithviMAE Enablement on vLLM v1 engine (#20577 ) Signed-off-by: Christian Pinto <christian.pinto@ibm.com>	2025-07-23 11:00:23 -07:00