p

Signed-off-by: khluu <khluu000@gmail.com>

.buildkite/hardware_tests/amd.yaml

@@ -5,7 +5,6 @@ steps:
     depends_on: []
     device: amd_cpu
     no_plugin: true
-    soft_fail: true
     commands:
     - >
       docker build

.buildkite/intel_jobs/test-intel.yaml

@@ -56,7 +56,7 @@ steps:
         'cd tests &&
         pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
         pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
-        pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py &&
+        pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
         pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py &&
         pytest -v -s v1/structured_output &&
         pytest -v -s v1/test_serial_utils.py &&

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -35,23 +35,6 @@ export PYTHONPATH=".."
 # Helper Functions
 ###############################################################################
 
-wait_for_clean_gpus() {
-  local timeout=${1:-300}
-  local start=$SECONDS
-  echo "--- Waiting for clean GPU state (timeout: ${timeout}s)"
-  while true; do
-    if grep -q clean /opt/amdgpu/etc/gpu_state; then
-      echo "GPUs state is \"clean\""
-      return
-    fi
-    if (( SECONDS - start >= timeout )); then
-      echo "Error: GPUs did not reach clean state within ${timeout}s" >&2
-      exit 1
-    fi
-    sleep 3
-  done
-}
-
 cleanup_docker() {
   # Get Docker's root directory
   docker_root=$(docker info -f '{{.DockerRootDir}}')
@@ -365,19 +348,12 @@ apply_rocm_test_overrides() {
 ###############################################################################
 
 # --- GPU initialization ---
-echo "--- Confirming Clean Initial State"
-wait_for_clean_gpus
-
 echo "--- ROCm info"
 rocminfo
 
 # --- Docker housekeeping ---
 cleanup_docker
 
-echo "--- Resetting GPUs"
-echo "reset" > /opt/amdgpu/etc/gpu_state
-wait_for_clean_gpus
-
 # --- Pull test image ---
 echo "--- Pulling container"
 image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"

.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh

@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
   exit 1
 fi
 
-echo "--- DP+TP"
-vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
-server_pid=$!
-timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
-vllm bench serve \
-    --backend vllm \
-    --dataset-name random \
-    --model meta-llama/Llama-3.2-3B-Instruct \
-    --num-prompts 20 \
-    --result-dir ./test_results \
-    --result-filename dp_pp.json \
-    --save-result \
-    --endpoint /v1/completions
-kill -s SIGTERM $server_pid; wait $server_pid || true
-failed_req=$(jq '.failed' ./test_results/dp_pp.json)
-if [ "$failed_req" -ne 0 ]; then
-  echo "Some requests were failed!"
-  exit 1
-fi
+#echo "--- DP+TP"
+#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
+#server_pid=$!
+#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+#vllm bench serve \
+#    --backend vllm \
+#    --dataset-name random \
+#    --model meta-llama/Llama-3.2-3B-Instruct \
+#    --num-prompts 20 \
+#    --result-dir ./test_results \
+#    --result-filename dp_pp.json \
+#    --save-result \
+#    --endpoint /v1/completions
+#kill -s SIGTERM $server_pid; wait $server_pid || true
+#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
+#if [ "$failed_req" -ne 0 ]; then
+#  echo "Some requests were failed!"
+#  exit 1
+#fi

.buildkite/test_areas/distributed.yaml

@@ -224,20 +224,6 @@ steps:
   commands:
     - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
 
-- label: MessageQueue TCP Multi-Node (2 GPUs)
-  timeout_in_minutes: 10
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 1
-  num_nodes: 2
-  no_plugin: true
-  optional: true
-  source_file_dependencies:
-  - vllm/distributed/device_communicators/shm_broadcast.py
-  - vllm/distributed/parallel_state.py
-  - tests/distributed/test_mq_tcp_multinode.py
-  commands:
-    - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 1 $IMAGE_TAG "torchrun --nnodes 2 --nproc-per-node=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_mq_tcp_multinode.py" "torchrun --nnodes 2 --nproc-per-node=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_mq_tcp_multinode.py"
-
 - label: Distributed NixlConnector PD accuracy (4 GPUs)
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/tests"

.buildkite/test_areas/model_runner_v2.yaml

@@ -78,7 +78,6 @@ steps:
     - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
     - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
 
-# These require fix https://github.com/vllm-project/vllm/pull/36280
 - label: Model Runner V2 Pipeline Parallelism (4 GPUs)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/tests"

csrc/cache_kernels.cu

@@ -91,9 +91,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
 
   if (n == 0) return;
 
-  const int64_t* src_data = src_ptrs.data_ptr<int64_t>();
-  const int64_t* dst_data = dst_ptrs.data_ptr<int64_t>();
-  const int64_t* size_data = sizes.data_ptr<int64_t>();
+  int64_t* src_data = src_ptrs.mutable_data_ptr<int64_t>();
+  int64_t* dst_data = dst_ptrs.mutable_data_ptr<int64_t>();
+  int64_t* size_data = sizes.mutable_data_ptr<int64_t>();
 
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
@@ -107,15 +107,24 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
   CUmemcpyAttributes attr = {};
   attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
   size_t attrs_idx = 0;
+  #if defined(CUDA_VERSION) && CUDA_VERSION >= 13000
+  CUresult result = cuMemcpyBatchAsync(
+      reinterpret_cast<CUdeviceptr*>(dst_data),
+      reinterpret_cast<CUdeviceptr*>(src_data),
+      reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
+      &attrs_idx, 1, static_cast<CUstream>(stream));
+  TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed with error ",
+              result);
+  #else
   size_t fail_idx = 0;
   CUresult result = cuMemcpyBatchAsync(
-      reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(dst_data)),
-      reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(src_data)),
-      reinterpret_cast<size_t*>(const_cast<int64_t*>(size_data)),
-      static_cast<size_t>(n), &attr, &attrs_idx, 1, &fail_idx,
-      static_cast<CUstream>(stream));
+      reinterpret_cast<CUdeviceptr*>(dst_data),
+      reinterpret_cast<CUdeviceptr*>(src_data),
+      reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
+      &attrs_idx, 1, &fail_idx, static_cast<CUstream>(stream));
   TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
               fail_idx, " with error ", result);
+  #endif
 #else
   // Fallback for CUDA < 12.8 and ROCm: individual async copies.
   // cudaMemcpyDefault lets the driver infer direction from pointer types.

csrc/cpu/torch_bindings.cpp

@@ -8,8 +8,6 @@
 // libraries use different ISAs.
 #define TORCH_EXTENSION_NAME _C
 
-std::string init_cpu_threads_env(const std::string& cpu_ids);
-
 void release_dnnl_matmul_handler(int64_t handler);
 
 int64_t create_onednn_scaled_mm_handler(const torch::Tensor& b,
@@ -354,7 +352,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "str act, str isa) -> ()");
   ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
 #endif
-  ops.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);
   ops.def(
       "mla_decode_kvcache("
       "   Tensor! out, Tensor query, Tensor kv_cache,"

csrc/cpu/utils.cpp

@@ -21,150 +21,6 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
 
 #endif
 
-#ifndef VLLM_NUMA_DISABLED
-std::string init_cpu_threads_env(const std::string& cpu_ids) {
-  bitmask* omp_cpu_mask = numa_parse_cpustring_all(cpu_ids.c_str());
-  TORCH_CHECK(omp_cpu_mask != nullptr,
-              "Failed to parse CPU string: " + cpu_ids);
-  TORCH_CHECK(omp_cpu_mask->size > 0);
-  std::vector<int> omp_cpu_ids;
-  omp_cpu_ids.reserve(omp_cpu_mask->size);
-
-  constexpr int group_size = 8 * sizeof(*omp_cpu_mask->maskp);
-
-  for (int offset = 0; offset < omp_cpu_mask->size; offset += group_size) {
-    unsigned long group_mask = omp_cpu_mask->maskp[offset / group_size];
-    int i = 0;
-    while (group_mask) {
-      if (group_mask & 1) {
-        omp_cpu_ids.emplace_back(offset + i);
-      }
-      ++i;
-      group_mask >>= 1;
-    }
-  }
-
-  // Memory node binding
-  if (numa_available() != -1) {
-    std::set<int> node_ids;
-    for (const auto& cpu_id : omp_cpu_ids) {
-      int node_id = numa_node_of_cpu(cpu_id);
-      if (node_id != -1) {
-        node_ids.insert(node_id);
-      }
-    }
-    // Concatenate all node_ids into a single comma-separated string
-    if (!node_ids.empty()) {
-      std::string node_ids_str;
-      for (const int node_id : node_ids) {
-        if (!node_ids_str.empty()) {
-          node_ids_str += ",";
-        }
-        node_ids_str += std::to_string(node_id);
-      }
-
-      bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
-      bitmask* src_mask = numa_get_mems_allowed();
-
-      int pid = getpid();
-
-      if (mask && src_mask) {
-        // move all existing pages to the specified numa node.
-        *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
-        int page_num = numa_migrate_pages(pid, src_mask, mask);
-        if (page_num == -1) {
-          TORCH_WARN("numa_migrate_pages failed. errno: " +
-                     std::to_string(errno));
-        }
-
-        // Restrict memory allocation to the selected NUMA node(s).
-        // Enhances memory locality for the threads bound to those NUMA CPUs.
-        if (node_ids.size() > 1) {
-          errno = 0;
-          numa_set_interleave_mask(mask);
-          if (errno != 0) {
-            TORCH_WARN("numa_set_interleave_mask failed. errno: " +
-                       std::to_string(errno));
-          } else {
-            TORCH_WARN(
-                "NUMA binding: Using INTERLEAVE policy for memory "
-                "allocation across multiple NUMA nodes (nodes: " +
-                node_ids_str +
-                "). Memory allocations will be "
-                "interleaved across the specified NUMA nodes.");
-          }
-        } else {
-          errno = 0;
-          numa_set_membind(mask);
-          if (errno != 0) {
-            TORCH_WARN("numa_set_membind failed. errno: " +
-                       std::to_string(errno));
-          } else {
-            TORCH_WARN(
-                "NUMA binding: Using MEMBIND policy for memory "
-                "allocation on the NUMA nodes (" +
-                node_ids_str +
-                "). Memory allocations will be "
-                "strictly bound to these NUMA nodes.");
-          }
-        }
-
-        numa_set_strict(1);
-
-        numa_free_nodemask(mask);
-        numa_free_nodemask(src_mask);
-      } else {
-        TORCH_WARN(
-            "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
-            std::to_string(errno));
-      }
-    }
-  }
-
-  // OMP threads binding
-  omp_set_num_threads((int)omp_cpu_ids.size());
-  torch::set_num_threads((int)omp_cpu_ids.size());
-  TORCH_CHECK_EQ(omp_cpu_ids.size(), torch::get_num_threads());
-  TORCH_CHECK_EQ(omp_cpu_ids.size(), omp_get_max_threads());
-
-  std::vector<std::pair<int, int>> thread_core_mapping;
-  thread_core_mapping.reserve(omp_cpu_ids.size());
-  omp_lock_t writelock;
-  omp_init_lock(&writelock);
-
-  #pragma omp parallel for schedule(static, 1)
-  for (size_t i = 0; i < omp_cpu_ids.size(); ++i) {
-    cpu_set_t mask;
-    CPU_ZERO(&mask);
-    CPU_SET(omp_cpu_ids[i], &mask);
-    int ret = sched_setaffinity(0, sizeof(cpu_set_t), &mask);
-    if (ret == -1) {
-      TORCH_CHECK(false,
-                  "sched_setaffinity failed. errno: " + std::to_string(errno));
-    }
-
-    omp_set_lock(&writelock);
-    thread_core_mapping.emplace_back(gettid(), omp_cpu_ids[i]);
-    omp_unset_lock(&writelock);
-  }
-
-  omp_destroy_lock(&writelock);
-
-  numa_free_nodemask(omp_cpu_mask);
-
-  std::stringstream ss;
-  ss << "OMP threads binding of Process " << getpid() << ":\n";
-  std::sort(thread_core_mapping.begin(), thread_core_mapping.end(),
-            [](auto&& a, auto&& b) { return a.second < b.second; });
-  for (auto&& item : thread_core_mapping) {
-    ss << "\t"
-       << "OMP tid: " << item.first << ", core " << item.second << "\n";
-  }
-
-  return ss.str();
-}
-#endif  // VLLM_NUMA_DISABLED
-
 namespace cpu_utils {
 ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
   this->realloc(allocation_unit * 128);

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh

@@ -26,8 +26,10 @@ using namespace cute;
 template <class OutType, int ScaleGranularityM,
           int ScaleGranularityN, int ScaleGranularityK,
           class MmaTileShape, class ClusterShape,
-          class EpilogueScheduler, class MainloopScheduler>
+          class EpilogueScheduler, class MainloopScheduler,
+          bool swap_ab_ = false>
 struct cutlass_3x_gemm_fp8_blockwise {
+  static constexpr bool swap_ab = swap_ab_;
   using ElementAB = cutlass::float_e4m3_t;
 
   using ElementA = ElementAB;
@@ -55,9 +57,13 @@ struct cutlass_3x_gemm_fp8_blockwise {
   using ElementCompute = float;
   using ElementBlockScale = float; 
 
-  using ScaleConfig = cutlass::detail::Sm120BlockwiseScaleConfig<
+  using ScaleConfig = conditional_t<swap_ab,
+      cutlass::detail::Sm120BlockwiseScaleConfig<
         ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
-        cute::UMMA::Major::MN, cute::UMMA::Major::K>;
+        cute::UMMA::Major::K, cute::UMMA::Major::MN>,
+      cutlass::detail::Sm120BlockwiseScaleConfig<
+        ScaleGranularityM, ScaleGranularityN, ScaleGranularityK,
+        cute::UMMA::Major::MN, cute::UMMA::Major::K>>;
 
   // layout_SFA and layout_SFB cannot be swapped since they are deduced.
   using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
@@ -78,17 +84,32 @@ struct cutlass_3x_gemm_fp8_blockwise {
       ElementAccumulator,
       ElementCompute,
       ElementC,
-      LayoutC,
+      conditional_t<swap_ab, LayoutC_Transpose, LayoutC>,
       AlignmentC,
       ElementD,
-      LayoutD,
+      conditional_t<swap_ab, LayoutD_Transpose, LayoutD>,
       AlignmentD,
       EpilogueScheduler,
       DefaultOperation
   >::CollectiveOp;
  
   using StageCountType = cutlass::gemm::collective::StageCountAuto; 
-  using CollectiveMainloop = 
+  using CollectiveMainloop = conditional_t<swap_ab,
+      typename cutlass::gemm::collective::CollectiveBuilder<
+          ArchTag,
+          OperatorClass,
+          ElementB,
+          cute::tuple<LayoutB_Transpose, LayoutSFA>,
+          AlignmentB,
+          ElementA,
+          cute::tuple<LayoutA_Transpose, LayoutSFB>,
+          AlignmentA,
+          ElementAccumulator,
+          MmaTileShape,
+          ClusterShape,
+          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+          MainloopScheduler
+      >::CollectiveOp,
       typename cutlass::gemm::collective::CollectiveBuilder<
           ArchTag,
           OperatorClass,
@@ -103,7 +124,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
           ClusterShape,
           cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
           MainloopScheduler
-      >::CollectiveOp;
+      >::CollectiveOp>;
 
   // SM12x family to support both SM120 (RTX 5090) and SM121 (DGX Spark)
   using KernelType = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
@@ -115,7 +136,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
 // Tile configurations for different M ranges
 template <typename OutType>
 struct sm120_blockwise_fp8_config_default {
-  // M > 256: use 128x128x128 tile with Cooperative (Auto) schedule
+  // use 128x128x128 tile with Cooperative (Auto) schedule
   using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto;
   using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
   using TileShape = Shape<_128, _128, _128>;
@@ -127,8 +148,8 @@ struct sm120_blockwise_fp8_config_default {
 };
 
 template <typename OutType>
-struct sm120_blockwise_fp8_config_M64 {
-  // M in [1, 256]: use 64x128x128 tile with Pingpong schedule
+struct sm120_blockwise_fp8_config_pingpong {
+  // use 64x128x128 tile with Pingpong schedule
   using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwisePingpongSm120;
   using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
   using TileShape = Shape<_64, _128, _128>;
@@ -139,11 +160,24 @@ struct sm120_blockwise_fp8_config_M64 {
       EpilogueSchedule, KernelSchedule>;
 };
 
+template <typename OutType>
+struct sm120_blockwise_fp8_config_swapab {
+  // use 128x32x128 tile with Cooperative schedule
+  using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedBlockwiseCooperativeSm120;
+  using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto;
+  using TileShape = Shape<_128, _32, _128>;
+  using ClusterShape = Shape<_1, _1, _1>;
+  using Gemm = cutlass_3x_gemm_fp8_blockwise<
+      OutType, 128, 1, 128, TileShape, ClusterShape,
+      EpilogueSchedule, KernelSchedule, true>;
+};
+
 template <typename Gemm>
 void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Tensor const& a,
                                    torch::stable::Tensor const& b,
                                    torch::stable::Tensor const& a_scales,
                                    torch::stable::Tensor const& b_scales) {
+  static constexpr bool swap_ab = Gemm::swap_ab;
   using GemmKernel = typename Gemm::GemmKernel;
   using StrideA = typename Gemm::GemmKernel::StrideA;
   using StrideB = typename Gemm::GemmKernel::StrideB;
@@ -167,11 +201,13 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
   b_stride =
       cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1));
   c_stride =
-      cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(m, n, 1));
+      cutlass::make_cute_packed_stride(StrideC{}, swap_ab ? cute::make_shape(n, m, 1) : cute::make_shape(m, n, 1));
 
-  LayoutSFA layout_SFA = 
+  LayoutSFA layout_SFA = swap_ab ?
+      ScaleConfig::tile_atom_to_shape_SFA(make_shape(n, m, k, 1)) :
       ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
-  LayoutSFB layout_SFB = 
+  LayoutSFB layout_SFB = swap_ab ?
+      ScaleConfig::tile_atom_to_shape_SFB(make_shape(n, m, k, 1)) :
       ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
 
   auto a_ptr = static_cast<ElementAB const*>(a.data_ptr());
@@ -180,15 +216,24 @@ void cutlass_gemm_caller_blockwise(torch::stable::Tensor& out, torch::stable::Te
   auto b_scales_ptr = static_cast<ElementBlockScale const*>(b_scales.data_ptr());
 
   typename GemmKernel::MainloopArguments mainloop_args{};
-  mainloop_args.ptr_A = a_ptr;
-  mainloop_args.dA = a_stride;
-  mainloop_args.ptr_B = b_ptr;
-  mainloop_args.dB = b_stride;
-  mainloop_args.ptr_SFA = a_scales_ptr;
   mainloop_args.layout_SFA = layout_SFA;
-  mainloop_args.ptr_SFB = b_scales_ptr;
   mainloop_args.layout_SFB = layout_SFB;
-  auto prob_shape = cute::make_shape(m, n, k, 1);
+  if (swap_ab) {
+    mainloop_args.ptr_A = b_ptr;
+    mainloop_args.dA = b_stride;
+    mainloop_args.ptr_B = a_ptr;
+    mainloop_args.dB = a_stride;
+    mainloop_args.ptr_SFA = b_scales_ptr;
+    mainloop_args.ptr_SFB = a_scales_ptr;
+  } else {
+    mainloop_args.ptr_A = a_ptr;
+    mainloop_args.dA = a_stride;
+    mainloop_args.ptr_B = b_ptr;
+    mainloop_args.dB = b_stride;
+    mainloop_args.ptr_SFA = a_scales_ptr;
+    mainloop_args.ptr_SFB = b_scales_ptr;
+  }
+  auto prob_shape = swap_ab ? cute::make_shape(n, m, k, 1) : cute::make_shape(m, n, k, 1);
 
   auto c_ptr = static_cast<ElementD*>(out.data_ptr());
   typename GemmKernel::EpilogueArguments epilogue_args{
@@ -204,15 +249,26 @@ void cutlass_gemm_blockwise_sm120_fp8_dispatch(torch::stable::Tensor& out,
                                                torch::stable::Tensor const& a_scales,
                                                torch::stable::Tensor const& b_scales) {
   int M = a.size(0);
-  if (M <= 256) {
-    using Gemm = typename sm120_blockwise_fp8_config_M64<OutType>::Gemm;
+  // more heuristic tuning can be done here by checking N/K dimensions as well
+  bool swap_ab = (M <= 64) || (M % 4 != 0);
+
+  if (!swap_ab) {
+    if (M <= 256) {
+      using Gemm = typename sm120_blockwise_fp8_config_pingpong<OutType>::Gemm;
+      return cutlass_gemm_caller_blockwise<Gemm>(
+          out, a, b, a_scales, b_scales);
+    }
+    // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
+    using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
+    return cutlass_gemm_caller_blockwise<Gemm>(
+        out, a, b, a_scales, b_scales);
+  } else {
+    // Swap A/B for small M to improve performance
+    // Use TILE_N=32 as the minimum compatible tile size.
+    using Gemm = typename sm120_blockwise_fp8_config_swapab<OutType>::Gemm;
     return cutlass_gemm_caller_blockwise<Gemm>(
         out, a, b, a_scales, b_scales);
   }
-  // M > 256: use default 128x128x128 config with Cooperative (Auto) schedule
-  using Gemm = typename sm120_blockwise_fp8_config_default<OutType>::Gemm;
-  return cutlass_gemm_caller_blockwise<Gemm>(
-      out, a, b, a_scales, b_scales);
 }
 
 }  // namespace vllm

docker/Dockerfile

@@ -685,6 +685,9 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
         --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
     fi
 
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system "transformers==5.4.0"
+
 RUN --mount=type=cache,target=/root/.cache/uv \
 . /etc/environment && \
 uv pip list

docker/Dockerfile.rocm

@@ -390,7 +390,21 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
 RUN mkdir src && mv vllm src/vllm
 
 # This is a workaround to ensure pytest exits with the correct status code in CI tests.
-RUN echo "import os\n\ndef pytest_sessionfinish(session, exitstatus):\n    os._exit(int(exitstatus))" > /vllm-workspace/conftest.py
+RUN printf '%s\n' \
+    'import os' \
+    '' \
+    '_exit_code = 1' \
+    '' \
+    'def pytest_sessionfinish(session, exitstatus):' \
+    '    global _exit_code' \
+    '    _exit_code = int(exitstatus)' \
+    '' \
+    'def pytest_unconfigure(config):' \
+    '    import sys' \
+    '    sys.stdout.flush()' \
+    '    sys.stderr.flush()' \
+    '    os._exit(_exit_code)' \
+    > /vllm-workspace/conftest.py
 
 # -----------------------
 # Final vLLM image

requirements/rocm-test.in

@@ -1,3 +1,5 @@
+-r common.txt
+
 # testing
 pytest
 tensorizer==2.10.1

requirements/rocm-test.txt

@@ -15,6 +15,7 @@ aiohappyeyeballs==2.6.1
 aiohttp==3.13.3
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   aiohttp-cors
     #   fsspec
     #   gpt-oss
@@ -38,20 +39,31 @@ annotated-doc==0.0.4
     #   typer
 annotated-types==0.7.0
     # via pydantic
+anthropic==0.89.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 antlr4-python3-runtime==4.9.3
     # via
     #   hydra-core
     #   omegaconf
-anyio==4.6.2.post1
+anyio==4.13.0
     # via
+    #   anthropic
     #   httpx
+    #   mcp
+    #   openai
+    #   sse-starlette
     #   starlette
+    #   watchfiles
 arctic-inference==0.1.1
     # via -r requirements/rocm-test.in
 argcomplete==3.6.3
     # via datamodel-code-generator
 arrow==1.4.0
     # via isoduration
+astor==0.8.1
+    # via depyf
 attrs==26.1.0
     # via
     #   aiohttp
@@ -83,6 +95,8 @@ bitsandbytes==0.49.2
     #   lightning
 black==26.3.1
     # via datamodel-code-generator
+blake3==1.0.8
+    # via -r requirements/common.txt
 blobfile==3.0.0
     # via -r requirements/rocm-test.in
 bm25s==0.2.13
@@ -99,6 +113,10 @@ bounded-pool-executor==0.0.3
     # via pqdm
 buildkite-test-collector==0.1.9
     # via -r requirements/rocm-test.in
+cachetools==7.0.5
+    # via -r requirements/common.txt
+cbor2==5.9.0
+    # via -r requirements/common.txt
 certifi==2026.2.25
     # via
     #   fiona
@@ -132,6 +150,7 @@ click==8.3.1
     #   nltk
     #   rasterio
     #   ray
+    #   rich-toolkit
     #   schemathesis
     #   typer
     #   uvicorn
@@ -142,6 +161,8 @@ cligj==0.7.2
     # via
     #   fiona
     #   rasterio
+cloudpickle==3.1.2
+    # via -r requirements/common.txt
 colorama==0.4.6
     # via
     #   perceptron
@@ -151,6 +172,10 @@ colorful==0.5.8
     # via ray
 colorlog==6.10.1
     # via optuna
+compressed-tensors==0.14.0.1
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 contourpy==1.3.3
     # via matplotlib
 coverage==7.13.5
@@ -182,24 +207,42 @@ decorator==5.2.1
     # via librosa
 decord==0.6.0
     # via -r requirements/rocm-test.in
+depyf==0.20.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 diffusers==0.37.0
     # via terratorch
 dill==0.3.8
     # via
     #   datasets
+    #   depyf
     #   evaluate
     #   lm-eval
     #   multiprocess
+diskcache==5.6.3
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 distlib==0.4.0
     # via virtualenv
+distro==1.9.0
+    # via
+    #   anthropic
+    #   openai
+dnspython==2.8.0
+    # via email-validator
 docker==7.1.0
     # via gpt-oss
 docopt==0.6.2
     # via num2words
 docstring-parser==0.17.0
-    # via jsonargparse
+    # via
+    #   anthropic
+    #   jsonargparse
 einops==0.8.2
     # via
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   encodec
     #   terratorch
@@ -208,6 +251,10 @@ einops==0.8.2
     #   vocos
 einx==0.4.2
     # via vector-quantize-pytorch
+email-validator==2.3.0
+    # via
+    #   fastapi
+    #   pydantic
 encodec==0.1.1
     # via vocos
 et-xmlfile==2.0.0
@@ -217,7 +264,15 @@ evaluate==0.4.6
 fastapi==0.135.2
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   gpt-oss
+    #   model-hosting-container-standards
+fastapi-cli==0.0.24
+    # via fastapi
+fastapi-cloud-cli==0.15.1
+    # via fastapi-cli
+fastar==0.9.0
+    # via fastapi-cloud-cli
 fastparquet==2026.3.0
     # via genai-perf
 fastsafetensors==0.2.2
@@ -225,6 +280,7 @@ fastsafetensors==0.2.2
 filelock==3.25.2
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   blobfile
     #   datasets
     #   diffusers
@@ -264,6 +320,10 @@ genson==1.3.0
     # via datamodel-code-generator
 geopandas==1.1.3
     # via terratorch
+gguf==0.18.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 gitdb==4.0.12
     # via gitpython
 gitpython==3.1.46
@@ -290,7 +350,10 @@ google-crc32c==1.8.0
 google-resumable-media==2.8.0
     # via google-cloud-storage
 googleapis-common-protos==1.73.0
-    # via google-api-core
+    # via
+    #   google-api-core
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
 gpt-oss==0.0.8
     # via -r requirements/rocm-test.in
 graphql-core==3.2.8
@@ -302,6 +365,7 @@ grpcio==1.78.0
     #   -c requirements/rocm.txt
     #   -r requirements/rocm-test.in
     #   grpcio-reflection
+    #   opentelemetry-exporter-otlp-proto-grpc
     #   ray
     #   tensorboard
 grpcio-reflection==1.78.0
@@ -328,12 +392,22 @@ html2text==2025.4.15
     # via gpt-oss
 httpcore==1.0.9
     # via httpx
+httptools==0.7.1
+    # via uvicorn
 httpx==0.27.2
     # via
     #   -r requirements/rocm-test.in
+    #   anthropic
     #   diffusers
+    #   fastapi
+    #   fastapi-cloud-cli
+    #   mcp
+    #   model-hosting-container-standards
+    #   openai
     #   perceptron
     #   schemathesis
+httpx-sse==0.4.3
+    # via mcp
 huggingface-hub==0.36.2
     # via
     #   -r requirements/rocm-test.in
@@ -370,10 +444,13 @@ hypothesis-jsonschema==0.23.1
 idna==3.11
     # via
     #   anyio
+    #   email-validator
     #   httpx
     #   jsonschema
     #   requests
     #   yarl
+ijson==3.5.0
+    # via -r requirements/common.txt
 imagehash==4.3.2
     # via -r requirements/rocm-test.in
 imageio==2.37.3
@@ -390,6 +467,8 @@ iniconfig==2.3.0
     # via pytest
 instanttensor==0.1.6
     # via -r requirements/rocm-test.in
+interegular==0.3.3
+    # via lm-format-enforcer
 isodate==0.7.2
     # via azure-storage-blob
 isoduration==20.11.0
@@ -399,15 +478,21 @@ isort==8.0.1
 jinja2==3.1.6
     # via
     #   datamodel-code-generator
+    #   fastapi
     #   genai-perf
     #   lm-eval
     #   torch
+jiter==0.13.0
+    # via
+    #   anthropic
+    #   openai
 jiwer==4.0.0
     # via -r requirements/rocm-test.in
 jmespath==1.1.0
     # via
     #   boto3
     #   botocore
+    #   model-hosting-container-standards
 joblib==1.5.3
     # via
     #   librosa
@@ -426,6 +511,7 @@ jsonpointer==3.1.0
 jsonschema==4.26.0
     # via
     #   hypothesis-jsonschema
+    #   mcp
     #   mistral-common
     #   ray
     #   schemathesis
@@ -443,6 +529,10 @@ kornia==0.8.2
     # via torchgeo
 kornia-rs==0.1.10
     # via kornia
+lark==1.2.2
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 lazy-loader==0.4
     # via
     #   librosa
@@ -466,14 +556,24 @@ lightning-utilities==0.15.3
     #   lightning
     #   pytorch-lightning
     #   torchmetrics
+llguidance==1.3.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 llvmlite==0.44.0
     # via numba
 lm-eval==0.4.11
     # via -r requirements/rocm-test.in
+lm-format-enforcer==0.11.3
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 logistro==2.0.1
     # via
     #   choreographer
     #   kaleido
+loguru==0.7.3
+    # via compressed-tensors
 lxml==6.0.2
     # via
     #   blobfile
@@ -500,12 +600,19 @@ mbstrdecoder==1.1.4
     #   dataproperty
     #   pytablewriter
     #   typepy
+mcp==1.27.0
+    # via -r requirements/common.txt
 mdurl==0.1.2
     # via markdown-it-py
 mistral-common==1.10.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
+model-hosting-container-standards==0.1.14
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 more-itertools==10.8.0
     # via
     #   inflect
@@ -522,6 +629,8 @@ msgpack==1.1.2
     # via
     #   librosa
     #   ray
+msgspec==0.20.0
+    # via -r requirements/common.txt
 mteb==2.11.5
     # via -r requirements/rocm-test.in
 multidict==6.7.1
@@ -541,6 +650,8 @@ networkx==3.6.1
     # via
     #   scikit-image
     #   torch
+ninja==1.13.0
+    # via -r requirements/common.txt
 nltk==3.9.3
     # via rouge-score
 num2words==0.5.14
@@ -555,6 +666,7 @@ numkong==7.1.1
     # via albucore
 numpy==2.2.6
     # via
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   accelerate
     #   albucore
@@ -572,6 +684,7 @@ numpy==2.2.6
     #   fastparquet
     #   genai-perf
     #   geopandas
+    #   gguf
     #   h5py
     #   imagehash
     #   imageio
@@ -620,15 +733,21 @@ numpy==2.2.6
     #   tritonclient
     #   vocos
     #   xarray
+    #   xgrammar
 omegaconf==2.3.0
     # via
     #   hydra-core
     #   lightning
 open-clip-torch==2.32.0
     # via -r requirements/rocm-test.in
+openai==2.30.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 openai-harmony==0.0.8
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   gpt-oss
 opencensus==0.11.4
     # via ray
@@ -637,6 +756,7 @@ opencensus-context==0.1.3
 opencv-python-headless==4.13.0.92
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   albumentations
     #   mistral-common
@@ -645,26 +765,59 @@ openpyxl==3.1.5
 opentelemetry-api==1.40.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-exporter-prometheus
     #   opentelemetry-sdk
     #   opentelemetry-semantic-conventions
+opentelemetry-exporter-otlp==1.40.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
+opentelemetry-exporter-otlp-proto-common==1.40.0
+    # via
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
+opentelemetry-exporter-otlp-proto-grpc==1.40.0
+    # via opentelemetry-exporter-otlp
+opentelemetry-exporter-otlp-proto-http==1.40.0
+    # via opentelemetry-exporter-otlp
 opentelemetry-exporter-prometheus==0.61b0
     # via ray
 opentelemetry-proto==1.40.0
-    # via ray
+    # via
+    #   opentelemetry-exporter-otlp-proto-common
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
+    #   ray
 opentelemetry-sdk==1.40.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-exporter-prometheus
+    #   opentelemetry-semantic-conventions-ai
     #   ray
 opentelemetry-semantic-conventions==0.61b0
-    # via opentelemetry-sdk
+    # via
+    #   opentelemetry-sdk
+    #   opentelemetry-semantic-conventions-ai
+opentelemetry-semantic-conventions-ai==0.5.1
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 optuna==3.6.1
     # via genai-perf
 orjson==3.11.7
     # via
     #   genai-perf
     #   kaleido
+outlines-core==0.2.11
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 packaging==26.0
     # via
     #   -c requirements/rocm.txt
@@ -682,6 +835,7 @@ packaging==26.0
     #   lazy-loader
     #   lightning
     #   lightning-utilities
+    #   lm-format-enforcer
     #   matplotlib
     #   optuna
     #   peft
@@ -713,6 +867,8 @@ pandas==3.0.1
     #   tacoreader
     #   torchgeo
     #   xarray
+partial-json-parser==0.2.1.1.post7
+    # via -r requirements/common.txt
 pathspec==1.0.4
     # via black
 pathvalidate==3.3.1
@@ -727,6 +883,7 @@ perf-analyzer==0.1.0
     # via genai-perf
 pillow==12.1.1
     # via
+    #   -r requirements/common.txt
     #   diffusers
     #   genai-perf
     #   imagehash
@@ -768,8 +925,14 @@ pqdm==0.2.0
 prometheus-client==0.24.1
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   opentelemetry-exporter-prometheus
+    #   prometheus-fastapi-instrumentator
     #   ray
+prometheus-fastapi-instrumentator==7.1.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 propcache==0.4.1
     # via
     #   aiohttp
@@ -779,6 +942,7 @@ proto-plus==1.27.1
 protobuf==6.33.6
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   google-api-core
     #   googleapis-common-protos
     #   grpcio-reflection
@@ -791,11 +955,14 @@ protobuf==6.33.6
     #   wandb
 psutil==7.2.2
     # via
+    #   -r requirements/common.txt
     #   accelerate
     #   peft
     #   tensorizer
 py==1.11.0
     # via pytest-forked
+py-cpuinfo==9.0.0
+    # via -r requirements/common.txt
 py-spy==0.4.1
     # via ray
 pyarrow==23.0.1
@@ -808,6 +975,8 @@ pyasn1==0.6.3
     # via pyasn1-modules
 pyasn1-modules==0.4.2
     # via google-auth
+pybase64==1.4.3
+    # via -r requirements/common.txt
 pycocotools==2.0.11
     # via terratorch
 pycountry==26.2.16
@@ -819,26 +988,44 @@ pycryptodomex==3.23.0
 pydantic==2.12.5
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   albumentations
+    #   anthropic
+    #   compressed-tensors
     #   datamodel-code-generator
     #   fastapi
+    #   fastapi-cloud-cli
     #   gpt-oss
     #   lightly
+    #   lm-format-enforcer
+    #   mcp
     #   mistral-common
+    #   model-hosting-container-standards
     #   mteb
+    #   openai
     #   openai-harmony
     #   pydantic-extra-types
+    #   pydantic-settings
     #   ray
     #   wandb
+    #   xgrammar
 pydantic-core==2.41.5
     # via pydantic
 pydantic-extra-types==2.11.1
-    # via mistral-common
+    # via
+    #   fastapi
+    #   mistral-common
+pydantic-settings==2.13.1
+    # via
+    #   fastapi
+    #   mcp
 pygments==2.19.2
     # via rich
 pyjwt==2.12.1
-    # via msal
+    # via
+    #   mcp
+    #   msal
 pyogrio==0.12.1
     # via geopandas
 pyparsing==3.3.2
@@ -898,6 +1085,16 @@ python-dateutil==2.9.0.post0
     #   typepy
 python-discovery==1.2.0
     # via virtualenv
+python-dotenv==1.2.2
+    # via
+    #   pydantic-settings
+    #   uvicorn
+python-json-logger==4.1.0
+    # via -r requirements/common.txt
+python-multipart==0.0.22
+    # via
+    #   fastapi
+    #   mcp
 python-rapidjson==1.23
     # via tritonclient
 pytokens==0.4.1
@@ -914,14 +1111,17 @@ pywavelets==1.9.0
     # via imagehash
 pyyaml==6.0.3
     # via
+    #   -r requirements/common.txt
     #   accelerate
     #   albumentations
     #   datamodel-code-generator
     #   datasets
     #   genai-perf
+    #   gguf
     #   huggingface-hub
     #   jsonargparse
     #   lightning
+    #   lm-format-enforcer
     #   omegaconf
     #   optuna
     #   peft
@@ -931,8 +1131,13 @@ pyyaml==6.0.3
     #   schemathesis
     #   timm
     #   transformers
+    #   uvicorn
     #   vocos
     #   wandb
+pyzmq==27.1.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 rapidfuzz==3.12.1
     # via
     #   -r requirements/rocm-test.in
@@ -952,6 +1157,7 @@ referencing==0.37.0
     #   jsonschema-specifications
 regex==2026.2.28
     # via
+    #   -r requirements/common.txt
     #   diffusers
     #   nltk
     #   open-clip-torch
@@ -961,12 +1167,14 @@ regex==2026.2.28
 requests==2.32.5
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   azure-core
     #   buildkite-test-collector
     #   datasets
     #   diffusers
     #   docker
     #   evaluate
+    #   gguf
     #   google-api-core
     #   google-cloud-storage
     #   gpt-oss
@@ -976,6 +1184,7 @@ requests==2.32.5
     #   mistral-common
     #   msal
     #   mteb
+    #   opentelemetry-exporter-otlp-proto-http
     #   pooch
     #   ray
     #   responses
@@ -999,8 +1208,15 @@ rich==14.3.3
     #   lightning
     #   mteb
     #   perceptron
+    #   rich-toolkit
     #   terratorch
     #   typer
+rich-toolkit==0.19.7
+    # via
+    #   fastapi-cli
+    #   fastapi-cloud-cli
+rignore==0.7.6
+    # via fastapi-cloud-cli
 rioxarray==0.22.0
     # via terratorch
 rouge-score==0.1.2
@@ -1070,12 +1286,20 @@ sentence-transformers==5.3.0
     # via
     #   -r requirements/rocm-test.in
     #   mteb
+sentencepiece==0.2.1
+    # via -r requirements/common.txt
 sentry-sdk==2.55.0
-    # via wandb
+    # via
+    #   fastapi-cloud-cli
+    #   wandb
+setproctitle==1.3.7
+    # via -r requirements/common.txt
 setuptools==79.0.1
     # via
     #   -c requirements/common.txt
     #   -c requirements/rocm.txt
+    #   -r requirements/common.txt
+    #   model-hosting-container-standards
     #   pytablewriter
     #   tensorboard
     #   torch
@@ -1092,6 +1316,7 @@ simplejson==3.20.2
 six==1.17.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   junit-xml
     #   lightly
     #   opencensus
@@ -1104,8 +1329,9 @@ smmap==5.0.3
     # via gitdb
 sniffio==1.3.1
     # via
-    #   anyio
+    #   anthropic
     #   httpx
+    #   openai
 sortedcontainers==2.4.0
     # via hypothesis
 soundfile==0.13.1
@@ -1124,10 +1350,16 @@ sqlalchemy==2.0.48
     #   optuna
 sqlitedict==2.1.0
     # via lm-eval
+sse-starlette==3.3.4
+    # via mcp
 starlette==0.52.1
     # via
     #   fastapi
+    #   mcp
+    #   model-hosting-container-standards
+    #   prometheus-fastapi-instrumentator
     #   schemathesis
+    #   sse-starlette
     #   starlette-testclient
 starlette-testclient==0.4.1
     # via schemathesis
@@ -1137,6 +1369,8 @@ stringzilla==4.6.0
     # via albucore
 structlog==25.5.0
     # via gpt-oss
+supervisor==4.3.0
+    # via model-hosting-container-standards
 sympy==1.14.0
     # via
     #   einx
@@ -1180,6 +1414,7 @@ tifffile==2026.3.3
 tiktoken==0.12.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   gpt-oss
     #   lm-eval
     #   mistral-common
@@ -1194,6 +1429,7 @@ timm==1.0.17
 tokenizers==0.22.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   transformers
 tomli==2.4.0
@@ -1212,8 +1448,10 @@ torchmetrics==1.9.0
     #   torchgeo
 tqdm==4.67.3
     # via
+    #   -r requirements/common.txt
     #   datasets
     #   evaluate
+    #   gguf
     #   huggingface-hub
     #   lightly
     #   lightning
@@ -1221,6 +1459,7 @@ tqdm==4.67.3
     #   mteb
     #   nltk
     #   open-clip-torch
+    #   openai
     #   optuna
     #   peft
     #   pqdm
@@ -1233,11 +1472,14 @@ tqdm==4.67.3
 transformers==4.57.5
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
+    #   compressed-tensors
     #   genai-perf
     #   peft
     #   sentence-transformers
     #   transformers-stream-generator
+    #   xgrammar
 transformers-stream-generator==0.0.5
     # via -r requirements/rocm-test.in
 tritonclient==2.66.0
@@ -1251,6 +1493,8 @@ typepy==1.3.4
     #   tabledata
 typer==0.24.1
     # via
+    #   fastapi-cli
+    #   fastapi-cloud-cli
     #   fastsafetensors
     #   perceptron
 typeshed-client==2.9.0
@@ -1258,9 +1502,12 @@ typeshed-client==2.9.0
 typing-extensions==4.15.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   aiosignal
     #   albumentations
     #   alembic
+    #   anthropic
+    #   anyio
     #   azure-core
     #   azure-identity
     #   azure-storage-blob
@@ -1272,9 +1519,13 @@ typing-extensions==4.15.0
     #   lightning
     #   lightning-utilities
     #   lm-eval
+    #   mcp
     #   mistral-common
     #   mteb
+    #   openai
     #   opentelemetry-api
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-sdk
     #   opentelemetry-semantic-conventions
     #   pqdm
@@ -1283,6 +1534,7 @@ typing-extensions==4.15.0
     #   pydantic-extra-types
     #   pytorch-lightning
     #   referencing
+    #   rich-toolkit
     #   sentence-transformers
     #   sqlalchemy
     #   starlette
@@ -1292,10 +1544,13 @@ typing-extensions==4.15.0
     #   typeshed-client
     #   typing-inspection
     #   wandb
+    #   xgrammar
 typing-inspection==0.4.2
     # via
     #   fastapi
+    #   mcp
     #   pydantic
+    #   pydantic-settings
 tzdata==2025.3
     # via arrow
 uri-template==1.3.0
@@ -1311,7 +1566,14 @@ urllib3==2.6.3
     #   sentry-sdk
     #   tritonclient
 uvicorn==0.42.0
-    # via gpt-oss
+    # via
+    #   fastapi
+    #   fastapi-cli
+    #   fastapi-cloud-cli
+    #   gpt-oss
+    #   mcp
+uvloop==0.22.1
+    # via uvicorn
 vector-quantize-pytorch==1.28.0
     # via -r requirements/rocm-test.in
 virtualenv==21.2.0
@@ -1320,10 +1582,16 @@ vocos==0.1.0
     # via -r requirements/rocm-test.in
 wandb==0.25.1
     # via terratorch
+watchfiles==1.1.1
+    # via
+    #   -r requirements/common.txt
+    #   uvicorn
 wcwidth==0.6.0
     # via ftfy
 webcolors==25.10.0
     # via jsonschema
+websockets==16.0
+    # via uvicorn
 werkzeug==3.1.6
     # via
     #   schemathesis
@@ -1334,6 +1602,10 @@ wrapt==2.1.2
     # via smart-open
 xarray==2026.2.0
     # via rioxarray
+xgrammar==0.1.33
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 xxhash==3.6.0
     # via
     #   datasets

requirements/xpu.txt

@@ -15,4 +15,4 @@ torch==2.10.0+xpu
 torchaudio
 torchvision
 
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.4/vllm_xpu_kernels-0.1.4-cp38-abi3-manylinux_2_28_x86_64.whl
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.5/vllm_xpu_kernels-0.1.5-cp38-abi3-manylinux_2_28_x86_64.whl

setup.py

@@ -1060,8 +1060,6 @@ setup(
         ],  # Required for audio processing
         "video": [],  # Kept for backwards compatibility
         "flashinfer": [],  # Kept for backwards compatibility
-        # Optional deps for AMD FP4 quantization support
-        "petit-kernel": ["petit-kernel"],
         # Optional deps for Helion kernel development
         # NOTE: When updating helion version, also update CI files:
         #   - .buildkite/test_areas/kernels.yaml

tests/distributed/test_mq_tcp_multinode.py

@@ -1,119 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Multi-node integration test for MessageQueue TCP fallback.
-
-Verifies that when writer and readers span separate nodes (Docker containers
-with isolated /dev/shm), `create_from_process_group` correctly detects
-cross-node ranks via `in_the_same_node_as()` and falls back to ZMQ TCP
-transport — and that data actually arrives.
-"""
-
-import numpy as np
-import torch.distributed as dist
-
-from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
-from vllm.distributed.parallel_state import in_the_same_node_as
-
-
-def main():
-    dist.init_process_group(backend="gloo")
-
-    rank = dist.get_rank()
-    world_size = dist.get_world_size()
-    assert world_size >= 2, (
-        f"Need at least 2 ranks across nodes, got world_size={world_size}"
-    )
-
-    # Verify that in_the_same_node_as detects cross-node correctly
-    status = in_the_same_node_as(dist.group.WORLD, source_rank=0)
-    local_count = sum(status)
-    print(
-        f"[Rank {rank}] in_the_same_node_as(source=0): {status}  "
-        f"(local={local_count}/{world_size})"
-    )
-    # With 2 Docker containers (1 proc each), rank 0 and rank 1
-    # should be on different nodes.
-    assert local_count < world_size, (
-        f"Expected cross-node ranks but all {world_size} ranks appear local."
-    )
-
-    # Create MessageQueue
-    writer_rank = 0
-    mq = MessageQueue.create_from_process_group(
-        dist.group.WORLD,
-        max_chunk_bytes=1024 * 1024,  # 1 MiB
-        max_chunks=10,
-        writer_rank=writer_rank,
-    )
-
-    # Verify the transport path selection
-    if rank == writer_rank:
-        print(
-            f"[Rank {rank}] Writer: n_local_reader={mq.n_local_reader}, "
-            f"n_remote_reader={mq.n_remote_reader}"
-        )
-        assert mq.n_remote_reader > 0, (
-            "Writer should have at least 1 remote (TCP) reader in a multi-node setup."
-        )
-    else:
-        if status[rank]:
-            assert mq._is_local_reader, (
-                f"Rank {rank} is on the same node as writer but is not a local reader."
-            )
-            print(f"[Rank {rank}] Reader: local (shared memory)")
-        else:
-            assert mq._is_remote_reader, (
-                f"Rank {rank} is on a different node but is not a remote (TCP) reader."
-            )
-            print(f"[Rank {rank}] Reader: remote (TCP)")
-
-    # Test data transfer: simple objects
-    dist.barrier()
-    if rank == writer_rank:
-        mq.enqueue("hello_from_node0")
-    else:
-        msg = mq.dequeue(timeout=10)
-        assert msg == "hello_from_node0"
-    dist.barrier()
-    print(f"[Rank {rank}] Simple object test passed")
-
-    # Test data transfer: numpy arrays
-    np.random.seed(42)
-    arrays = [
-        np.random.randint(0, 100, size=np.random.randint(100, 5000)) for _ in range(100)
-    ]
-
-    dist.barrier()
-    if rank == writer_rank:
-        for arr in arrays:
-            mq.enqueue(arr)
-    else:
-        for i, expected in enumerate(arrays):
-            received = mq.dequeue(timeout=10)
-            assert np.array_equal(expected, received), (
-                f"Array mismatch at index {i}: "
-                f"expected shape {expected.shape}, got shape {received.shape}"
-            )
-    dist.barrier()
-    print(f"[Rank {rank}] Numpy array test passed")
-
-    # Test data transfer: large payload (> max_chunk_bytes)
-    dist.barrier()
-    big_array = np.zeros(200_000, dtype=np.int64)  # ~1.6 MiB > 1 MiB chunk
-    if rank == writer_rank:
-        mq.enqueue(big_array)
-    else:
-        received = mq.dequeue(timeout=10)
-        assert np.array_equal(big_array, received)
-    dist.barrier()
-    print(f"[Rank {rank}] Large payload test passed")
-
-    # Done -- cleanup
-    dist.barrier()
-    print(f"[Rank {rank}] All MessageQueue TCP multi-node tests passed!")
-    dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    main()

tests/entrypoints/serve/disagg/test_generate_stream.py

@@ -0,0 +1,474 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from dataclasses import dataclass, field
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from vllm.config.multimodal import MultiModalConfig
+from vllm.entrypoints.openai.engine.protocol import StreamOptions
+from vllm.entrypoints.openai.models.protocol import BaseModelPath
+from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+from vllm.entrypoints.serve.disagg.protocol import GenerateRequest
+from vllm.entrypoints.serve.disagg.serving import ServingTokens
+from vllm.entrypoints.serve.render.serving import OpenAIServingRender
+from vllm.logprobs import Logprob
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.renderers import renderer_from_config
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine.async_llm import AsyncLLM
+
+MODEL_NAME = "openai-community/gpt2"
+BASE_MODEL_PATHS = [
+    BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
+]
+
+
+@dataclass
+class MockHFConfig:
+    model_type: str = "any"
+
+
+@dataclass
+class MockModelConfig:
+    task = "generate"
+    runner_type = "generate"
+    model = MODEL_NAME
+    tokenizer = MODEL_NAME
+    trust_remote_code = False
+    tokenizer_mode = "auto"
+    max_model_len = 100
+    tokenizer_revision = None
+    multimodal_config = MultiModalConfig()
+    hf_config = MockHFConfig()
+    hf_text_config = MockHFConfig()
+    logits_processors: list[str] | None = None
+    diff_sampling_param: dict | None = None
+    allowed_local_media_path: str = ""
+    allowed_media_domains: list[str] | None = None
+    encoder_config = None
+    generation_config: str = "auto"
+    media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
+    skip_tokenizer_init = False
+    is_encoder_decoder: bool = False
+    is_multimodal_model: bool = False
+    renderer_num_workers: int = 1
+
+    def get_diff_sampling_param(self):
+        return self.diff_sampling_param or {}
+
+
+@dataclass
+class MockParallelConfig:
+    _api_process_rank: int = 0
+
+
+@dataclass
+class MockVllmConfig:
+    model_config: MockModelConfig
+    parallel_config: MockParallelConfig
+
+
+def _build_renderer(model_config: MockModelConfig):
+    return renderer_from_config(
+        MockVllmConfig(model_config, parallel_config=MockParallelConfig()),
+    )
+
+
+def _build_serving_tokens(engine: AsyncLLM, **kwargs) -> ServingTokens:
+    models = OpenAIServingModels(
+        engine_client=engine,
+        base_model_paths=BASE_MODEL_PATHS,
+    )
+    serving_render = OpenAIServingRender(
+        model_config=engine.model_config,
+        renderer=engine.renderer,
+        io_processor=engine.io_processor,
+        model_registry=models.registry,
+        request_logger=None,
+        chat_template=None,
+        chat_template_content_format="auto",
+    )
+    serving = ServingTokens(
+        engine,
+        models,
+        openai_serving_render=serving_render,
+        request_logger=None,
+        **kwargs,
+    )
+
+    async def _fake_preprocess(*args, **kwargs):
+        return [{"prompt_token_ids": [1, 2, 3]}]
+
+    serving.openai_serving_render.preprocess_completion = AsyncMock(
+        side_effect=_fake_preprocess
+    )
+    return serving
+
+
+def _make_request_output(
+    request_id: str,
+    token_ids: list[int],
+    finish_reason: str | None = None,
+    finished: bool = False,
+    prompt_token_ids: list[int] | None = None,
+    logprobs: list[dict[int, Any] | None] | None = None,
+    num_cached_tokens: int | None = None,
+    index: int = 0,
+) -> RequestOutput:
+    return RequestOutput(
+        request_id=request_id,
+        prompt=None,
+        prompt_token_ids=prompt_token_ids or [1, 2, 3],
+        prompt_logprobs=None,
+        outputs=[
+            CompletionOutput(
+                index=index,
+                text="",
+                token_ids=token_ids,
+                cumulative_logprob=None,
+                logprobs=logprobs,
+                finish_reason=finish_reason,
+            )
+        ],
+        finished=finished,
+        metrics=None,
+        lora_request=None,
+        encoder_prompt=None,
+        encoder_prompt_token_ids=None,
+        num_cached_tokens=num_cached_tokens,
+    )
+
+
+def _mock_engine() -> MagicMock:
+    engine = MagicMock(spec=AsyncLLM)
+    engine.errored = False
+    engine.model_config = MockModelConfig()
+    engine.input_processor = MagicMock()
+    engine.io_processor = MagicMock()
+    engine.renderer = _build_renderer(engine.model_config)
+    return engine
+
+
+def _parse_sse_chunks(chunks: list[str]) -> list[Any]:
+    """Parse SSE chunks into dicts (JSON) or raw strings ([DONE])."""
+    parsed: list[Any] = []
+    for chunk in chunks:
+        assert chunk.startswith("data: ") and chunk.endswith("\n\n")
+        payload = chunk[len("data: ") : -len("\n\n")]
+        if payload == "[DONE]":
+            parsed.append("[DONE]")
+        else:
+            parsed.append(json.loads(payload))
+    return parsed
+
+
+@pytest.mark.asyncio
+async def test_stream_basic():
+    """Streaming returns SSE chunks with correct token_ids and ends with [DONE]."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10])
+        yield _make_request_output("req-1", token_ids=[20, 30])
+        yield _make_request_output(
+            "req-1", token_ids=[40], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10),
+        model=MODEL_NAME,
+        stream=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    parsed = _parse_sse_chunks(chunks)
+
+    # 3 data chunks + [DONE]
+    assert parsed[-1] == "[DONE]"
+    data_chunks = [c for c in parsed if c != "[DONE]"]
+    assert len(data_chunks) == 3
+
+    assert data_chunks[0]["choices"][0]["token_ids"] == [10]
+    assert data_chunks[1]["choices"][0]["token_ids"] == [20, 30]
+    assert data_chunks[2]["choices"][0]["token_ids"] == [40]
+    assert data_chunks[2]["choices"][0]["finish_reason"] == "stop"
+
+
+@pytest.mark.asyncio
+async def test_stream_error_mid_generation():
+    """finish_reason='error' mid-stream yields error chunk then [DONE]."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10])
+        yield _make_request_output(
+            "req-1", token_ids=[20], finish_reason="error", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10),
+        model=MODEL_NAME,
+        stream=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    assert len(chunks) >= 2
+    assert any("Internal server error" in chunk for chunk in chunks), (
+        f"Expected error message in chunks: {chunks}"
+    )
+    assert chunks[-1] == "data: [DONE]\n\n"
+
+
+@pytest.mark.asyncio
+async def test_stream_error_with_empty_delta():
+    """finish_reason='error' with empty delta_token_ids still raises."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10])
+        yield _make_request_output(
+            "req-1", token_ids=[], finish_reason="error", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10),
+        model=MODEL_NAME,
+        stream=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    assert any("Internal server error" in chunk for chunk in chunks), (
+        f"Expected error message in chunks: {chunks}"
+    )
+    assert chunks[-1] == "data: [DONE]\n\n"
+
+
+@pytest.mark.asyncio
+async def test_stream_skips_empty_token_output():
+    """Outputs with empty token_ids are skipped (no chunk emitted)."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10])
+        yield _make_request_output("req-1", token_ids=[])
+        yield _make_request_output(
+            "req-1", token_ids=[20], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10),
+        model=MODEL_NAME,
+        stream=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    parsed = _parse_sse_chunks(chunks)
+    assert parsed[-1] == "[DONE]"
+    data_chunks = [c for c in parsed if c != "[DONE]"]
+
+    # Only 2 data chunks — the empty one is skipped
+    assert len(data_chunks) == 2
+    assert data_chunks[0]["choices"][0]["token_ids"] == [10]
+    assert data_chunks[1]["choices"][0]["token_ids"] == [20]
+
+
+@pytest.mark.asyncio
+async def test_stream_include_usage():
+    """stream_options.include_usage emits a final usage-only chunk."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10])
+        yield _make_request_output(
+            "req-1", token_ids=[20], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10),
+        model=MODEL_NAME,
+        stream=True,
+        stream_options=StreamOptions(include_usage=True),
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    parsed = _parse_sse_chunks(chunks)
+    assert parsed[-1] == "[DONE]"
+
+    # The chunk before [DONE] should be the usage-only chunk
+    usage_chunk = parsed[-2]
+    assert usage_chunk["choices"] == []
+    assert usage_chunk["usage"]["prompt_tokens"] == 3
+    assert usage_chunk["usage"]["completion_tokens"] == 2
+    assert usage_chunk["usage"]["total_tokens"] == 5
+
+
+@pytest.mark.asyncio
+async def test_stream_continuous_usage():
+    """continuous_usage_stats adds usage to every data chunk."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output("req-1", token_ids=[10])
+        yield _make_request_output(
+            "req-1", token_ids=[20], finish_reason="stop", finished=True
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10),
+        model=MODEL_NAME,
+        stream=True,
+        stream_options=StreamOptions(
+            include_usage=True,
+            continuous_usage_stats=True,
+        ),
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    parsed = _parse_sse_chunks(chunks)
+    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+
+    # Every data chunk should have usage
+    for i, dc in enumerate(data_chunks):
+        assert dc["usage"] is not None, f"chunk {i} missing usage"
+        assert dc["usage"]["prompt_tokens"] == 3
+
+    # First chunk: 1 completion token
+    assert data_chunks[0]["usage"]["completion_tokens"] == 1
+    assert data_chunks[0]["usage"]["total_tokens"] == 4
+
+    # Second chunk: 2 completion tokens (cumulative)
+    assert data_chunks[1]["usage"]["completion_tokens"] == 2
+    assert data_chunks[1]["usage"]["total_tokens"] == 5
+
+
+@pytest.mark.asyncio
+async def test_stream_with_logprobs():
+    """Streaming with logprobs includes logprob data in each chunk."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1",
+            token_ids=[10],
+            logprobs=[{10: Logprob(logprob=-0.5)}],
+        )
+        yield _make_request_output(
+            "req-1",
+            token_ids=[20],
+            logprobs=[{20: Logprob(logprob=-1.0)}],
+            finish_reason="stop",
+            finished=True,
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10, logprobs=1),
+        model=MODEL_NAME,
+        stream=True,
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    parsed = _parse_sse_chunks(chunks)
+    data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
+
+    for dc in data_chunks:
+        lp = dc["choices"][0]["logprobs"]
+        assert lp is not None
+        assert len(lp["content"]) == 1
+        assert lp["content"][0]["token"].startswith("token_id:")
+
+
+@pytest.mark.asyncio
+async def test_stream_prompt_tokens_details():
+    """enable_prompt_tokens_details includes cached_tokens in final usage."""
+    engine = _mock_engine()
+
+    async def mock_generate(*args, **kwargs):
+        yield _make_request_output(
+            "req-1",
+            token_ids=[10],
+            finish_reason="stop",
+            finished=True,
+            num_cached_tokens=2,
+        )
+
+    engine.generate = MagicMock(side_effect=mock_generate)
+    serving = _build_serving_tokens(engine, enable_prompt_tokens_details=True)
+
+    request = GenerateRequest(
+        token_ids=[1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=10),
+        model=MODEL_NAME,
+        stream=True,
+        stream_options=StreamOptions(include_usage=True),
+    )
+
+    response = await serving.serve_tokens(request)
+    chunks = []
+    async for chunk in response:
+        chunks.append(chunk)
+
+    parsed = _parse_sse_chunks(chunks)
+    # Usage-only chunk (before [DONE])
+    usage_chunk = parsed[-2]
+    assert usage_chunk["choices"] == []
+    assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 2

tests/entrypoints/serve/disagg/test_serving_tokens.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 import os
 
 import httpx
@@ -113,6 +114,54 @@ async def test_generate_endpoint(client):
     assert "choices" in data
 
 
+@pytest.mark.asyncio
+async def test_generate_stream(client):
+    payload = {
+        "model": MODEL_NAME,
+        "token_ids": [1, 2, 3],
+        "sampling_params": {"max_tokens": 5},
+        "stream": True,
+    }
+    async with client.stream("POST", GEN_ENDPOINT, json=payload) as resp:
+        resp.raise_for_status()
+        chunks = []
+        async for line in resp.aiter_lines():
+            if not line.startswith("data: "):
+                continue
+            payload_str = line[len("data: ") :]
+            if payload_str == "[DONE]":
+                break
+            chunks.append(json.loads(payload_str))
+
+    assert len(chunks) > 0
+    # Every chunk has choices with token_ids
+    all_token_ids = []
+    for chunk in chunks:
+        assert "choices" in chunk
+        assert len(chunk["choices"]) == 1
+        choice = chunk["choices"][0]
+        assert "token_ids" in choice
+        assert len(choice["token_ids"]) > 0
+        all_token_ids.extend(choice["token_ids"])
+
+    # Last chunk should have a finish_reason
+    assert chunks[-1]["choices"][0]["finish_reason"] is not None
+
+    # Streaming should produce the same tokens as non-streaming
+    non_stream_resp = await client.post(
+        GEN_ENDPOINT,
+        json={
+            "model": MODEL_NAME,
+            "token_ids": [1, 2, 3],
+            "sampling_params": {"max_tokens": 5, "temperature": 0.0},
+            "stream": False,
+        },
+    )
+    non_stream_data = non_stream_resp.json()
+    # Just verify we got the right number of tokens
+    assert len(all_token_ids) == len(non_stream_data["choices"][0]["token_ids"])
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("logprobs_value", [0, 1, 5])
 async def test_generate_logprobs(client, logprobs_value):

tests/models/registry.py

@@ -7,6 +7,7 @@ from typing import Any, Literal
 
 import pytest
 from packaging.version import Version
+from transformers import PretrainedConfig
 from transformers import __version__ as TRANSFORMERS_VERSION
 
 from vllm.config.model import ModelDType, TokenizerMode
@@ -1004,7 +1005,26 @@ _MULTIMODAL_EXAMPLE_MODELS = {
         trust_remote_code=True,
     ),
     "NemotronH_Nano_VL_V2": _HfExamplesInfo(
-        "nano_vl_dummy", is_available_online=False, trust_remote_code=True
+        "nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16",
+        max_model_len=4096,
+        # NemotronH layers are constructed via `hybrid_override_pattern`:
+        use_original_num_layers=True,
+        hf_overrides={
+            "vision_config": PretrainedConfig(
+                args={
+                    "min_num_patches": 1,  # Trigger image dynamic res
+                    "max_num_patches": 12,
+                    "model": "vit_huge_patch16_224",
+                },
+                # Trigger conv3d:
+                video_temporal_patch_size=2,
+            ),
+            "text_config": {
+                "num_hidden_layers": 2,
+                "hybrid_override_pattern": "M*",
+            },
+        },
+        trust_remote_code=True,
     ),
     "OpenCUAForConditionalGeneration": _HfExamplesInfo(
         "xlangai/OpenCUA-7B", trust_remote_code=True

tests/models/utils.py

@@ -447,9 +447,16 @@ def dummy_hf_overrides(
     Dummy HF overrides function used to create dummy model
     with only minimum nums of layer.
     """
-    hf_config.update(exist_overrides or {})
+    # Copy because this helper is called more than once
+    # while loading config, and we `.pop()`
+    exist_overrides = (exist_overrides or {}).copy()
+    text_config_override = exist_overrides.pop("text_config", None)
+    hf_config.update(exist_overrides)
 
     text_config = hf_config.get_text_config()
+    if text_config_override is not None:
+        # multimodal test models may override *some* text-model fields
+        text_config.update(text_config_override)
 
     # Ensure at least 2 expert per group
     # Since `grouped_topk` assumes top-2

tests/quantization/test_online.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests online quantization."""
+
+import pytest
+import torch
+
+from tests.quantization.utils import (
+    _test_online_quant_peak_mem_impl,
+    is_quant_method_supported,
+)
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.online.fp8 import (
+    Fp8PerBlockOnlineLinearMethod,
+    Fp8PerBlockOnlineMoEMethod,
+    Fp8PerTensorOnlineLinearMethod,
+    Fp8PerTensorOnlineMoEMethod,
+)
+from vllm.platforms import current_platform
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("fp8"),
+    reason="FP8 is not supported on this GPU type.",
+)
+@pytest.mark.parametrize(
+    "quant_scheme,online_quant_args,expected_linear_cls,expected_moe_cls",
+    [
+        # simple case - quantization='fp8_per_tensor'
+        (
+            "fp8_per_tensor",
+            None,
+            Fp8PerTensorOnlineLinearMethod,
+            Fp8PerTensorOnlineMoEMethod,
+        ),
+        # simple case - quantization='fp8_per_block'
+        (
+            "fp8_per_block",
+            None,
+            Fp8PerBlockOnlineLinearMethod,
+            Fp8PerBlockOnlineMoEMethod,
+        ),
+        # quantization='online with linear_scheme_override and
+        # moe_scheme_override
+        (
+            "online",
+            {
+                "linear_scheme_override": "fp8_per_block",
+                "moe_scheme_override": "fp8_per_tensor",
+            },
+            Fp8PerBlockOnlineLinearMethod,
+            Fp8PerTensorOnlineMoEMethod,
+        ),
+        # ignore with direct layer name
+        (
+            "fp8_per_tensor",
+            # qkv_proj is fused from q_proj/k_proj/v_proj, so currently the
+            # ignore regex must match the unfused shard names
+            # TODO(future PR): also make 're:.*qkv_proj.*' work
+            {"ignore": ["model.layers.1.self_attn.o_proj", "re:.*[qkv]_proj"]},
+            Fp8PerTensorOnlineLinearMethod,
+            Fp8PerTensorOnlineMoEMethod,
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
+)
+def test_online_quantization(
+    vllm_runner,
+    quant_scheme: str,
+    online_quant_args: dict | None,
+    expected_linear_cls,
+    expected_moe_cls,
+    use_rocm_aiter: bool,
+    monkeypatch,
+) -> None:
+    """
+    Tests that online quantization frontend configuration works -
+    selecting quant schemes, overriding quant schemes by type, ignoring
+    layers.
+
+    Does not test performance, peak memory usage, etc.
+    """
+
+    if use_rocm_aiter:
+        monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+
+    # `LLM.apply_model` requires pickling a function.
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+
+    # a tiny model with both dense and MoE layers
+    model_name = "ibm-granite/granite-3.0-1b-a400m-base"
+
+    runner_kwargs = dict(
+        quantization=quant_scheme,
+        enforce_eager=True,
+    )
+    if online_quant_args is not None:
+        runner_kwargs["quantization_config"] = online_quant_args
+
+    with vllm_runner(
+        model_name,
+        **runner_kwargs,
+    ) as llm:
+
+        def check_model(model):
+            # checks further down in the test case are hardcoded for this
+            # model
+            assert model_name == "ibm-granite/granite-3.0-1b-a400m-base"
+
+            o_proj = model.model.layers[0].self_attn.o_proj
+            moe = model.model.layers[0].block_sparse_moe.experts
+
+            # o_proj and moe in layer 0 are always quantized (never ignored)
+            # because of how we craft the test case inputs
+            assert isinstance(o_proj.quant_method, expected_linear_cls)
+            if moe is not None:
+                assert isinstance(moe.quant_method, expected_moe_cls)
+
+            if current_platform.is_cuda():
+                assert o_proj.weight.dtype == torch.float8_e4m3fn
+            elif current_platform.is_rocm():
+                assert o_proj.weight.dtype == current_platform.fp8_dtype()
+            else:
+                pytest.skip("Only runs on CUDA and ROCm.")
+
+            # Verify ignored layers are unquantized.
+            if isinstance(online_quant_args, dict) and "ignore" in online_quant_args:
+                # only .*1.self_attn_o_proj is skipped
+                for layer_idx in range(len(model.model.layers)):
+                    o_proj = model.model.layers[layer_idx].self_attn.o_proj
+                    if layer_idx == 1:
+                        assert isinstance(o_proj.quant_method, UnquantizedLinearMethod)
+                    else:
+                        assert isinstance(o_proj.quant_method, expected_linear_cls)
+
+                # every .*self_attn.qkv_proj is skipped
+                for layer_idx in range(len(model.model.layers)):
+                    qkv_proj = model.model.layers[layer_idx].self_attn.qkv_proj
+                    assert isinstance(qkv_proj.quant_method, UnquantizedLinearMethod)
+
+        llm.apply_model(check_model)
+
+        outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
+        print(outputs[0][1])
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("fp8"),
+    reason="FP8 is not supported on this GPU type.",
+)
+def test_online_quant_peak_mem(
+    vllm_runner,
+    caplog_mp_spawn,
+    monkeypatch,
+) -> None:
+    _test_online_quant_peak_mem_impl(
+        "fp8_per_tensor", vllm_runner, caplog_mp_spawn, monkeypatch
+    )
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("fp8"),
+    reason="FP8 is not supported on this GPU type.",
+)
+def test_online_quant_load_format_dummy(
+    vllm_runner,
+    monkeypatch,
+    caplog,
+) -> None:
+    with vllm_runner(
+        "ibm-granite/granite-3.0-1b-a400m-base",
+        quantization="fp8_per_tensor",
+        enforce_eager=True,
+        load_format="dummy",
+    ) as llm:
+        outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
+        print(outputs[0][1])

tests/quantization/utils.py

@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import logging
+
+import regex as re
+
 from vllm.model_executor.layers.quantization import get_quantization_config
 from vllm.platforms import current_platform
 
@@ -21,3 +25,74 @@ def is_quant_method_supported(quant_method: str) -> bool:
     min_capability = get_quantization_config(quant_method).get_min_capability()
 
     return capability.to_int() >= min_capability
+
+
+def _test_online_quant_peak_mem_impl(
+    quantization_arg_value,
+    vllm_runner,
+    caplog_mp_spawn,
+    monkeypatch,
+) -> None:
+    # Note: `allenai/OLMoE-1B-7B-0125-Instruct` was selected because:
+    # 1. it covers both Linear and MoE paths
+    # 2. it is already used by other tests in CI, so adding it here
+    #    does not increase disk space for CI runners
+    # I really wanted to use `ibm-granite/granite-3.0-1b-a400m-base`
+    # which I think is the smallest MoE model in vLLM (2.5 GiB bf16,
+    # 1.3 GiB fp8), but could not as adding one more model makes CI
+    # run out of disk space.
+    model_name = "allenai/OLMoE-1B-7B-0125-Instruct"
+
+    # Force spawn to ensure caplog_mp_spawn works consistently
+    # (it relies on VLLM_LOGGING_CONFIG_PATH which spawn reads but fork ignores)
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+    with (
+        caplog_mp_spawn(logging.DEBUG) as log_holder,
+        vllm_runner(
+            model_name,
+            quantization=quantization_arg_value,
+            enforce_eager=True,
+        ) as llm,
+    ):
+        outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
+        print(outputs[0][1])
+
+    log_text = log_holder.text
+
+    # Parse memory usage from captured logs
+    model_memory_gib = None
+    peak_memory_gib = None
+    for line in log_text.splitlines():
+        if model_memory_gib is None:
+            match = re.search(r"Model loading took ([\d.]+) GiB memory", line)
+            if match:
+                model_memory_gib = float(match.group(1))
+        if peak_memory_gib is None:
+            match = re.search(
+                r"Peak GPU memory after loading weights: ([\d.]+) GiB", line
+            )
+            if match:
+                peak_memory_gib = float(match.group(1))
+
+    assert model_memory_gib is not None, "Could not find model loading memory log"
+    assert peak_memory_gib is not None, "Could not find peak memory log"
+    print(f"GPU memory used after loading weights: {model_memory_gib} GiB")
+    print(f"Peak GPU memory usage while loading weights: {peak_memory_gib} GiB")
+
+    # model specific, allenai/OLMoE-1B-7B-0125-Instruct fp8 online quant
+    # uses 6.65 GiB for weight loading (bf16 checkpoint is ~12.89 GiB)
+    expected_model_memory_gib = 6.7
+
+    # for allenai/OLMoE-1B-7B-0125-Instruct the number we see today is 9.06
+    # GiB, which is 1.36x above model_memory_gib. A slightly higher number is
+    # expected as when we load and quantize weights in a streaming fashion we
+    # need to have individual weights in bf16 + fp8 alive at the same time.
+    expected_peak_memory_gib = expected_model_memory_gib * 1.4
+
+    assert model_memory_gib < expected_model_memory_gib, (
+        f"{model_memory_gib=} higher than {expected_model_memory_gib}"
+    )
+    assert peak_memory_gib < expected_peak_memory_gib, (
+        f"{peak_memory_gib=} higher than {expected_peak_memory_gib}"
+    )

tests/v1/determinism/test_batch_invariance.py

@@ -187,7 +187,7 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
         tensor_parallel_size=tp_size,
         max_num_seqs=128,
         max_model_len=8192,
-        dtype="bfloat16",  # not everything is supported
+        dtype="auto",  # not everything is supported
         gpu_memory_utilization=0.9,
         enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config={"backend": backend},
@@ -400,7 +400,7 @@ def test_simple_generation(backend):
         tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
         gpu_memory_utilization=0.9,
         max_model_len=2048,
-        dtype="bfloat16",
+        dtype="auto",
         enable_prefix_caching=False,
         enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config={"backend": backend},
@@ -466,7 +466,7 @@ def test_logprobs_without_batch_invariance_should_fail(
         tensor_parallel_size=tp_size,
         max_num_seqs=32,
         max_model_len=8192,
-        dtype="bfloat16",
+        dtype="auto",
         enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config={"backend": backend},
     )
@@ -686,7 +686,7 @@ def test_decode_logprobs_match_prefill_logprobs(
         tensor_parallel_size=tp_size,
         max_num_seqs=32,
         max_model_len=8192,
-        dtype="bfloat16",
+        dtype="auto",
         enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config={"backend": backend},
     )
@@ -931,7 +931,7 @@ def LLM_with_max_seqs(
         max_num_seqs=max_num_seqs,
         gpu_memory_utilization=gpu_memory_utilization,
         max_model_len=max_model_len,
-        dtype="bfloat16",
+        dtype="auto",
         tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
         enable_prefix_caching=False,
         enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,

vllm/config/kernel.py

@@ -68,10 +68,13 @@ class IrOpPriorityConfig:
     def set_priority(self):
         """
         Context manager to set the IR op priority for all op members.
-        It also imports vllm.kernels to ensure all implementations are made available.
+        It also imports IR kernel implementations for the current platform
+        to ensure all implementations are made available.
         """
-        import vllm.kernels  # noqa: F401, registers IR op implementations
         from vllm.ir.op import IrOp
+        from vllm.platforms import current_platform
+
+        current_platform.import_ir_kernels()
 
         with contextlib.ExitStack() as stack:
             for field in fields(self):

vllm/config/model.py

@@ -21,6 +21,7 @@ from vllm.config.multimodal import (
     MultiModalConfig,
 )
 from vllm.config.pooler import PoolerConfig
+from vllm.config.quantization import OnlineQuantizationConfigArgs
 from vllm.config.scheduler import RunnerType
 from vllm.config.utils import config, getattr_iter
 from vllm.logger import init_logger
@@ -199,6 +200,10 @@ class ModelConfig:
     `quantization_config` attribute in the model config file. If that is
     `None`, we assume the model weights are not quantized and use `dtype` to
     determine the data type of the weights."""
+    quantization_config: dict[str, Any] | OnlineQuantizationConfigArgs | None = None
+    """Arguments for online quantization.
+    Auto-created when `quantization` equals to one of the string values of
+    the `OnlineQuantScheme` enum."""
     allow_deprecated_quantization: bool = False
     """Whether to allow deprecated quantization methods."""
     enforce_eager: bool = False
@@ -943,7 +948,6 @@ class ModelConfig:
                 "modelopt_fp4",
                 "modelopt_mxfp8",
                 "modelopt_mixed",
-                "petit_nvfp4",
                 # Ensure heavy backends are probed last to avoid unnecessary
                 # imports during override detection (e.g., MXFP4 imports Triton)
                 "mxfp4",

vllm/config/quantization.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from enum import Enum
+from typing import Any
+
+from pydantic import Field, field_validator
+
+from vllm.config.utils import config
+
+
+class OnlineQuantScheme(Enum):
+    """Supported online quantization schemes."""
+
+    # fp8, weights and activations scaled per-tensor
+    FP8_PER_TENSOR = "fp8_per_tensor"
+
+    # fp8, activations scaled in blocks of 1x128 elements, weights scaled in
+    # blocks of 128x128 elements (popularized by DeepSeek)
+    FP8_PER_BLOCK = "fp8_per_block"
+
+    # TODO(future PRs): add more online quant schemes here: mxfp8, etc
+
+
+@config
+class OnlineQuantizationConfigArgs:
+    """Configuration for online quantization.
+
+    Controls how ``OnlineQuantizationConfig`` is applied to a model.
+    At least one of ``global_scheme``, ``linear_scheme_override``, or
+    ``moe_scheme_override`` must be set.
+    """
+
+    global_scheme: OnlineQuantScheme | None = None
+    """Quantization scheme applied to every supported layer."""
+
+    linear_scheme_override: OnlineQuantScheme | None = None
+    """Quantization scheme override for ``LinearBase`` layers."""
+
+    moe_scheme_override: OnlineQuantScheme | None = None
+    """Quantization scheme override for ``FusedMoE`` layers."""
+
+    ignore: list[str] = Field(default_factory=list)
+    """Layers to skip quantization for. Supports exact names and regex
+    patterns with ``re:`` prefix (e.g. ``re:.*attn.*``), consistent with
+    compressed_tensors layer skipping."""
+
+    @field_validator(
+        "global_scheme", "linear_scheme_override", "moe_scheme_override", mode="before"
+    )
+    @classmethod
+    def _coerce_scheme(
+        cls, v: str | OnlineQuantScheme | None
+    ) -> OnlineQuantScheme | None:
+        if isinstance(v, str):
+            return OnlineQuantScheme(v)
+        return v
+
+
+def resolve_online_quant_config(
+    quantization: str | None,
+    quantization_config: dict[str, Any] | OnlineQuantizationConfigArgs | None,
+) -> OnlineQuantizationConfigArgs | None:
+    """Resolve online quant scheme shorthand into a quantization config.
+
+    If ``quantization`` is an online quant scheme (e.g. ``'fp8_per_tensor'``),
+    ensures ``quantization_config`` has a matching ``global_scheme`` and casts
+    it to :class:`OnlineQuantizationConfigArgs` if needed.
+    """
+    online_quant_values = {s.value for s in OnlineQuantScheme}
+    valid_quantization_values = online_quant_values | {"online"}
+    if quantization not in valid_quantization_values:
+        if quantization_config is not None:
+            raise ValueError(
+                f"quantization_config is only supported when quantization "
+                f"is one of {sorted(valid_quantization_values)}, "
+                f"got quantization={quantization!r}"
+            )
+        return None
+
+    if quantization in online_quant_values:
+        scheme = OnlineQuantScheme(quantization)
+
+        if quantization_config is None:
+            quantization_config = {
+                "global_scheme": scheme.value,
+            }
+        elif isinstance(quantization_config, OnlineQuantizationConfigArgs):
+            if quantization_config.global_scheme is None:
+                quantization_config.global_scheme = scheme
+            elif quantization_config.global_scheme != scheme:
+                raise ValueError(
+                    f"quantization={quantization!r} conflicts with "
+                    f"quantization_config.global_scheme="
+                    f"{quantization_config.global_scheme.value!r}. "
+                    f"These must match when both are specified."
+                )
+        elif isinstance(quantization_config, dict):
+            existing = quantization_config.get("global_scheme")
+            if existing is None:
+                quantization_config["global_scheme"] = scheme.value
+            else:
+                # Coerce to enum for comparison
+                existing_scheme = (
+                    OnlineQuantScheme(existing)
+                    if isinstance(existing, str)
+                    else existing
+                )
+                if existing_scheme != scheme:
+                    raise ValueError(
+                        f"quantization={quantization!r} conflicts "
+                        f"with quantization_config"
+                        f"['global_scheme']={existing!r}. "
+                        f"These must match when both are specified."
+                    )
+
+    # Cast dict to OnlineQuantizationConfigArgs
+    if isinstance(quantization_config, dict):
+        quantization_config = OnlineQuantizationConfigArgs(**quantization_config)
+
+    return quantization_config

vllm/config/vllm.py

@@ -1106,6 +1106,9 @@ class VllmConfig:
             )
         current_platform.check_and_update_config(self)
 
+        if envs.VLLM_USE_V2_MODEL_RUNNER:
+            self._validate_v2_model_runner()
+
         # Re-compute compile ranges after platform-specific config updates
         # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled)
         self._set_compile_ranges()
@@ -1713,6 +1716,7 @@ class VllmConfig:
             f"dcp_comm_backend={self.parallel_config.dcp_comm_backend}, "  # noqa
             f"disable_custom_all_reduce={self.parallel_config.disable_custom_all_reduce}, "  # noqa
             f"quantization={self.model_config.quantization}, "
+            f"quantization_config={self.model_config.quantization_config}, "  # noqa
             f"enforce_eager={self.model_config.enforce_eager}, "
             f"enable_return_routed_experts={self.model_config.enable_return_routed_experts}, "  # noqa
             f"kv_cache_dtype={self.cache_config.cache_dtype}, "
@@ -1728,6 +1732,49 @@ class VllmConfig:
             f"kernel_config={self.kernel_config!r}"
         )
 
+    def _validate_v2_model_runner(self) -> None:
+        """Check for features not yet supported by the V2 model runner."""
+        unsupported: list[str] = []
+
+        if self.model_config is not None and self.model_config.has_inner_state:
+            unsupported.append("hybrid/mamba models")
+
+        if self.parallel_config.prefill_context_parallel_size > 1:
+            unsupported.append("prefill context parallelism")
+
+        if (
+            self.speculative_config is not None
+            and self.speculative_config.method not in ("eagle", "eagle3", "mtp")
+        ):
+            unsupported.append(f"speculative method '{self.speculative_config.method}'")
+
+        if self.parallel_config.enable_dbo:
+            unsupported.append("dual batch overlap")
+
+        if (
+            self.model_config is not None
+            and self.model_config.enable_return_routed_experts
+        ):
+            # Will be added by https://github.com/vllm-project/vllm/pull/38163
+            unsupported.append("routed experts capture")
+
+        if self.model_config is not None and self.model_config.logits_processors:
+            unsupported.append("custom logits processors")
+
+        if self.cache_config.kv_sharing_fast_prefill:
+            # Will be added by https://github.com/vllm-project/vllm/pull/35045
+            unsupported.append("KV sharing fast prefill")
+
+        if self.ec_transfer_config is not None:
+            # Will be added by https://github.com/vllm-project/vllm/pull/38390
+            unsupported.append("EC transfer")
+
+        if unsupported:
+            raise ValueError(
+                "VLLM_USE_V2_MODEL_RUNNER does not yet support: "
+                + ", ".join(unsupported)
+            )
+
     def validate_block_size(self) -> None:
         """Validate block_size against DCP and mamba constraints.
 

vllm/engine/arg_utils.py

@@ -112,6 +112,7 @@ from vllm.v1.sample.logits_processor import LogitsProcessor
 from vllm.version import __version__ as VLLM_VERSION
 
 if TYPE_CHECKING:
+    from vllm.config.quantization import OnlineQuantizationConfigArgs
     from vllm.model_executor.layers.quantization import QuantizationMethods
     from vllm.model_executor.model_loader import LoadFormats
     from vllm.usage.usage_lib import UsageContext
@@ -483,6 +484,7 @@ class EngineArgs:
     hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides")
     tokenizer_revision: str | None = ModelConfig.tokenizer_revision
     quantization: QuantizationMethods | str | None = ModelConfig.quantization
+    quantization_config: "dict[str, Any] | OnlineQuantizationConfigArgs | None" = None
     allow_deprecated_quantization: bool = ModelConfig.allow_deprecated_quantization
     enforce_eager: bool = ModelConfig.enforce_eager
     disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce
@@ -661,6 +663,12 @@ class EngineArgs:
         if isinstance(self.ir_op_priority, dict):
             self.ir_op_priority = IrOpPriorityConfig(**self.ir_op_priority)
 
+        from vllm.config.quantization import resolve_online_quant_config
+
+        self.quantization_config = resolve_online_quant_config(
+            self.quantization, self.quantization_config
+        )
+
         # Setup plugins
         from vllm.plugins import load_general_plugins
 
@@ -1431,6 +1439,7 @@ class EngineArgs:
             tokenizer_revision=self.tokenizer_revision,
             max_model_len=self.max_model_len,
             quantization=self.quantization,
+            quantization_config=self.quantization_config,
             allow_deprecated_quantization=self.allow_deprecated_quantization,
             enforce_eager=self.enforce_eager,
             enable_return_routed_experts=self.enable_return_routed_experts,

vllm/entrypoints/llm.py

@@ -34,6 +34,9 @@ from vllm.config.model import (
     RunnerOption,
     TokenizerMode,
 )
+from vllm.config.quantization import (
+    OnlineQuantizationConfigArgs,
+)
 from vllm.distributed.weight_transfer.base import (
     WeightTransferInitRequest,
     WeightTransferUpdateRequest,
@@ -247,6 +250,9 @@ class LLM:
         attention_config: dict[str, Any] | AttentionConfig | None = None,
         kv_cache_memory_bytes: int | None = None,
         compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
+        quantization_config: dict[str, Any]
+        | OnlineQuantizationConfigArgs
+        | None = None,
         logits_processors: list[str | type[LogitsProcessor]] | None = None,
         **kwargs: Any,
     ) -> None:
@@ -367,6 +373,7 @@ class LLM:
             profiler_config=profiler_config_instance,
             attention_config=attention_config_instance,
             compilation_config=compilation_config_instance,
+            quantization_config=quantization_config,
             logits_processors=logits_processors,
             **kwargs,
         )

vllm/entrypoints/serve/disagg/protocol.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel, Field, field_validator
 
 from vllm.config import ModelConfig
 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionLogProbs
-from vllm.entrypoints.openai.engine.protocol import StreamOptions
+from vllm.entrypoints.openai.engine.protocol import StreamOptions, UsageInfo
 from vllm.logprobs import Logprob
 from vllm.renderers import TokenizeParams
 from vllm.sampling_params import SamplingParams
@@ -122,6 +122,26 @@ class GenerateResponseChoice(BaseModel):
     token_ids: list[int] | None = None
 
 
+class GenerateResponseStreamChoice(BaseModel):
+    index: int
+    logprobs: ChatCompletionLogProbs | None = None
+    finish_reason: str | None = None
+    token_ids: list[int] | None = None
+
+
+class GenerateStreamResponse(BaseModel):
+    request_id: str = Field(
+        default_factory=lambda: f"{random_uuid()}",
+        description=(
+            "The request_id related to this request. If the caller does "
+            "not set it, a random_uuid will be generated. This id is used "
+            "through out the inference process and return in response."
+        ),
+    )
+    choices: list[GenerateResponseStreamChoice]
+    usage: UsageInfo | None = Field(default=None)
+
+
 class GenerateResponse(BaseModel):
     request_id: str = Field(
         default_factory=lambda: f"{random_uuid()}",

vllm/entrypoints/serve/disagg/serving.py

@@ -18,6 +18,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
 )
 from vllm.entrypoints.openai.engine.protocol import (
     ErrorResponse,
+    GenerationError,
     PromptTokenUsageInfo,
     RequestResponseMetadata,
     UsageInfo,
@@ -28,12 +29,15 @@ from vllm.entrypoints.serve.disagg.protocol import (
     GenerateRequest,
     GenerateResponse,
     GenerateResponseChoice,
+    GenerateResponseStreamChoice,
+    GenerateStreamResponse,
 )
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
+from vllm.entrypoints.utils import should_include_usage
 from vllm.logger import init_logger
 from vllm.logprobs import Logprob
 from vllm.outputs import RequestOutput
-from vllm.sampling_params import SamplingParams
+from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.utils.collection_utils import as_list
 
 logger = init_logger(__name__)
@@ -74,7 +78,7 @@ class ServingTokens(OpenAIServing):
         self,
         request: GenerateRequest,
         raw_request: Request | None = None,
-    ) -> GenerateResponse | ErrorResponse:
+    ) -> GenerateResponse | ErrorResponse | AsyncGenerator[str, None]:
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
             logger.error("Error with model %s", error_check_ret)
@@ -110,6 +114,8 @@ class ServingTokens(OpenAIServing):
         sampling_params = request.sampling_params
         if self.force_no_detokenize:
             sampling_params.detokenize = False
+        if request.stream:
+            sampling_params.output_kind = RequestOutputKind.DELTA
 
         self._log_inputs(
             request_id,
@@ -133,9 +139,17 @@ class ServingTokens(OpenAIServing):
             priority=request.priority,
         )
 
-        # TODO(NickLucche): Implement streaming response
-
         assert result_generator is not None
+
+        if request.stream:
+            return self.serve_tokens_stream_generator(
+                request,
+                result_generator,
+                request_id,
+                model_name,
+                request_metadata,
+            )
+
         return await self.serve_tokens_full_generator(
             request, result_generator, request_id, model_name, request_metadata
         )
@@ -236,6 +250,109 @@ class ServingTokens(OpenAIServing):
 
         return response
 
+    async def serve_tokens_stream_generator(
+        self,
+        request: GenerateRequest,
+        result_generator: AsyncGenerator[RequestOutput, None],
+        request_id: str,
+        model_name: str,
+        request_metadata: RequestResponseMetadata,
+    ) -> AsyncGenerator[str, None]:
+        num_prompt_tokens = 0
+        num_generated_tokens: list[int] = []
+        first_iteration = True
+        num_cached_tokens = None
+        sampling_params: SamplingParams = request.sampling_params
+
+        include_usage, include_continuous_usage = should_include_usage(
+            request.stream_options, False
+        )
+
+        try:
+            async for res in result_generator:
+                if first_iteration:
+                    if res.prompt_token_ids is not None:
+                        num_prompt_tokens = len(res.prompt_token_ids)
+                    if res.encoder_prompt_token_ids is not None:
+                        num_prompt_tokens += len(res.encoder_prompt_token_ids)
+                    num_cached_tokens = res.num_cached_tokens
+                    num_generated_tokens = [0] * len(res.outputs)
+                    first_iteration = False
+
+                for output in res.outputs:
+                    i = output.index
+                    delta_token_ids = output.token_ids
+                    num_generated_tokens[i] += len(delta_token_ids)
+
+                    finish_reason = output.finish_reason
+                    self._raise_if_error(finish_reason, request_id)
+
+                    if not delta_token_ids:
+                        continue
+
+                    if sampling_params.logprobs is not None:
+                        out_logprobs = output.logprobs
+                        assert out_logprobs is not None, "Did not output logprobs"
+                        logprobs = self._create_tokens_logprobs(
+                            token_ids=delta_token_ids,
+                            top_logprobs=out_logprobs,
+                            num_output_top_logprobs=sampling_params.logprobs,
+                        )
+                    else:
+                        logprobs = None
+
+                    chunk = GenerateStreamResponse(
+                        request_id=request_id,
+                        choices=[
+                            GenerateResponseStreamChoice(
+                                index=i,
+                                logprobs=logprobs,
+                                finish_reason=finish_reason,
+                                token_ids=as_list(delta_token_ids),
+                            )
+                        ],
+                    )
+                    if include_continuous_usage:
+                        chunk.usage = UsageInfo(
+                            prompt_tokens=num_prompt_tokens,
+                            completion_tokens=num_generated_tokens[i],
+                            total_tokens=(num_prompt_tokens + num_generated_tokens[i]),
+                        )
+
+                    yield f"data: {chunk.model_dump_json()}\n\n"
+
+            total_completion_tokens = sum(num_generated_tokens)
+            final_usage_info = UsageInfo(
+                prompt_tokens=num_prompt_tokens,
+                completion_tokens=total_completion_tokens,
+                total_tokens=num_prompt_tokens + total_completion_tokens,
+            )
+
+            if self.enable_prompt_tokens_details and num_cached_tokens:
+                final_usage_info.prompt_tokens_details = PromptTokenUsageInfo(
+                    cached_tokens=num_cached_tokens
+                )
+
+            if include_usage:
+                final_chunk = GenerateStreamResponse(
+                    request_id=request_id,
+                    choices=[],
+                    usage=final_usage_info,
+                )
+                yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
+
+            request_metadata.final_usage_info = final_usage_info
+
+        except GenerationError as e:
+            yield (
+                f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n"
+            )
+        except Exception as e:
+            logger.exception("Error in token generation stream.")
+            data = self.create_streaming_error_response(e)
+            yield f"data: {data}\n\n"
+        yield "data: [DONE]\n\n"
+
     def _create_tokens_logprobs(
         self,
         token_ids: GenericSequence[int],

vllm/model_executor/layers/batch_invariant.py

@@ -10,6 +10,7 @@ import vllm.envs as envs
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.utils.mem_utils import get_max_shared_memory_bytes
 from vllm.utils.platform_utils import num_compute_units
 from vllm.utils.torch_utils import is_torch_equal_or_newer
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -177,7 +178,7 @@ def matmul_persistent(
         },
         torch.float16: {
             "BLOCK_SIZE_M": 128,
-            "BLOCK_SIZE_N": 256,
+            "BLOCK_SIZE_N": _fp16_block_size_n,
             "BLOCK_SIZE_K": 64,
             "GROUP_SIZE_M": 8,
             "num_stages": 3,
@@ -700,7 +701,7 @@ def bmm_batch_invariant(a, b, *, out=None):
         },
         torch.float16: {
             "BLOCK_SIZE_M": 128,
-            "BLOCK_SIZE_N": 256,
+            "BLOCK_SIZE_N": _fp16_block_size_n,
             "BLOCK_SIZE_K": 64,
             "num_stages": 3,
             "num_warps": 8,
@@ -752,7 +753,8 @@ def addmm_batch_invariant(bias, a, b):
 
 
 def _log_softmax_batch_invariant(input, dim, _half_to_float):
-    assert not _half_to_float, "not implemented"
+    if _half_to_float:
+        return log_softmax(input.float(), dim=dim)
     return log_softmax(input, dim=dim)
 
 
@@ -923,12 +925,15 @@ _original_fp16_reduction_precision = None
 _original_bf16_reduction_precision = None
 _original_cublas_workspace_cfg = None
 _original_cublaslt_workspace_size = None
+_fp16_block_size_n = 256
 
 
 def enable_batch_invariant_mode():
     global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
     global _original_fp16_reduction_precision, _original_bf16_reduction_precision
     global _original_cublas_workspace_cfg, _original_cublaslt_workspace_size
+    global _fp16_block_size_n
+
     if _batch_invariant_MODE:
         return
 
@@ -944,6 +949,10 @@ def enable_batch_invariant_mode():
         _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
         _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
         _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
+
+        # Query the shared memory size and set block size
+        # accordingly to avoid triton OutOfResources
+        _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
     else:
         # Only source of batch invariance for Hopper is split-k, can disable through
         # cuBLAS workspace config

vllm/model_executor/layers/fla/ops/chunk.py

@@ -16,7 +16,7 @@ from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
 from .cumsum import chunk_local_cumsum
 from .l2norm import l2norm_fwd
 from .solve_tril import solve_tril
-from .utils import SUPPRESS_LEVEL, input_guard
+from .utils import FLA_CHUNK_SIZE, SUPPRESS_LEVEL, input_guard
 from .wy_fast import recompute_w_u_fwd
 
 
@@ -30,13 +30,24 @@ def chunk_gated_delta_rule_fwd(
     initial_state: torch.Tensor,
     output_final_state: bool,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
 ):
-    g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)
+    g = chunk_local_cumsum(
+        g, chunk_size=FLA_CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
+    )
     # obtain WY representation. u is actually the new v.
     A = chunk_scaled_dot_kkt_fwd(
-        k=k, beta=beta, g=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32
+        k=k,
+        beta=beta,
+        g=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        output_dtype=torch.float32,
+    )
+    A = solve_tril(
+        A=A, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, output_dtype=k.dtype
     )
-    A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
     w, u = recompute_w_u_fwd(
         k=k,
         v=v,
@@ -44,6 +55,7 @@ def chunk_gated_delta_rule_fwd(
         A=A,
         g_cumsum=g,
         cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
     )
     h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
         k=k,
@@ -53,6 +65,8 @@ def chunk_gated_delta_rule_fwd(
         initial_state=initial_state,
         output_final_state=output_final_state,
         cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        chunk_offsets=chunk_offsets,
     )
     o = chunk_fwd_o(
         q=q,
@@ -62,6 +76,7 @@ def chunk_gated_delta_rule_fwd(
         g=g,
         scale=scale,
         cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
     )
     if SUPPRESS_LEVEL < 3:
         return g, o, A, final_state, None, None, None
@@ -84,6 +99,8 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
         initial_state: torch.Tensor,
         output_final_state: bool,
         cu_seqlens: torch.Tensor | None = None,
+        chunk_indices: torch.Tensor | None = None,
+        chunk_offsets: torch.Tensor | None = None,
         use_qk_l2norm_in_kernel: bool = False,
     ):
         if use_qk_l2norm_in_kernel:
@@ -100,6 +117,8 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
             initial_state=initial_state,
             output_final_state=output_final_state,
             cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            chunk_offsets=chunk_offsets,
         )
         ctx.scale = scale
         ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
@@ -117,6 +136,8 @@ def chunk_gated_delta_rule(
     initial_state: torch.Tensor = None,
     output_final_state: bool = False,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
     use_qk_l2norm_in_kernel: bool = False,
 ):
     r"""
@@ -206,6 +227,8 @@ def chunk_gated_delta_rule(
         initial_state,
         output_final_state,
         cu_seqlens,
+        chunk_indices,
+        chunk_offsets,
         use_qk_l2norm_in_kernel,
     )
     return o, final_state

vllm/model_executor/layers/fla/ops/chunk_delta_h.py

@@ -14,7 +14,7 @@ from vllm.triton_utils import tl, triton
 
 from .index import prepare_chunk_indices, prepare_chunk_offsets
 from .op import exp
-from .utils import use_cuda_graph
+from .utils import FLA_CHUNK_SIZE, use_cuda_graph
 
 NUM_WARPS = [2, 4, 8, 16]
 
@@ -286,9 +286,11 @@ def chunk_gated_delta_rule_fwd_h(
     gk: torch.Tensor | None = None,
     initial_state: torch.Tensor | None = None,
     output_final_state: bool = False,
-    chunk_size: int = 64,  # SY: remove this argument and force chunk size 64?
+    chunk_size: int = FLA_CHUNK_SIZE,
     save_new_value: bool = True,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     # This kernel is slightly different from fla to support Q/K with different head numbers.
     # In fla, Q/K always have the same head number, so Hg is always equal to H.
@@ -296,20 +298,15 @@ def chunk_gated_delta_rule_fwd_h(
     H = u.shape[-2]
     BT = chunk_size
 
-    chunk_indices = (
-        prepare_chunk_indices(cu_seqlens, chunk_size)
-        if cu_seqlens is not None
-        else None
-    )
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
     # N: the actual number of sequences in the batch with either equal or variable lengths
     if cu_seqlens is None:
         N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
     else:
-        N, NT, chunk_offsets = (
-            len(cu_seqlens) - 1,
-            len(chunk_indices),
-            prepare_chunk_offsets(cu_seqlens, BT),
-        )
+        N, NT = len(cu_seqlens) - 1, len(chunk_indices)
+        if chunk_offsets is None:
+            chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT)
     assert K <= 256, "current kernel does not support head dimension larger than 256."
 
     h = k.new_empty(B, NT, H, V, K)

vllm/model_executor/layers/fla/ops/chunk_o.py

@@ -146,14 +146,14 @@ def chunk_fwd_o(
     g: torch.Tensor | None = None,  # cumsum of log decay
     scale: float | None = None,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
     chunk_size: int = FLA_CHUNK_SIZE,
 ) -> torch.Tensor:
     B, T, Hg, K, V = *q.shape, v.shape[-1]
     H = v.shape[-2]
     BT = chunk_size
-    chunk_indices = (
-        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
-    )
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
     NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
     if scale is None:
         scale = k.shape[-1] ** -0.5

vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py

@@ -14,6 +14,7 @@ from vllm.triton_utils import tl, triton
 
 from .index import prepare_chunk_indices
 from .op import exp
+from .utils import FLA_CHUNK_SIZE
 
 
 @triton.heuristics(
@@ -103,7 +104,8 @@ def chunk_scaled_dot_kkt_fwd(
     g: torch.Tensor | None = None,
     beta: torch.Tensor | None = None,
     cu_seqlens: torch.Tensor | None = None,
-    chunk_size: int = 64,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = FLA_CHUNK_SIZE,
     output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
     r"""
@@ -119,6 +121,9 @@ def chunk_scaled_dot_kkt_fwd(
         cu_seqlens (torch.Tensor):
             The cumulative sequence lengths of the input tensor.
             Default: None
+        chunk_indices (torch.Tensor):
+            Pre-computed chunk indices. If None and cu_seqlens is provided,
+            computed internally. Default: None
         chunk_size (int):
             The chunk size. Default: 64.
         output_dtype (torch.dtype):
@@ -132,9 +137,8 @@ def chunk_scaled_dot_kkt_fwd(
     B, T, Hg, K = k.shape
     H = beta.shape[-1]
     BT = chunk_size
-    chunk_indices = (
-        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
-    )
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
     NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
 
     A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)

vllm/model_executor/layers/fla/ops/cumsum.py

@@ -162,6 +162,7 @@ def chunk_local_cumsum_scalar(
     chunk_size: int,
     reverse: bool = False,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
     head_first: bool = False,
     output_dtype: torch.dtype | None = torch.float,
 ) -> torch.Tensor:
@@ -172,10 +173,9 @@ def chunk_local_cumsum_scalar(
     assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
         "chunk_size must be a power of 2"
     )
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
     BT = chunk_size
-    chunk_indices = (
-        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
-    )
     NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
     g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
     grid = (NT, B * H)
@@ -199,6 +199,7 @@ def chunk_local_cumsum_vector(
     chunk_size: int,
     reverse: bool = False,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
     head_first: bool = False,
     output_dtype: torch.dtype | None = torch.float,
 ) -> torch.Tensor:
@@ -206,16 +207,13 @@ def chunk_local_cumsum_vector(
         B, H, T, S = g.shape
     else:
         B, T, H, S = g.shape
-    BT = chunk_size
-    chunk_indices = (
-        prepare_chunk_indices(cu_seqlens, chunk_size)
-        if cu_seqlens is not None
-        else None
-    )
-    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
     assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
         "chunk_size must be a power of 2"
     )
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    BT = chunk_size
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
 
     g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
 
@@ -247,6 +245,7 @@ def chunk_local_cumsum(
     chunk_size: int,
     reverse: bool = False,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
     head_first: bool = False,
     output_dtype: torch.dtype | None = torch.float,
     **kwargs,
@@ -257,11 +256,23 @@ def chunk_local_cumsum(
         )
     if len(g.shape) == 3:
         return chunk_local_cumsum_scalar(
-            g, chunk_size, reverse, cu_seqlens, head_first, output_dtype
+            g,
+            chunk_size,
+            reverse,
+            cu_seqlens,
+            chunk_indices,
+            head_first,
+            output_dtype,
         )
     elif len(g.shape) == 4:
         return chunk_local_cumsum_vector(
-            g, chunk_size, reverse, cu_seqlens, head_first, output_dtype
+            g,
+            chunk_size,
+            reverse,
+            cu_seqlens,
+            chunk_indices,
+            head_first,
+            output_dtype,
         )
     else:
         raise ValueError(

vllm/model_executor/layers/fla/ops/kda.py

@@ -23,7 +23,7 @@ from .index import prepare_chunk_indices
 from .l2norm import l2norm_fwd
 from .op import exp, log
 from .solve_tril import solve_tril
-from .utils import is_amd
+from .utils import FLA_CHUNK_SIZE, is_amd
 
 BT_LIST_AUTOTUNE = [32, 64, 128]
 NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]
@@ -721,7 +721,7 @@ def chunk_kda_scaled_dot_kkt_fwd(
     beta: torch.Tensor | None = None,
     scale: float | None = None,
     cu_seqlens: torch.Tensor | None = None,
-    chunk_size: int = 64,
+    chunk_size: int = FLA_CHUNK_SIZE,
     output_dtype: torch.dtype = torch.float32,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     r"""
@@ -1178,7 +1178,7 @@ def chunk_kda_fwd(
     output_final_state: bool,
     cu_seqlens: torch.Tensor | None = None,
 ):
-    chunk_size = 64
+    chunk_size = FLA_CHUNK_SIZE
     g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens)
     # the intra Aqk is kept in fp32
     # the computation has very marginal effect on the entire throughput
@@ -1189,6 +1189,7 @@ def chunk_kda_fwd(
         beta=beta,
         scale=scale,
         cu_seqlens=cu_seqlens,
+        chunk_size=chunk_size,
         output_dtype=torch.float32,
     )
     A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)

vllm/model_executor/layers/fla/ops/solve_tril.py

@@ -507,6 +507,7 @@ def merge_16x16_to_64x64_inverse_kernel(
 def solve_tril(
     A: torch.Tensor,
     cu_seqlens: torch.Tensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
     output_dtype: torch.dtype = torch.float,
 ) -> torch.Tensor:
     """
@@ -518,6 +519,8 @@ def solve_tril(
             [B, T, H, BT], where BT should only be 16, 32, or 64.
         cu_seqlens (torch.Tensor):
             The cumulative sequence lengths of the input tensor. Default: `None`.
+        chunk_indices (torch.Tensor):
+            Pre-computed chunk indices. Default: `None`.
         output_dtype (torch.dtype):
             The dtype of the output tensor. Default: `torch.float`.
             If `None`, the output dtype will be the same as the input dtype.
@@ -529,9 +532,8 @@ def solve_tril(
     output_dtype = A.dtype if output_dtype is None else output_dtype
 
     B, T, H, BT = A.shape
-    chunk_indices = (
-        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
-    )
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
     NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
 
     Ai = torch.zeros_like(A, dtype=output_dtype)

vllm/model_executor/layers/fla/ops/utils.py

@@ -154,9 +154,13 @@ is_nvidia_hopper = is_nvidia and (
 )
 use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1"
 is_gather_supported = hasattr(triton.language, "gather")
-is_tma_supported = (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9) and (
-    hasattr(triton.language, "_experimental_make_tensor_descriptor")
-    or hasattr(triton.language, "make_tensor_descriptor")
+is_tma_supported = (
+    is_nvidia_hopper
+    and os.getenv("FLA_USE_TMA", "0") == "1"
+    and (
+        hasattr(triton.language, "_experimental_make_tensor_descriptor")
+        or hasattr(triton.language, "make_tensor_descriptor")
+    )
 )
 
 

vllm/model_executor/layers/fla/ops/wy_fast.py

@@ -123,14 +123,14 @@ def recompute_w_u_fwd(
     g_cumsum: torch.Tensor,
     A: torch.Tensor,
     cu_seqlens: torch.Tensor | None,
+    chunk_indices: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     B, T, Hg, K, V = *k.shape, v.shape[-1]
     H = v.shape[-2]
     BT = A.shape[-1]
 
-    chunk_indices = (
-        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
-    )
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
     NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
     BK = 64
     BV = 64

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

@@ -79,11 +79,8 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
             RoutingMethodType.Default,
             RoutingMethodType.DeepSeekV3,
             RoutingMethodType.Llama4,
-            # NOTE: TRTLLM Kernel has issue with Qwen3.5 router.
-            # Re-enable once the issue is resolved.
-            # https://github.com/vllm-project/vllm/issues/37591
-            # RoutingMethodType.Renormalize,
-            # RoutingMethodType.RenormalizeNaive
+            RoutingMethodType.Renormalize,
+            RoutingMethodType.RenormalizeNaive,
         ]
 
     @staticmethod

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

@@ -277,13 +277,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        """Monolithic kernels need to express router support.
-        Renormalize/RenormalizeNaive are excluded: the monolithic kernel's
-        internal routing for these methods produces output uncorrelated
-        with the modular kernel's output and with Triton kernel's output
-        for Qwen3.5-35B-A3B-FP8.
-        See: https://github.com/vllm-project/vllm/issues/37591
-        """
+        """Monolithic kernels need to express router support."""
         # NOTE(dbari): TopK routing could also be enabled, but need to validate models
         # NOTE(dbari): Default is not implemented and should not be enabled until it is
 
@@ -295,6 +289,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
             return routing_method in [
                 RoutingMethodType.DeepSeekV3,
                 RoutingMethodType.Simulated,
+                RoutingMethodType.Renormalize,
+                RoutingMethodType.RenormalizeNaive,
             ]
         elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):
             # NOTE(dbari): as above, potentially allow others here.
@@ -302,6 +298,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
                 RoutingMethodType.DeepSeekV3,
                 RoutingMethodType.Llama4,
                 RoutingMethodType.Simulated,
+                RoutingMethodType.Renormalize,
+                RoutingMethodType.RenormalizeNaive,
             ]
         else:
             raise ValueError("Unsupported quantization scheme.")

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

@@ -222,6 +222,18 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
                     self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
             else:
                 self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
+        elif current_platform.is_xpu():
+            w13 = layer.w13_weight
+            w2 = layer.w2_weight
+
+            w13.data = w13.transpose(-1, -2).contiguous()
+            w2.data = w2.transpose(-1, -2).contiguous()
+
+            self._setup_kernel(
+                layer=layer,
+                w13=w13,
+                w2=w2,
+            )
         else:
             self._setup_kernel(
                 layer=layer,

vllm/model_executor/layers/layernorm.py

@@ -241,8 +241,12 @@ class RMSNorm(CustomOp):
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """PyTorch-native implementation equivalent to forward()."""
         if residual is None:
+            # TODO(luka): address the weight=None passing issue more generally
             return ir.ops.rms_norm(
-                x, self.weight.data, self.variance_epsilon, self.variance_size_override
+                x,
+                self.weight.data if self.has_weight else None,
+                self.variance_epsilon,
+                self.variance_size_override,
             )
 
         return self.forward_static(

vllm/model_executor/layers/linear.py

@@ -60,7 +60,6 @@ WEIGHT_LOADER_V2_SUPPORTED = [
     "ModelOptFp8PbWoLinearMethod",
     "QuarkLinearMethod",
     "ModelOptNvFp4LinearMethod",
-    "PetitNvFp4LinearMethod",
 ]
 
 

vllm/model_executor/layers/mamba/gdn_linear_attn.py

@@ -163,6 +163,8 @@ class ChunkGatedDeltaRule(CustomOp):
         initial_state: torch.Tensor,
         output_final_state: bool,
         cu_seqlens: torch.Tensor | None = None,
+        chunk_indices: torch.Tensor | None = None,
+        chunk_offsets: torch.Tensor | None = None,
         use_qk_l2norm_in_kernel: bool = True,
     ):
         return fi_chunk_gated_delta_rule(
@@ -187,6 +189,8 @@ class ChunkGatedDeltaRule(CustomOp):
         initial_state: torch.Tensor,
         output_final_state: bool,
         cu_seqlens: torch.Tensor | None = None,
+        chunk_indices: torch.Tensor | None = None,
+        chunk_offsets: torch.Tensor | None = None,
         use_qk_l2norm_in_kernel: bool = True,
     ):
         return fla_chunk_gated_delta_rule(
@@ -198,6 +202,8 @@ class ChunkGatedDeltaRule(CustomOp):
             initial_state=initial_state,
             output_final_state=output_final_state,
             cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            chunk_offsets=chunk_offsets,
             use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
         )
 
@@ -959,6 +965,8 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
                 initial_state=initial_state,
                 output_final_state=True,
                 cu_seqlens=non_spec_query_start_loc,
+                chunk_indices=attn_metadata.chunk_indices,
+                chunk_offsets=attn_metadata.chunk_offsets,
                 use_qk_l2norm_in_kernel=False,
             )
             # Init cache

vllm/model_executor/layers/quantization/__init__.py

@@ -31,8 +31,14 @@ QuantizationMethods = Literal[
     "inc",
     "mxfp4",
     "mxfp8",
-    "petit_nvfp4",
     "cpu_awq",
+    "online",
+    # Below are values of the OnlineQuantScheme enum, specified as strings to
+    # avoid circular import issues. This is here to provide a shortcut where
+    # the user can specify "LLM(..., quantization='fp8_per_tensor')" as
+    # shorthand for creating a more complicated online quant config object
+    "fp8_per_tensor",
+    "fp8_per_block",
 ]
 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
 
@@ -41,7 +47,6 @@ DEPRECATED_QUANTIZATION_METHODS = [
     "fbgemm_fp8",
     "fp_quant",
     "experts_int8",
-    "petit_nvfp4",
 ]
 
 # The customized quantization methods which will be added to this dict.
@@ -103,6 +108,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
         raise ValueError(f"Invalid quantization method: {quantization}")
 
     # lazy import to avoid triggering `torch.compile` too early
+    from vllm.config.quantization import OnlineQuantScheme
     from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
 
     from .awq import AWQConfig
@@ -129,7 +135,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
     from .moe_wna16 import MoeWNA16Config
     from .mxfp4 import Mxfp4Config
     from .mxfp8 import Mxfp8Config
-    from .petit import PetitNvFp4Config
+    from .online.base import OnlineQuantizationConfig
     from .torchao import TorchAOConfig
 
     method_to_config: dict[str, type[QuantizationConfig]] = {
@@ -155,9 +161,21 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
         "inc": INCConfig,
         "mxfp4": Mxfp4Config,
         "mxfp8": Mxfp8Config,
-        "petit_nvfp4": PetitNvFp4Config,
         "cpu_awq": CPUAWQConfig,
+        "online": OnlineQuantizationConfig,
     }
+
+    # Below are values of the OnlineQuantScheme enum. This is here to provide
+    # a shortcut where the user can specify
+    # "LLM(..., quantization='fp8_per_tensor')" as shorthand for creating a
+    # more complicated online quant config object
+    for scheme in OnlineQuantScheme:
+        assert scheme.value not in method_to_config, (
+            f"Online quant scheme {scheme.value!r} conflicts with an "
+            f"existing quantization method"
+        )
+        method_to_config[scheme.value] = OnlineQuantizationConfig
+
     # Update the `method_to_config` with customized quantization methods.
     method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
 

vllm/model_executor/layers/quantization/awq.py

@@ -8,6 +8,7 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
 from transformers import PretrainedConfig
 
 from vllm import _custom_ops as ops
+from vllm import envs
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.layer import FusedMoE
 from vllm.model_executor.layers.linear import (
@@ -273,8 +274,9 @@ class AWQLinearMethod(LinearMethodBase):
 
         # num_tokens >= threshold
         FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
-
-        if FP16_MATMUL_HEURISTIC_CONDITION:
+        # Batch invariant mode requires torch.matmul path
+        # for Triton override
+        if FP16_MATMUL_HEURISTIC_CONDITION or envs.VLLM_BATCH_INVARIANT:
             out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
             out = torch.matmul(reshaped_x, out)
         else:

vllm/model_executor/layers/quantization/awq_marlin.py

@@ -10,6 +10,7 @@ from transformers import PretrainedConfig
 
 import vllm.model_executor.layers.fused_moe  # noqa
 from vllm import _custom_ops as ops
+from vllm import envs
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import (
     MPLinearLayerConfig,
@@ -233,6 +234,11 @@ class AWQMarlinConfig(QuantizationConfig):
     def override_quantization_method(
         cls, hf_quant_cfg, user_quant
     ) -> "QuantizationMethods | None":
+        # Skip override to marlin kernels, as they are not
+        # batch invariant
+        if envs.VLLM_BATCH_INVARIANT:
+            return None
+
         can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg)
         is_valid_user_quant = (
             user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin"

vllm/model_executor/layers/quantization/fp8.py

@@ -497,6 +497,8 @@ class Fp8LinearMethod(LinearMethodBase):
         return self.fp8_linear.apply_weights(layer, x, bias)
 
 
+# TODO(future PR): remove this class in favor of
+# online/fp8.py::Fp8PerTensorOnlineLinearMethod
 class Fp8OnlineLinearMethod(Fp8LinearMethod):
     """Online version of Fp8LinearMethod which loads a full precision checkpoint
     and quantizes weights during loading."""
@@ -919,6 +921,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
         )
 
 
+# TODO(future PR): remove this class in favor of
+# online/fp8.py::Fp8PerTensorOnlineMoEMethod
 class Fp8OnlineMoEMethod(Fp8MoEMethod):
     """MoE method for online FP8 quantization.
     Supports loading quantized FP16/BF16 model checkpoints with dynamic
@@ -1028,6 +1032,10 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
                 layer.w2_weight[expert, :, :]
             )
 
+        if current_platform.is_xpu():
+            w13.data = w13.transpose(-1, -2).contiguous()
+            w2.data = w2.transpose(-1, -2).contiguous()
+
         # Shuffle weights to runtime format and setup kernel.
         self._setup_kernel(
             layer,

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

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

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

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+
+from vllm.config.quantization import (
+    OnlineQuantizationConfigArgs,
+    OnlineQuantScheme,
+)
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoE,
+)
+from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+    UnquantizedFusedMoEMethod,
+)
+from vllm.model_executor.layers.linear import (
+    LinearBase,
+    UnquantizedLinearMethod,
+)
+from vllm.model_executor.layers.quantization import QuantizationMethods
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
+    should_ignore_layer,
+)
+from vllm.model_executor.layers.quantization.online.fp8 import (
+    Fp8PerBlockOnlineLinearMethod,
+    Fp8PerBlockOnlineMoEMethod,
+    Fp8PerTensorOnlineLinearMethod,
+    Fp8PerTensorOnlineMoEMethod,
+)
+
+
+class OnlineQuantizationConfig(QuantizationConfig):
+    """Model-level config class for online quantization (quantize fp16/bf16 weights
+    during model loading, without requiring a pre-quantized checkpoint)."""
+
+    def __init__(
+        self,
+        args: OnlineQuantizationConfigArgs,
+    ) -> None:
+        super().__init__()
+        if (
+            args.global_scheme is None
+            and args.linear_scheme_override is None
+            and args.moe_scheme_override is None
+        ):
+            raise ValueError(
+                "OnlineQuantizationConfig requires at least one of "
+                "global_scheme, linear_scheme_override, or "
+                "moe_scheme_override to be set."
+            )
+        self.args = args
+        self.quant_scheme = args.global_scheme
+        self.ignored_layers: list[str] = args.ignore
+
+    @classmethod
+    def get_name(cls) -> QuantizationMethods:
+        return "online"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        # Note: as more online quant schemes will be added, this
+        # value will become the minimum across all supported schemes.
+        return 75
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "OnlineQuantizationConfig":
+        raise NotImplementedError(
+            "OnlineQuantizationConfig does not support loading from a "
+            "checkpoint config. Use quantization_config or "
+            "quantization='fp8_per_tensor'/'fp8_per_block' instead."
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> "QuantizeMethodBase | None":
+        if isinstance(layer, LinearBase):
+            if should_ignore_layer(
+                prefix,
+                ignore=self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedLinearMethod()
+
+            linear_scheme = self.args.linear_scheme_override or self.args.global_scheme
+            if linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
+                return Fp8PerBlockOnlineLinearMethod()
+            else:
+                return Fp8PerTensorOnlineLinearMethod()
+        elif isinstance(layer, FusedMoE):
+            if should_ignore_layer(
+                prefix,
+                ignore=self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedFusedMoEMethod(layer.moe_config)
+
+            moe_scheme = self.args.moe_scheme_override or self.args.global_scheme
+            if moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
+                return Fp8PerBlockOnlineMoEMethod(layer=layer)
+            else:
+                return Fp8PerTensorOnlineMoEMethod(layer=layer)
+        return None

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

@@ -0,0 +1,632 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch.nn import Module
+
+if TYPE_CHECKING:
+    import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+    from vllm.model_executor.layers.fused_moe import FusedMoE
+    from vllm.model_executor.layers.fused_moe.config import (
+        FusedMoEConfig,
+        FusedMoEQuantConfig,
+    )
+    from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEMethodBase,
+)
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+    select_fp8_moe_backend,
+)
+from vllm.model_executor.layers.linear import (
+    LinearMethodBase,
+)
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    W8A8BlockFp8LinearOp,
+    maybe_post_process_fp8_weight_block,
+    process_fp8_weight_block_strategy,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    kFp8Dynamic128Sym,
+    kFp8DynamicTensorSym,
+    kFp8DynamicTokenSym,
+    kFp8Static128BlockSym,
+    kFp8StaticTensorSym,
+)
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    cutlass_block_fp8_supported,
+    cutlass_fp8_supported,
+)
+from vllm.model_executor.model_loader.reload.layerwise import (
+    initialize_online_processing,
+)
+from vllm.model_executor.parameter import ModelWeightParameter
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import is_deep_gemm_supported, per_block_cast_to_fp8
+
+# ---------------------------------------------------------------------------
+# Online FP8 Linear Methods
+# ---------------------------------------------------------------------------
+
+
+class _Fp8OnlineLinearBase(LinearMethodBase):
+    """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
+    weights onto meta device and materializes them just-in-time."""
+
+    uses_meta_device: bool = True
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+        layer.weight_block_size = None
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                device="meta",  # materialized and processed during loading
+                dtype=params_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        initialize_online_processing(layer)
+
+
+class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
+    """Online tensorwise FP8 linear quantization.
+    Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
+
+    def __init__(self):
+        self.out_dtype = torch.get_default_dtype()
+
+        # Use per-token quantization for better perf if dynamic and cutlass
+        if cutlass_fp8_supported():
+            activation_quant_key = kFp8DynamicTokenSym
+        else:
+            activation_quant_key = kFp8DynamicTensorSym
+
+        self.fp8_linear = init_fp8_linear_kernel(
+            activation_quant_key=activation_quant_key,
+            weight_quant_key=kFp8StaticTensorSym,
+            out_dtype=torch.get_default_dtype(),
+            module_name=self.__class__.__name__,
+        )
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        layer.input_scale = None
+        qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
+
+        # Update layer with new values.
+        replace_parameter(layer, "weight", qweight.t().data)
+        replace_parameter(layer, "weight_scale", weight_scale.data)
+
+        # Prevent duplicate processing (e.g., during weight reload)
+        layer._already_called_process_weights_after_loading = True
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # if batch invariant mode is enabled, use BF16 dequant
+        if envs.VLLM_BATCH_INVARIANT:
+            weight_fp8 = layer.weight.to(torch.bfloat16)
+            weight_scale = layer.weight_scale.to(torch.bfloat16)
+            if weight_scale.numel() == 1:
+                # Per-tensor: simple scalar multiplication
+                weight_bf16 = weight_fp8 * weight_scale
+            else:
+                # Multiple scales (fused modules like QKV)
+                if (
+                    weight_scale.dim() == 1
+                    and weight_scale.shape[0] == weight_fp8.shape[0]
+                ):
+                    # Per-row scaling
+                    weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1)
+                else:
+                    # Fallback
+                    weight_bf16 = weight_fp8 * weight_scale
+            return torch.nn.functional.linear(x, weight_bf16.t(), bias)
+
+        return self.fp8_linear.apply_weights(layer, x, bias)
+
+
+class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
+    """Online blockwise FP8 linear quantization.
+    Loads fp16/bf16 weights and quantizes them per-block during loading."""
+
+    def __init__(self):
+        self.out_dtype = torch.get_default_dtype()
+        self.weight_block_size = [128, 128]
+
+        self.use_deep_gemm = is_deep_gemm_supported()
+        self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
+        self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
+
+        self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
+            weight_group_shape=GroupShape(*self.weight_block_size),
+            act_quant_group_shape=GroupShape(1, self.weight_block_size[0]),
+            cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
+            use_aiter_and_is_supported=self.use_aiter_and_is_supported,
+            use_deep_gemm=self.use_deep_gemm,
+        )
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        super().create_weights(
+            layer,
+            input_size_per_partition,
+            output_partition_sizes,
+            input_size,
+            output_size,
+            params_dtype,
+            **extra_weight_attrs,
+        )
+        layer.weight_block_size = self.weight_block_size
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        layer.input_scale = None
+        block_size = self.weight_block_size
+
+        qweight, weight_scale_inv = per_block_cast_to_fp8(
+            layer.weight, block_size=block_size, use_ue8m0=False
+        )
+
+        qweight, weight_scale_inv = process_fp8_weight_block_strategy(
+            qweight, weight_scale_inv
+        )
+
+        replace_parameter(layer, "weight", qweight.data)
+        replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
+
+        maybe_post_process_fp8_weight_block(layer)
+
+        # Prevent duplicate processing (e.g., during weight reload)
+        layer._already_called_process_weights_after_loading = True
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        assert self.weight_block_size is not None
+
+        # Note: batch invariance already handled in the function below
+        return self.w8a8_block_fp8_linear.apply(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale_inv,
+            input_scale=layer.input_scale,
+            bias=bias,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Online FP8 MoE Methods
+# ---------------------------------------------------------------------------
+
+
+class _Fp8OnlineMoEBase(FusedMoEMethodBase):
+    """Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint
+    weights onto meta device and materializes them just-in-time."""
+
+    uses_meta_device: bool = True
+
+    # Declared here for mypy; actual values are set in __init__.
+    fp8_backend: "Fp8MoeBackend"
+    experts_cls: "type[mk.FusedMoEExperts] | None"
+    weight_scale_name: str
+    weight_block_size: list[int] | None
+    moe: "FusedMoEConfig"
+    is_monolithic: bool
+    moe_quant_config: "FusedMoEQuantConfig | None"
+    moe_kernel: "mk.FusedMoEKernel | None"
+
+    def __init__(
+        self,
+        *,
+        weight_block_size: list[int] | None,
+        layer: torch.nn.Module,
+    ):
+        super().__init__(layer.moe_config)
+        self.weight_block_size = weight_block_size
+        self.block_quant: bool = self.weight_block_size is not None
+        self.weight_scale_name = (
+            "weight_scale_inv" if self.block_quant else "weight_scale"
+        )
+
+        # Set weight key and activation key for kernel compatibility
+        if self.block_quant:
+            weight_key = kFp8Static128BlockSym
+            activation_key = kFp8Dynamic128Sym
+        else:
+            weight_key = kFp8StaticTensorSym
+            activation_key = kFp8DynamicTensorSym
+
+        # Select Fp8 MoE backend
+        self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
+            config=self.moe,
+            weight_key=weight_key,
+            activation_key=activation_key,
+            allow_vllm_cutlass=False,
+        )
+
+    def create_weights(
+        self,
+        layer: Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        layer.num_experts = num_experts
+        layer.orig_dtype = params_dtype
+        layer.weight_block_size = None
+
+        # WEIGHTS
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size,
+                device="meta",
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition,
+                device="meta",  # materialized and processed during loading
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        # BIASES (for models like GPT-OSS that have biased MoE)
+        if self.moe.has_bias:
+            w13_bias = torch.nn.Parameter(
+                torch.zeros(
+                    num_experts,
+                    2 * intermediate_size_per_partition,
+                    device="meta",  # materialized and processed during loading
+                    dtype=layer.orig_dtype,
+                ),
+                requires_grad=False,
+            )
+            layer.register_parameter("w13_bias", w13_bias)
+            set_weight_attrs(w13_bias, extra_weight_attrs)
+
+            w2_bias = torch.nn.Parameter(
+                torch.zeros(
+                    num_experts,
+                    hidden_size,
+                    device="meta",  # materialized and processed during loading
+                    dtype=layer.orig_dtype,
+                ),
+                requires_grad=False,
+            )
+            layer.register_parameter("w2_bias", w2_bias)
+            set_weight_attrs(w2_bias, extra_weight_attrs)
+
+        layer.w13_input_scale = None
+        layer.w2_input_scale = None
+
+        initialize_online_processing(layer)
+
+    def _setup_kernel(
+        self,
+        layer: "FusedMoE",
+        w13: torch.Tensor,
+        w2: torch.Tensor,
+        w13_scale: torch.Tensor,
+        w2_scale: torch.Tensor,
+        w13_input_scale: torch.Tensor | None,
+        w2_input_scale: torch.Tensor | None,
+    ) -> None:
+        from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+            convert_to_fp8_moe_kernel_format,
+            make_fp8_moe_kernel,
+        )
+
+        # Shuffle weights to runtime format.
+        w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
+            fp8_backend=self.fp8_backend,
+            layer=layer,
+            w13=w13,
+            w2=w2,
+            w13_scale=w13_scale,
+            w2_scale=w2_scale,
+            w13_input_scale=w13_input_scale,
+            w2_input_scale=w2_input_scale,
+        )
+
+        # Replace parameters with updated versions. Note that this helper
+        # function ensures the replacement is compatible with RL weight reloads.
+        replace_parameter(layer, "w13_weight", w13)
+        replace_parameter(layer, "w2_weight", w2)
+        replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
+        replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
+
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        if self.moe_quant_config:
+            assert self.experts_cls is not None
+            self.moe_kernel = make_fp8_moe_kernel(
+                moe_quant_config=self.moe_quant_config,
+                moe_config=self.moe,
+                fp8_backend=self.fp8_backend,
+                experts_cls=self.experts_cls,
+                routing_tables=layer._maybe_init_expert_routing_tables(),
+                shared_experts=layer.shared_experts,
+            )
+
+    def maybe_make_prepare_finalize(
+        self,
+        routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+    ) -> "mk.FusedMoEPrepareAndFinalizeModular | None":
+        raise ValueError(
+            f"{self.__class__.__name__} uses the new modular kernel "
+            "initialization logic. This function should not be called."
+        )
+
+    def get_fused_moe_quant_config(
+        self, layer: torch.nn.Module
+    ) -> "FusedMoEQuantConfig":
+        from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+            make_fp8_moe_quant_config,
+        )
+
+        w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
+        w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
+        a1_scale = layer.w13_input_scale
+        a2_scale = layer.w2_input_scale
+
+        quant_config = make_fp8_moe_quant_config(
+            fp8_backend=self.fp8_backend,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=self.weight_block_size,
+        )
+
+        # Inject biases into the quant config if the model has them
+        # (e.g. GPT-OSS biased MoE)
+        if quant_config is not None and self.moe.has_bias:
+            w13_bias = getattr(layer, "w13_bias", None)
+            w2_bias = getattr(layer, "w2_bias", None)
+            if w13_bias is not None:
+                quant_config._w1.bias = w13_bias
+            if w2_bias is not None:
+                quant_config._w2.bias = w2_bias
+
+        return quant_config
+
+    @property
+    def supports_eplb(self) -> bool:
+        return True
+
+    def apply_monolithic(
+        self,
+        layer: "FusedMoE",
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        assert self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply_monolithic(
+            x,
+            layer.w13_weight,
+            layer.w2_weight,
+            router_logits,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            num_expert_group=layer.num_expert_group,
+            topk_group=layer.topk_group,
+            e_score_correction_bias=layer.e_score_correction_bias,
+            routed_scaling_factor=layer.routed_scaling_factor,
+        )
+
+    def apply(
+        self,
+        layer: "FusedMoE",
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        assert not self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply(
+            x,
+            layer.w13_weight,
+            layer.w2_weight,
+            topk_weights,
+            topk_ids,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            shared_experts_input=shared_experts_input,
+        )
+
+
+class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
+    """Online tensorwise FP8 MoE quantization.
+    Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
+
+    def __init__(
+        self,
+        *,
+        layer: torch.nn.Module,
+    ):
+        super().__init__(
+            weight_block_size=None,
+            layer=layer,
+        )
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        # TODO(@ksayers): inplace fp8 quant kernel, initialize scales with ones
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        # If checkpoint is fp16, quantize in place.
+        fp8_dtype = current_platform.fp8_dtype()
+        w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
+        w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
+        w13_scale = torch.ones(
+            layer.num_experts, device=w13.device, dtype=torch.float32
+        )
+        w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32)
+        layer.w13_input_scale = None
+        layer.w2_input_scale = None
+
+        for expert in range(layer.local_num_experts):
+            w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
+                layer.w13_weight[expert, :, :]
+            )
+            w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
+                layer.w2_weight[expert, :, :]
+            )
+
+        # Shuffle weights to runtime format and setup kernel.
+        self._setup_kernel(
+            layer,
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            w13_input_scale=layer.w13_input_scale,
+            w2_input_scale=layer.w2_input_scale,
+        )
+
+        # Prevent duplicate processing (e.g., during weight reload)
+        layer._already_called_process_weights_after_loading = True
+
+
+class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
+    """Online blockwise FP8 MoE quantization.
+    Loads fp16/bf16 weights and quantizes them per-block during loading."""
+
+    def __init__(
+        self,
+        *,
+        layer: torch.nn.Module,
+    ):
+        super().__init__(
+            weight_block_size=[128, 128],
+            layer=layer,
+        )
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        fp8_dtype = current_platform.fp8_dtype()
+        w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
+        w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
+
+        block_size = self.weight_block_size
+        assert block_size is not None
+        block_n, block_k = block_size
+
+        # Create block-shaped scales (computed here rather than in
+        # create_weights because online quant doesn't need them until now).
+        num_experts = layer.local_num_experts
+        _, w13_out, w13_in = layer.w13_weight.shape
+        _, w2_out, w2_in = layer.w2_weight.shape
+
+        w13_scale = torch.ones(
+            num_experts,
+            (w13_out + block_n - 1) // block_n,
+            (w13_in + block_k - 1) // block_k,
+            dtype=torch.float32,
+            device=w13.device,
+        )
+        w2_scale = torch.ones(
+            num_experts,
+            (w2_out + block_n - 1) // block_n,
+            (w2_in + block_k - 1) // block_k,
+            dtype=torch.float32,
+            device=w2.device,
+        )
+
+        for expert in range(num_experts):
+            w13[expert], w13_scale[expert] = per_block_cast_to_fp8(
+                layer.w13_weight[expert],
+                block_size=block_size,
+                use_ue8m0=False,
+            )
+            w2[expert], w2_scale[expert] = per_block_cast_to_fp8(
+                layer.w2_weight[expert],
+                block_size=block_size,
+                use_ue8m0=False,
+            )
+
+        layer.weight_block_size = block_size
+
+        # Shuffle weights to runtime format and setup kernel.
+        self._setup_kernel(
+            layer,
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            layer.w13_input_scale,
+            layer.w2_input_scale,
+        )
+
+        # Prevent duplicate processing (e.g., during weight reload)
+        layer._already_called_process_weights_after_loading = True

vllm/model_executor/layers/quantization/petit.py

@@ -1,319 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py
-
-from typing import Any
-
-import regex as re
-import torch
-from torch.nn.parameter import Parameter
-
-from vllm.logger import init_logger
-from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.linear import (
-    LinearBase,
-    LinearMethodBase,
-    UnquantizedLinearMethod,
-)
-from vllm.model_executor.layers.quantization import QuantizationMethods
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig,
-    QuantizeMethodBase,
-)
-from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
-from vllm.model_executor.layers.quantization.utils.petit_utils import (
-    apply_petit_nvfp4_linear,
-    prepare_nvfp4_layer_for_petit,
-    verify_petit_nvfp4_supported,
-)
-from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
-from vllm.model_executor.parameter import ModelWeightParameter, PerTensorScaleParameter
-from vllm.platforms import current_platform
-
-# Initialize logger for the module
-logger = init_logger(__name__)
-
-
-# Configuration class to support the NVFP4 quantized model
-# generated by the ModelOpt quantization tool
-class PetitNvFp4Config(QuantizationConfig):
-    """Config class for Petit FP4."""
-
-    def __init__(
-        self,
-        is_checkpoint_nvfp4_serialized: bool = False,
-        kv_cache_quant_algo: str | None = None,
-        group_size: int | None = None,
-        exclude_modules: list[str] | None = None,
-    ) -> None:
-        self._check_hardware_support()
-        self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
-        if is_checkpoint_nvfp4_serialized:
-            logger.warning(
-                "Detected nvfp4 checkpoint. Please note that the "
-                "format is experimental and subject to change."
-            )
-        self.group_size = group_size
-        self.kv_cache_quant_algo = kv_cache_quant_algo
-        self.exclude_modules = exclude_modules
-
-    def _check_hardware_support(self) -> None:
-        """
-        Verifies that the current hardware is supported by the Petit backend.
-        This backend is specifically designed for AMD GPUs and is not
-        supported on the CUDA platform.
-        """
-        # This check ensures the code is NOT running on an NVIDIA GPU.
-        if current_platform.is_cuda():
-            raise ValueError(
-                "The 'petit' quantization backend is designed for AMD GPUs "
-                "and is not supported on the CUDA platform. For NVIDIA GPUs, "
-                "please use a different quantization method such as FP8, AWQ, "
-                "or GPTQ."
-            )
-
-    @classmethod
-    def get_name(cls) -> QuantizationMethods:
-        return "petit_nvfp4"
-
-    @classmethod
-    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
-        return [torch.bfloat16, torch.half]
-
-    @classmethod
-    def get_min_capability(cls) -> int:
-        # Petit supports the gfx90a and gfx942 GPUs
-        return 90
-
-    @classmethod
-    def get_config_filenames(cls) -> list[str]:
-        return ["hf_quant_config.json"]
-
-    @classmethod
-    def from_config(cls, config: dict[str, Any]) -> "PetitNvFp4Config":
-        qc = cls.get_from_keys(config, ["quantization"])
-
-        quant_method_raw = qc.get("quant_algo")
-        if not isinstance(quant_method_raw, str) or not quant_method_raw:
-            raise ValueError("Missing or invalid 'quant_algo' in quantization config.")
-        quant_method = quant_method_raw.upper()
-
-        group_size_raw = qc.get("group_size")
-        if not isinstance(group_size_raw, int):
-            raise ValueError(
-                "Missing or invalid 'group_size' (int) in hf_quant_config.json."
-            )
-        group_size = group_size_raw
-
-        verify_petit_nvfp4_supported(quant_method, group_size)
-
-        kv_cache_quant_algo_raw = qc.get("kv_cache_quant_algo") or "auto"
-        if not isinstance(kv_cache_quant_algo_raw, str):
-            raise ValueError("'kv_cache_quant_algo' must be a string if provided.")
-        kv_cache_quant_algo = kv_cache_quant_algo_raw
-
-        exclude_raw = qc.get("exclude_modules", [])
-        if exclude_raw is None:
-            exclude_modules: list[str] = []
-        elif isinstance(exclude_raw, list) and all(
-            isinstance(x, str) for x in exclude_raw
-        ):
-            exclude_modules = exclude_raw
-        else:
-            raise ValueError("'exclude_modules' must be a list[str] (or omitted).")
-
-        is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method
-
-        return cls(
-            is_checkpoint_nvfp4_serialized=is_checkpoint_nvfp4_serialized,
-            kv_cache_quant_algo=kv_cache_quant_algo,
-            group_size=group_size,
-            exclude_modules=exclude_modules,
-        )
-
-    @classmethod
-    def override_quantization_method(
-        cls, hf_quant_cfg, user_quant
-    ) -> QuantizationMethods | None:
-        if not current_platform.is_rocm():
-            return None
-
-        qc = hf_quant_cfg.get("quantization", hf_quant_cfg)
-        algo = (qc.get("quant_algo") or qc.get("quant_method") or "").upper()
-        if algo in ("NVFP4", "MODELOPT_FP4", "MODELOPT"):
-            return cls.get_name()  # "petit_nvfp4"
-        return None
-
-    @classmethod
-    def is_petit_nvfp4_compatible(cls, quant_config: dict[str, Any]) -> bool:
-        qc = quant_config.get("quantization", quant_config)
-        algo = (qc.get("quant_algo") or qc.get("quant_method") or "").upper()
-        return algo == "NVFP4"
-
-    def is_layer_excluded(self, prefix: str, exclude_modules: list[str]) -> bool:
-        for pattern in exclude_modules:
-            regex_str = pattern.replace(".", r"\.").replace("*", r".*")
-            if re.fullmatch(regex_str, prefix):
-                return True
-        return False
-
-    def get_quant_method(
-        self, layer: torch.nn.Module, prefix: str
-    ) -> "QuantizeMethodBase | None":
-        exclude = self.require_exclude_modules()
-
-        if isinstance(layer, LinearBase):
-            if is_layer_skipped(prefix, exclude) or self.is_layer_excluded(
-                prefix, exclude
-            ):
-                return UnquantizedLinearMethod()
-            return PetitNvFp4LinearMethod(self)
-        elif isinstance(layer, Attention):
-            return PetitFp8KVCacheMethod(self)
-        return None
-
-    def get_scaled_act_names(self) -> list[str]:
-        return []
-
-    def require_group_size(self) -> int:
-        if self.group_size is None:
-            logger.warning("group_size not set; defaulting to 16 for NVFP4.")
-            return 16
-        return self.group_size
-
-    def require_kv_cache_quant_algo(self) -> str:
-        return self.kv_cache_quant_algo or "auto"
-
-    def require_exclude_modules(self) -> list[str]:
-        return list(self.exclude_modules or [])
-
-
-class PetitFp8KVCacheMethod(BaseKVCacheMethod):
-    """
-    Supports loading kv-cache scaling factors from FP8 checkpoints.
-    """
-
-    def __init__(self, quant_config: PetitNvFp4Config):
-        super().__init__(quant_config)
-
-
-class PetitNvFp4LinearMethod(LinearMethodBase):
-    """Linear method for NVFP4.
-    Supports loading NVFP4 checkpoints with the following structure:
-
-    |Tensor Name           | datatype      |  shape      |
-    |----------------------------------------------------|
-    |input_scale           | torch.float32 | scalar      |
-    |weight                | NVFP4(SE2M1)  | [1, X, y/2] |
-    |weight_scale          | FP8-E4M3      | [X, Y]      |
-    |weight_scale_2        | torch.float32 | scalar      |
-
-    The weights are quantized per block of 16 elements.
-    Args: quant_config: The ModelOpt quantization config.
-    """
-
-    def __init__(self, quant_config: PetitNvFp4Config):
-        self.quant_config = quant_config
-
-    def create_weights(
-        self,
-        layer: torch.nn.Module,
-        input_size_per_partition: int,
-        output_partition_sizes: list[int],
-        input_size: int,
-        output_size: int,
-        params_dtype: torch.dtype,
-        **extra_weight_attrs,
-    ):
-        del input_size, output_size
-        if not self.quant_config.is_checkpoint_nvfp4_serialized:
-            raise ValueError(
-                "NVFP4 quantization was selected, "
-                " dynamic quantization is not supported."
-            )
-
-        output_size_per_partition = sum(output_partition_sizes)
-        weight_loader = extra_weight_attrs.get("weight_loader")
-
-        layer.logical_widths = output_partition_sizes
-
-        layer.input_size_per_partition = input_size_per_partition
-        layer.output_size_per_partition = output_size_per_partition
-        if input_size_per_partition % 16 != 0:
-            raise ValueError(
-                "Unsupported model when in features size is not multiple of 16"
-            )
-
-        weight_dtype = (
-            torch.float8_e4m3fn
-            if self.quant_config.is_checkpoint_nvfp4_serialized
-            else params_dtype
-        )
-
-        weight = ModelWeightParameter(
-            data=torch.empty(
-                # 2 fp4 data is packed in one uint8 in the input dimension
-                output_size_per_partition,
-                input_size_per_partition // 2,
-                dtype=torch.uint8,
-            ),
-            input_dim=1,
-            output_dim=0,
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight", weight)
-
-        input_scale = PerTensorScaleParameter(
-            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
-            weight_loader=weight_loader,
-        )
-
-        layer.register_parameter("input_scale", input_scale)
-
-        weight_scale_2 = PerTensorScaleParameter(
-            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight_scale_2", weight_scale_2)
-
-        group_size = self.quant_config.require_group_size()
-        weight_scale = ModelWeightParameter(
-            data=torch.empty(
-                output_size_per_partition,
-                input_size_per_partition // group_size,
-                dtype=weight_dtype,
-            ),
-            input_dim=1,
-            output_dim=0,
-            weight_loader=weight_loader,
-        )
-
-        layer.register_parameter("weight_scale", weight_scale)
-
-    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        input_scale_2 = layer.input_scale.max().to(torch.float32)
-        weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
-        layer.input_scale = Parameter(input_scale_2, requires_grad=False)
-        layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False)
-        layer.alpha = Parameter(
-            layer.input_scale * layer.weight_scale_2, requires_grad=False
-        )
-
-        prepare_nvfp4_layer_for_petit(layer)
-        del layer.input_scale
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        return apply_petit_nvfp4_linear(
-            input=x,
-            weight=layer.weight,
-            weight_scale=layer.weight_scale,
-            weight_scale_2=layer.weight_scale_2,
-            size_n=layer.output_size_per_partition,
-            size_k=layer.input_size_per_partition,
-            bias=bias,
-        )

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

@@ -88,6 +88,13 @@ if current_platform.is_rocm():
     def round_int8(x):
         return tl.extra.hip.libdevice.round(x).to(tl.int8)
 
+
+elif current_platform.is_xpu():
+
+    @triton.jit
+    def round_int8(x):
+        return tl.extra.intel.libdevice.round(x).to(tl.int8)
+
 else:
 
     @triton.jit

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

@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING
-
-import torch
-
-# TYPE_CHECKING is used for static type analysis to prevent circular imports.
-if TYPE_CHECKING:
-    from types import ModuleType
-
-# 1. Create a global variable as a placeholder for the module
-_petit_kernel: "ModuleType | None" = None
-
-_PETIT_INSTALL_MSG = (
-    "Petit is not installed. Please install it with `pip install petit-kernel`."
-)
-
-
-def _import_petit_kernel() -> "ModuleType":
-    """
-    A helper function to handle the lazy import.
-    The first time this function is called, it will import the petit_kernel
-    library and store it in the global _petit_kernel variable.
-    Subsequent calls will return the already-loaded module directly.
-    """
-    global _petit_kernel
-    if _petit_kernel is not None:
-        return _petit_kernel
-
-    try:
-        import petit_kernel
-
-        _petit_kernel = petit_kernel
-        return _petit_kernel
-    except ImportError:
-        # The 'from None' syntax prevents chaining the original ImportError,
-        # making the traceback cleaner.
-        raise ImportError(_PETIT_INSTALL_MSG) from None
-
-
-def _check_petit_nvfp4_supported(
-    quant_method: str, group_size: int | None
-) -> tuple[bool, str | None]:
-    if quant_method != "NVFP4":
-        return (
-            False,
-            (
-                "Petit currently only supports: NVFP4 quantizations in sglang. "
-                "Please check the `hf_quant_config.json` file for your model's "
-                "quant configuration."
-            ),
-        )
-    if group_size is not None and group_size != 16:
-        return (
-            False,
-            "Petit currently only supports: group_size=16 quantizations.",
-        )
-    return (True, None)
-
-
-def verify_petit_nvfp4_supported(quant_method: str, group_size: int | None) -> None:
-    supported, error_msg = _check_petit_nvfp4_supported(quant_method, group_size)
-    if not supported:
-        assert error_msg is not None
-        raise ValueError(error_msg)
-
-
-def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None:
-    # 2. Call _import_petit_kernel() to trigger (or get) the import.
-    petit_kernel = _import_petit_kernel()
-
-    # Repack weights to petit format
-    part_size_n = layer.output_size_per_partition
-    part_size_k = layer.input_size_per_partition
-    qweight = layer.weight.view(torch.int32).contiguous()
-
-    # 3. Call functions through the imported module variable.
-    petit_qweight = petit_kernel.repack_nvfp4(
-        qweight, size_n=part_size_n, size_k=part_size_k
-    )
-    layer.weight = torch.nn.Parameter(petit_qweight, requires_grad=False)
-
-    # Permute scales
-    weight_scale = petit_kernel.process_nvfp4_scales(
-        scales=layer.weight_scale, size_k=part_size_k, size_n=part_size_n
-    )
-    layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
-
-
-def apply_petit_nvfp4_linear(
-    input: torch.Tensor,
-    weight: torch.Tensor,
-    weight_scale: torch.Tensor,
-    weight_scale_2: torch.Tensor,
-    size_n: int,
-    size_k: int,
-    bias: torch.Tensor | None = None,
-) -> torch.Tensor:
-    # Trigger (or get) the import here as well.
-    petit_kernel = _import_petit_kernel()
-
-    reshaped_x = input.reshape(-1, input.shape[-1])
-    out_shape = input.shape[:-1] + (size_n,)
-
-    # TODO: Use auto-tuning to find the performant solution_id
-    # Call the function via the module variable.
-    output = petit_kernel.mul_nvfp4_a16(
-        a=reshaped_x,
-        b=weight,
-        s=weight_scale,
-        global_scale=weight_scale_2,
-        size_m=reshaped_x.size(0),
-        size_n=size_n,
-        size_k=size_k,
-        solution_id=-1,
-    )
-    if bias is not None:
-        output.add_(bias)  # In-place add
-
-    return output.reshape(out_shape)

vllm/model_executor/model_loader/weight_utils.py

@@ -296,6 +296,13 @@ def get_quant_config(
         )
 
     if hf_quant_config is not None:
+        if model_config.quantization_config is not None:
+            raise ValueError(
+                "Setting `quantization_config` for online "
+                "quantization when the model checkpoint already "
+                "has a `quantization_config` is not supported"
+            )
+
         # For modelopt_mixed, config.json's quantization_config may or may
         # not contain the per-layer quantized_layers map.  Newer checkpoints
         # embed it directly; older ones keep it only in hf_quant_config.json.
@@ -319,6 +326,12 @@ def get_quant_config(
     quantization_config_file = hf_overrides.get("quantization_config_file", None)
     if quantization_config_file is not None:
         if hasattr(quant_cls, "from_config_file"):
+            if model_config.quantization_config is not None:
+                raise ValueError(
+                    "Setting `quantization_config` for online "
+                    "quantization when the model checkpoint already "
+                    "has a `quantization_config` is not supported"
+                )
             return quant_cls.from_config_file(quantization_config_file)
         else:
             raise NotImplementedError(
@@ -329,6 +342,12 @@ def get_quant_config(
     quantization_config_json = hf_overrides.get("quantization_config_dict_json", None)
     if quantization_config_json is not None:
         if hasattr(quant_cls, "from_config_dict_json"):
+            if model_config.quantization_config is not None:
+                raise ValueError(
+                    "Setting `quantization_config` for online "
+                    "quantization when the model checkpoint already "
+                    "has a `quantization_config` is not supported"
+                )
             return quant_cls.from_config_dict_json(quantization_config_json)
         else:
             raise NotImplementedError(
@@ -337,6 +356,19 @@ def get_quant_config(
                 f"{quant_cls}"
             )
 
+    # Online quantization doesn't read from checkpoint configs — it quantizes
+    # fp16/bf16 weights on the fly during loading.
+    if model_config.quantization_config is not None:
+        from vllm.config.quantization import OnlineQuantizationConfigArgs
+        from vllm.model_executor.layers.quantization.online.base import (
+            OnlineQuantizationConfig,
+        )
+
+        assert isinstance(
+            model_config.quantization_config, OnlineQuantizationConfigArgs
+        )
+        return OnlineQuantizationConfig(args=model_config.quantization_config)
+
     # Inflight BNB quantization
     if model_config.quantization == "bitsandbytes":
         return quant_cls.from_config({})

vllm/model_executor/models/arctic.py

@@ -16,7 +16,6 @@ from vllm.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
-from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
@@ -42,6 +41,7 @@ from vllm.transformers_utils.configs.arctic import ArcticConfig
 
 from .interfaces import SupportsPP, SupportsQuant
 from .utils import (
+    AutoWeightsLoader,
     extract_layer_index,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
@@ -49,8 +49,6 @@ from .utils import (
     maybe_prefix,
 )
 
-logger = init_logger(__name__)
-
 
 class ArcticMLP(nn.Module):
     def __init__(
@@ -384,6 +382,7 @@ class ArcticModel(nn.Module):
         cache_config = vllm_config.cache_config
         quant_config = vllm_config.quant_config
 
+        self.config = config
         self.vocab_size = config.vocab_size
         self.embed_tokens = VocabParallelEmbedding(
             self.vocab_size, config.hidden_size, org_num_embeddings=self.vocab_size
@@ -426,6 +425,116 @@ class ArcticModel(nn.Module):
         hidden_states = self.norm(hidden_states)
         return hidden_states
 
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+
+        mlp_params_mapping: list[tuple[str, str, int]] = []
+        expert_params_mapping: list[tuple[str, str, int]] = []
+
+        for layer in range(self.config.num_hidden_layers):
+            is_moe_layer = (layer + 1) % self.config.moe_layer_frequency == 0
+            if is_moe_layer and self.config.use_residual:
+                mlp_params_mapping.append(
+                    (
+                        f"layers.{layer}.residual_mlp.w13.weight",
+                        f"layers.{layer}.residual_mlp.w1.weight",
+                        0,
+                    )
+                )
+                mlp_params_mapping.append(
+                    (
+                        f"layers.{layer}.residual_mlp.w13.weight",
+                        f"layers.{layer}.residual_mlp.w3.weight",
+                        1,
+                    )
+                )
+
+            if is_moe_layer:
+                for expert_id in range(self.config.num_local_experts):
+                    expert_params_mapping.append(
+                        ("ws", f"experts.{expert_id}.w1.weight", expert_id)
+                    )
+                    expert_params_mapping.append(
+                        ("w2s", f"experts.{expert_id}.w2.weight", expert_id)
+                    )
+                    expert_params_mapping.append(
+                        ("ws", f"experts.{expert_id}.w3.weight", expert_id)
+                    )
+            else:
+                mlp_params_mapping.append(
+                    (
+                        f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
+                        f"layers.{layer}.block_sparse_moe.mlp.w1.weight",
+                        0,
+                    )
+                )
+                mlp_params_mapping.append(
+                    (
+                        f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
+                        f"layers.{layer}.block_sparse_moe.mlp.w3.weight",
+                        1,
+                    )
+                )
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                for param_name, weight_name, shard_id in mlp_params_mapping:
+                    if weight_name not in name:
+                        continue
+                    name = name.replace(weight_name, param_name)
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    param = params_dict[name]
+                    weight_loader = param.weight_loader
+                    weight_loader(param, loaded_weight, shard_id)
+                    break
+                else:
+                    for param_name, weight_name, shard_id in expert_params_mapping:
+                        if weight_name not in name:
+                            continue
+                        name = name.replace(weight_name, param_name)
+                        if is_pp_missing_parameter(name, self):
+                            continue
+                        param = params_dict[name]
+                        weight_loader = param.weight_loader
+                        weight_loader(
+                            param, loaded_weight, weight_name, expert_id=shard_id
+                        )
+                        break
+                    else:
+                        if name.endswith(".bias") and name not in params_dict:
+                            continue
+                        if is_pp_missing_parameter(name, self):
+                            continue
+                        param = params_dict[name]
+                        weight_loader = getattr(
+                            param, "weight_loader", default_weight_loader
+                        )
+                        weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
 
 class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
     packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
@@ -478,117 +587,8 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-
-        mlp_params_mapping: list[tuple[str, str, int]] = []
-        expert_params_mapping: list[tuple[str, str, int]] = []
-        num_layers = self.config.num_hidden_layers
-
-        for layer in range(num_layers):
-            mlp_params_mapping.append(
-                (
-                    f"layers.{layer}.residual_mlp.w13.weight",
-                    f"layers.{layer}.residual_mlp.w1.weight",
-                    0,
-                )
-            )
-            mlp_params_mapping.append(
-                (
-                    f"layers.{layer}.residual_mlp.w13.weight",
-                    f"layers.{layer}.residual_mlp.w3.weight",
-                    1,
-                )
-            )
-            if layer % 2 == 0:
-                # MLP layers
-                mlp_params_mapping.append(
-                    (
-                        f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
-                        f"layers.{layer}.block_sparse_moe.mlp.w1.weight",
-                        0,
-                    )
-                )
-                mlp_params_mapping.append(
-                    (
-                        f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
-                        f"layers.{layer}.block_sparse_moe.mlp.w3.weight",
-                        1,
-                    )
-                )
-            else:
-                # MoE layers
-                for expert_id in range(self.config.num_local_experts):
-                    expert_params_mapping.append(
-                        ("ws", f"experts.{expert_id}.w1.weight", expert_id)
-                    )
-                    expert_params_mapping.append(
-                        ("w2s", f"experts.{expert_id}.w2.weight", expert_id)
-                    )
-                    expert_params_mapping.append(
-                        ("ws", f"experts.{expert_id}.w3.weight", expert_id)
-                    )
-
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-
-        logger.info(
-            "It will take ~10 minutes loading from the 16-bit weights. "
-            "Alternatively, use the prequantized 8-bit weights of arctic "
-            "and set load-format to `sharded_state` will accelerate loading."
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
         )
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                for param_name, weight_name, shard_id in mlp_params_mapping:
-                    if weight_name not in name:
-                        continue
-                    name = name.replace(weight_name, param_name)
-                    if is_pp_missing_parameter(name, self):
-                        continue
-                    param = params_dict[name]
-                    weight_loader = param.weight_loader
-                    weight_loader(param, loaded_weight, shard_id)
-                    break
-                else:
-                    for param_name, weight_name, shard_id in expert_params_mapping:
-                        if weight_name not in name:
-                            continue
-                        name = name.replace(weight_name, param_name)
-                        if is_pp_missing_parameter(name, self):
-                            continue
-                        param = params_dict[name]
-                        weight_loader = param.weight_loader
-                        weight_loader(
-                            param, loaded_weight, weight_name, expert_id=shard_id
-                        )
-                        break
-                    else:
-                        if name.endswith(".bias") and name not in params_dict:
-                            continue
-                        if is_pp_missing_parameter(name, self):
-                            continue
-                        param = params_dict[name]
-
-                        weight_loader = getattr(
-                            param, "weight_loader", default_weight_loader
-                        )
-                        weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        return loader.load_weights(weights)

vllm/model_executor/models/deepseek_mtp.py

@@ -184,11 +184,16 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
+        self.quant_config = vllm_config.quant_config
         self.model = DeepSeekMultiTokenPredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
         # Set MoE hyperparameters
         self.set_moe_parameters()
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
 
     def set_moe_parameters(self):
         self.expert_weights = []
@@ -241,11 +246,16 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
             ("gate_up_proj", "up_proj", 1),
             ("fused_qkv_a_proj", "q_a_proj", 0),
             ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
-            # Fused indexer wk + weights_proj
-            ("wk_weights_proj", "wk", 0),
-            ("wk_weights_proj", "weights_proj", 1),
         ]
 
+        if self.is_fp4_ckpt:
+            # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
+            indexer_fused_mapping = [
+                ("wk_weights_proj", "wk", 0),
+                ("wk_weights_proj", "weights_proj", 1),
+            ]
+            stacked_params_mapping.extend(indexer_fused_mapping)
+
         expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
             self,
             ckpt_gate_proj_name="gate_proj",

vllm/model_executor/models/deepseek_v2.py

@@ -625,6 +625,11 @@ class Indexer(nn.Module):
         super().__init__()
         self.vllm_config = vllm_config
         self.config = config
+        self.quant_config = quant_config
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
         # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
         self.topk_tokens = config.index_topk
         self.n_head = config.index_n_heads  # 64
@@ -639,18 +644,36 @@ class Indexer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.wq_b",
         )
-        # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
-        # weights_proj does not get quantized, so we run both with quant_config=None
-        # wk may be upcasted from the default quant; experiments show fusion is always
-        # faster unless WK proj is in FP4, which is not the case for all known quants.
-        self.wk_weights_proj = MergedColumnParallelLinear(
-            hidden_size,
-            [self.head_dim, self.n_head],
-            bias=False,
-            quant_config=None,
-            disable_tp=True,
-            prefix=f"{prefix}.wk_weights_proj",
-        )
+        if self.is_fp4_ckpt:
+            # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
+            # weights_proj does not get quantized,
+            # so we run both with quant_config=None
+            # wk may be upcasted from the default quant;
+            # experiments show fusion is always faster unless WK proj is in FP4,
+            # which is not the case for all known quants.
+            self.wk_weights_proj = MergedColumnParallelLinear(
+                hidden_size,
+                [self.head_dim, self.n_head],
+                bias=False,
+                quant_config=None,
+                disable_tp=True,
+                prefix=f"{prefix}.wk_weights_proj",
+            )
+        else:
+            self.wk = ReplicatedLinear(
+                hidden_size,
+                self.head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.wk",
+            )
+            self.weights_proj = ReplicatedLinear(
+                hidden_size,
+                self.n_head,
+                bias=False,
+                quant_config=None,
+                prefix=f"{prefix}.weights_proj",
+            )
         self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
         self.softmax_scale = self.head_dim**-0.5
 
@@ -691,11 +714,14 @@ class Indexer(nn.Module):
         q_pe, q_nope = torch.split(
             q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
         )
-
-        # Fused wk + weights_proj: one GEMM, then split
-        kw, _ = self.wk_weights_proj(hidden_states)
-        k = kw[:, : self.head_dim]
-        weights_raw = kw[:, self.head_dim :]
+        if self.is_fp4_ckpt:
+            # Fused wk + weights_proj: one GEMM, then split
+            kw, _ = self.wk_weights_proj(hidden_states)
+            k = kw[:, : self.head_dim]
+            weights = kw[:, self.head_dim :]
+        else:
+            k, _ = self.wk(hidden_states)
+            weights, _ = self.weights_proj(hidden_states)
 
         k = self.k_norm(k)
         k_pe, k_nope = torch.split(
@@ -726,7 +752,7 @@ class Indexer(nn.Module):
         q_scale = q_scale.view(-1, self.n_head, 1)
 
         weights = (
-            weights_raw.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
+            weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
         )
         weights = weights.squeeze(-1)
 
@@ -1314,6 +1340,10 @@ class DeepseekV2ForCausalLM(
         quant_config = vllm_config.quant_config
         self.config = config
         self.quant_config = quant_config
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
 
         qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
         qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
@@ -1439,12 +1469,13 @@ class DeepseekV2ForCausalLM(
             ("qkv_proj", "k_proj", "k"),
             ("qkv_proj", "v_proj", "v"),
         ]
-        # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
-        indexer_fused_mapping = [
-            ("wk_weights_proj", "wk", 0),
-            ("wk_weights_proj", "weights_proj", 1),
-        ]
-        stacked_params_mapping.extend(indexer_fused_mapping)
+        if self.is_fp4_ckpt:
+            # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
+            indexer_fused_mapping = [
+                ("wk_weights_proj", "wk", 0),
+                ("wk_weights_proj", "weights_proj", 1),
+            ]
+            stacked_params_mapping.extend(indexer_fused_mapping)
 
         if self.use_mha:
             stacked_params_mapping.extend(mha_params_mapping)

vllm/model_executor/models/nano_nemotron_vl.py

@@ -7,7 +7,6 @@
 #     LICENSE is in root directory.
 # --------------------------------------------------------
 
-import copy
 import math
 import warnings
 from collections.abc import Iterable, Mapping, Sequence
@@ -17,7 +16,7 @@ from typing import Annotated, Literal, TypeAlias
 
 import torch
 import torch.nn as nn
-from transformers import BatchFeature
+from transformers import BatchFeature, PretrainedConfig
 
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
@@ -210,11 +209,15 @@ class NanoNemotronVLProcessingInfo(BaseProcessingInfo):
 
     @cached_property
     def is_dynamic_tiler(self) -> bool:
-        return self.get_hf_processor().dynamic_tiler is not None
+        return BaseNanoNemotronVLProcessor.use_dynamic_resolution(self.get_hf_config())
 
-    @cached_property
+    @property
     def supports_video(self):
-        return self.get_hf_processor().supports_video
+        return True
+
+    @property
+    def supports_audio(self) -> bool:
+        return self.sound_config is not None
 
     def get_video_token(self) -> str | None:
         return IMG_CONTEXT
@@ -223,8 +226,8 @@ class NanoNemotronVLProcessingInfo(BaseProcessingInfo):
         return self.ctx.get_mm_config().video_pruning_rate
 
     @property
-    def audio_extractor(self) -> ParakeetExtractor | None:
-        return self.get_hf_processor().audio_extractor
+    def sound_config(self) -> PretrainedConfig | None:
+        return getattr(self.get_hf_config(), "sound_config", None)
 
     def get_default_tok_params(self) -> TokenizeParams:
         return super().get_default_tok_params().with_kwargs(add_special_tokens=False)
@@ -232,14 +235,14 @@ class NanoNemotronVLProcessingInfo(BaseProcessingInfo):
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         image_limit = {"image": None}
         video_limit = {"video": None} if self.supports_video else {}
-        audio_limit = {"audio": None} if self.audio_extractor is not None else {}
+        audio_limit = {"audio": None} if self.supports_audio else {}
         return {**image_limit, **video_limit, **audio_limit}
 
     def get_data_parser(self):
         target_sr = None
         target_channels = None
-        if extractor := self.audio_extractor:
-            target_sr = extractor.sampling_rate
+        if self.sound_config:
+            target_sr = self.sound_config.sampling_rate
             target_channels = 1
 
         return MultiModalDataParser(
@@ -371,7 +374,7 @@ class NanoNemotronVLMultiModalProcessor(
         fields = self._get_image_fields_config(hf_inputs)
         if self.info.supports_video:
             fields |= self._get_video_fields_config(hf_inputs)
-        if self.info.audio_extractor:
+        if self.info.supports_audio:
             fields |= self._get_audio_fields_config(hf_inputs)
 
         return fields
@@ -399,9 +402,8 @@ class NanoNemotronVLMultiModalProcessor(
 
             if isinstance(images, ImageEmbeddingItems):
                 feature_size = images.get_feature_size(item_idx)
-            elif tiler := hf_processor.dynamic_tiler:
-                image = images.get(item_idx)
-                feature_size = tiler.get_cached_feature_size(image)
+            elif self.info.is_dynamic_tiler:
+                feature_size = out_mm_data["num_tokens_per_image"][item_idx]
             else:
                 image_size = images.get_image_size(item_idx)
                 max_num_tiles = hf_processor.max_num_tiles
@@ -536,7 +538,7 @@ class NanoNemotronVLMultiModalProcessor(
             prompt_repls.append(
                 self._get_prompt_repl_video(mm_items, hf_processor, out_mm_data)
             )
-        if self.info.audio_extractor:
+        if self.info.supports_audio:
             prompt_repls.append(
                 self._get_prompt_repl_audio(mm_items, hf_processor, out_mm_data)
             )
@@ -772,12 +774,14 @@ class NanoNemotronVLDummyInputsBuilder(
         else:
             dummy_video = {}
 
-        if extractor := self.info.audio_extractor:
+        if sound_config := self.info.sound_config:
             num_audios = mm_counts.get("audio", 0)
             audio_overrides = mm_options.get("audio") if mm_options else None
             tokens_per_audio = max(1, seq_len // max(num_audios, 1))
-            max_audio_num_samples = MAX_AUDIO_LEN_S * extractor.sampling_rate
-            calculated_max_audio_num_samples = extractor.audio_length(tokens_per_audio)
+            max_audio_num_samples = MAX_AUDIO_LEN_S * sound_config.sampling_rate
+            calculated_max_audio_num_samples = ParakeetExtractor.audio_length(
+                sound_config, tokens_per_audio
+            )
             audio_len = min(max_audio_num_samples, calculated_max_audio_num_samples)
             dummy_audio = {
                 "audio": self._get_dummy_audios(
@@ -1029,9 +1033,13 @@ class NemotronH_Nano_VL_V2(
                 data=image_embeds,
             )
 
+        pixel_values_flat = kwargs.pop("pixel_values_flat", None)
+        if pixel_values_flat is None:
+            return None
+
         if self.dynamic_resolution:
             pixel_values_flat = DynamicResolutionImageTiler.stack(
-                kwargs.pop("pixel_values_flat"), self.patch_size
+                pixel_values_flat, self.patch_size
             )
             return NanoNemotronVLImagePixelInputsDynamic(
                 pixel_values_flat=pixel_values_flat, **kwargs
@@ -1497,15 +1505,13 @@ class NemotronH_Nano_VL_V2(
     @classmethod
     def get_mamba_state_shape_from_config(cls, vllm_config: "VllmConfig"):
         text_config = vllm_config.model_config.hf_config.text_config
-        temp_vllm_config = copy.deepcopy(vllm_config)
-        temp_vllm_config.model_config.hf_config = text_config
+        temp_vllm_config = vllm_config.with_hf_config(text_config)
         return NemotronHForCausalLM.get_mamba_state_shape_from_config(temp_vllm_config)
 
     @classmethod
     def get_mamba_state_dtype_from_config(cls, vllm_config: "VllmConfig"):
         text_config = vllm_config.model_config.hf_config.text_config
-        temp_vllm_config = copy.deepcopy(vllm_config)
-        temp_vllm_config.model_config.hf_config = text_config
+        temp_vllm_config = vllm_config.with_hf_config(text_config)
         return NemotronHForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config)
 
     @classmethod

vllm/model_executor/models/parakeet.py

@@ -159,5 +159,7 @@ class ParakeetExtractor(ParakeetFeatureExtractor):
         outputs["audio_num_clips"] = audio_num_clips
         return outputs
 
-    def audio_length(self, audio_tokens: int) -> int:
-        return int(audio_tokens * self.config.subsampling_factor * self.hop_length)
+    @staticmethod
+    def audio_length(raw_config: PretrainedConfig, audio_tokens: int) -> int:
+        config = ExtractorConfig.from_hf_config(raw_config)
+        return int(audio_tokens * config.subsampling_factor * config.hop_length)

vllm/model_executor/models/qwen3_5.py

@@ -620,6 +620,7 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
             self.packed_modules_mapping = {k: list(v) for k, v in base.items()}
             self.packed_modules_mapping.pop("in_proj_qkvz", None)
             self.packed_modules_mapping["in_proj_qkv"] = ["in_proj_qkv"]
+            self.packed_modules_mapping["in_proj_z"] = ["in_proj_z"]
 
     def embed_input_ids(
         self,

vllm/model_executor/models/radio.py

@@ -176,7 +176,6 @@ class ViTPatchGenerator(nn.Module):
                 temporal_patch_size=temporal_patch_size,
                 **factory,
             )
-            self._video_embedder_loaded = False
 
         if abs_pos:
             scale = embed_dim**-0.5
@@ -225,12 +224,7 @@ class ViTPatchGenerator(nn.Module):
         Returns:
             Embedded patches with temporal compression applied.
         """
-        if not self._video_embedder_loaded:
-            raise ValueError(
-                "Temporal compression (video_temporal_patch_size > 1) requires "
-                "video_embedder weights, but they were never loaded. "
-                "Ensure the checkpoint was trained with temporal compression."
-            )
+        assert self.temporal_patch_size > 1
         T = self.temporal_patch_size
         input_size = x.shape[2:]
 
@@ -794,9 +788,6 @@ class RadioModel(nn.Module):
                 weight_loader(param, weight)
                 loaded_params.add(vllm_key)
 
-        if "model.patch_generator.video_embedder.weight" in loaded_params:
-            self.model.patch_generator._video_embedder_loaded = True
-
         return loaded_params
 
     def _extract_final(

vllm/platforms/cpu.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import glob
-import json
 import os
 import platform
 import subprocess
@@ -11,11 +10,11 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 import psutil
-import regex as re
 import torch
 
 from vllm import envs
 from vllm.logger import init_logger
+from vllm.utils.ompmultiprocessing import OMPProcessManager
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
@@ -76,6 +75,10 @@ class CpuPlatform(Platform):
     dispatch_key: str = "CPU"
     dist_backend: str = "gloo"
     device_control_env_var = "CPU_VISIBLE_MEMORY_NODES"
+    omp_process_manager = None
+    smt = 1  # SMT level for OMP - 4 threads on PowerPC, 1 on others
+    global_cpu_mask = None
+    simulate_numa = int(os.environ.get("_SIM_MULTI_NUMA", 0))
 
     @property
     def supported_dtypes(self) -> list[torch.dtype]:
@@ -191,26 +194,10 @@ class CpuPlatform(Platform):
 
         cache_config.cpu_kvcache_space_bytes = CpuPlatform.get_device_total_memory()
 
-        # reserve at least one core for nixl_connector under p/d case
-        if vllm_config.kv_transfer_config and (
-            envs.VLLM_CPU_NUM_OF_RESERVED_CPU == 0
-            or envs.VLLM_CPU_NUM_OF_RESERVED_CPU is None
-        ):
-            os.environ["VLLM_CPU_NUM_OF_RESERVED_CPU"] = "1"
-
         parallel_config = vllm_config.parallel_config
-        if (
-            parallel_config.world_size > 1
-            and parallel_config.distributed_executor_backend is not None
-            and parallel_config.distributed_executor_backend != "mp"
-        ):
-            logger.warning(
-                (
-                    "%s is not supported on CPU, fallback to mp "
-                    "distributed executor backend."
-                ),
-                parallel_config.distributed_executor_backend,
-            )
+        # OMP requires the MP executor to function correctly, UniProc is not
+        # supported as it is not possible to set the OMP environment correctly
+        if parallel_config.distributed_executor_backend == "uni":
             parallel_config.distributed_executor_backend = "mp"
         if parallel_config.worker_cls == "auto":
             parallel_config.worker_cls = "vllm.v1.worker.cpu_worker.CPUWorker"
@@ -267,14 +254,6 @@ class CpuPlatform(Platform):
         # variable "NUMEXPR_MAX_THREADS" (64)'.
         os.environ["NUMEXPR_MAX_THREADS"] = str(get_max_threads())
 
-        if envs.VLLM_CPU_OMP_THREADS_BIND != "nobind":
-            # Set default threads num for OpenMP parallel
-            os.environ["OMP_NUM_THREADS"] = str(torch.get_num_threads())
-        else:
-            # In this case, setting the OpenMP configuration via
-            # OMP_NUM_THREADS is up to the user.
-            logger.info("Disabling binding processes to CPU cores...")
-
         # Disable torch async compiling which won't work with daemonic processes
         os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
 
@@ -286,8 +265,8 @@ class CpuPlatform(Platform):
 
         ld_preload_str = os.getenv("LD_PRELOAD", "")
 
-        # Intel OpenMP setting
-        if "libiomp5.so" in ld_preload_str:
+        # Intel and CLANG OpenMP setting
+        if "libiomp5.so" in ld_preload_str or "libomp5" in ld_preload_str:
             # The time(milliseconds) that a thread should wait after
             # completing the execution of a parallel region, before sleeping.
             os.environ["KMP_BLOCKTIME"] = "1"
@@ -324,37 +303,6 @@ class CpuPlatform(Platform):
                     ld_preload_str = tcmalloc_so
                 os.environ["LD_PRELOAD"] = ld_preload_str
 
-        if (
-            platform.system() == "Linux"
-            and cpu_architecture in (CpuArchEnum.ARM, CpuArchEnum.POWERPC)
-            and not ("libomp" in ld_preload_str or "libgomp" in ld_preload_str)
-        ):
-            # We need to LD_PRELOAD PyTorch's libgomp, otherwise only
-            # one core will be properly utilized when we thread-bind
-            # See: https://github.com/vllm-project/vllm/issues/27369
-            # TODO: Remove once:
-            # https://github.com/pytorch/pytorch/issues/166087 is fixed
-
-            # We need to find the location of PyTorch's libgomp
-            torch_pkg = os.path.dirname(torch.__file__)
-            site_root = os.path.dirname(torch_pkg)
-            # Search both torch.libs and torch/lib - See: https://github.com/vllm-project/vllm/issues/30470
-            torch_libs_paths = [
-                os.path.join(site_root, "torch.libs"),
-                os.path.join(torch_pkg, "lib"),
-            ]
-            pytorch_libgomp_so_candidates = []
-            for torch_libs in torch_libs_paths:
-                pytorch_libgomp_so_candidates.extend(
-                    glob.glob(os.path.join(torch_libs, "libgomp*.so*"))
-                )
-            if pytorch_libgomp_so_candidates:
-                pytorch_libgomp_so = pytorch_libgomp_so_candidates[0]
-                if ld_preload_str:
-                    ld_preload_str += ":"
-                ld_preload_str += pytorch_libgomp_so
-                os.environ["LD_PRELOAD"] = ld_preload_str
-
         os.environ["LOCAL_WORLD_SIZE"] = str(
             vllm_config.parallel_config.tensor_parallel_size
         )
@@ -369,6 +317,13 @@ class CpuPlatform(Platform):
                 vllm_config.model_config.max_model_len,
                 vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,
             )
+        # CI specific "quick" NUMA simulation - split all available CPUs
+        # into a fake NUMA topology
+        if os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", None) is not None:
+            os.environ["_SIM_MULTI_NUMA"] = str(
+                vllm_config.parallel_config.world_size
+                * vllm_config.parallel_config._api_process_count
+            )
 
     @classmethod
     def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
@@ -377,46 +332,76 @@ class CpuPlatform(Platform):
         pass
 
     @classmethod
-    def get_allowed_cpu_core_node_list(cls) -> tuple[list[int], list[LogicalCPUInfo]]:
-        assert platform.system() == "Linux"
+    def get_omp_manager(cls) -> OMPProcessManager:
+        # initialise the OMP resource management if need be and return the manager
+        if cls.omp_process_manager is None:
+            if cls.get_cpu_architecture() == CpuArchEnum.POWERPC:
+                cls.smt = 4
+            cls.omp_process_manager = OMPProcessManager(
+                affinity=cls.get_global_cpu_mask(), smt=cls.smt
+            )
+            # we need to fix up the topology returned by the OMP Manager for
+            # simulated NUMA environments in CI
+            if cls.simulate_numa > 0:
+                logger.info(
+                    "Adjusting numa topology to resemble at least %d nodes",
+                    int(cls.simulate_numa),
+                )
+                om = cls.omp_process_manager
+                while len(om.omp_places) < cls.simulate_numa:
+                    new_omp_places = []
+                    touched = False
+                    for omp_place in om.omp_places:
+                        if len(omp_place["mask"]) > 1:
+                            touched = True
+                            cpu_list = sorted(list(omp_place["mask"]))
+                            new_omp_places.append(
+                                {
+                                    "mask": set(cpu_list[0 : int(len(cpu_list) / 2)]),
+                                    "available": True,
+                                }
+                            )
+                            new_omp_places.append(
+                                {
+                                    "mask": set(cpu_list[int(len(cpu_list) / 2) :]),
+                                    "available": True,
+                                }
+                            )
+                    if touched:
+                        om.omp_places = new_omp_places
+                    else:
+                        raise ValueError(
+                            "Cannot split the existing NUMA topology to match "
+                            "simulation requirements"
+                        )
 
-        # Init LogicalCPUInfo from lscpu
-        lscpu_output = subprocess.check_output(
-            "lscpu -J -e=CPU,CORE,NODE", shell=True, text=True
+        return cls.omp_process_manager
+
+    @classmethod
+    def get_global_cpu_mask(cls) -> set[int]:
+        # get global cpu mask
+        if cls.global_cpu_mask is None:
+            if hasattr(os, "sched_getaffinity"):
+                cls.global_cpu_mask = os.sched_getaffinity(0)
+            else:
+                # macOS does not support sched_getaffinity
+                cpu_count = os.cpu_count() or 1
+                cls.global_cpu_mask = set(range(cpu_count))
+        return cls.global_cpu_mask
+
+    @classmethod
+    def reserve_cpus(cls, reserve: set[int]) -> bool:
+        # remove CPUs from global mask, for now there is no "release" mechanism
+        if cls.omp_process_manager is not None:
+            for place in cls.omp_process_manager.omp_places:
+                if not place["available"]:
+                    return False
+        cls.global_cpu_mask = cls.get_global_cpu_mask() - reserve
+        # reinitialize OMP resource management
+        cls.omp_process_manager = OMPProcessManager(
+            affinity=cls.global_cpu_mask, smt=cls.smt
         )
-        lscpu_output = re.sub(r'"node":\s*-\s*(,|\n)', r'"node": 0\1', lscpu_output)
-        logical_cpu_list: list[LogicalCPUInfo] = json.loads(
-            lscpu_output, object_hook=LogicalCPUInfo.json_decoder
-        )["cpus"]
-
-        # Filter CPUs with invalid attributes
-        logical_cpu_list = [
-            x
-            for x in logical_cpu_list
-            if -1 not in (x.id, x.physical_core, x.numa_node)
-        ]
-
-        # Filter allowed CPUs
-        if hasattr(os, "sched_getaffinity"):
-            allowed_cpu_id_list = os.sched_getaffinity(0)
-        else:
-            raise NotImplementedError("Unsupported OS")
-        logical_cpu_list = [x for x in logical_cpu_list if x.id in allowed_cpu_id_list]
-
-        # Get allowed NUMA nodes
-        allowed_numa_nodes = set()
-        for x in logical_cpu_list:
-            allowed_numa_nodes.add(x.numa_node)  # type: ignore
-        allowed_numa_nodes_list = sorted(allowed_numa_nodes)
-
-        env_key = CpuPlatform.device_control_env_var
-        if env_key in os.environ and os.environ[env_key] != "":
-            visible_nodes = [int(s) for s in os.environ[env_key].split(",")]
-            allowed_numa_nodes_list = [
-                x for x in sorted(list(set(visible_nodes))) if x in allowed_numa_nodes
-            ]
-
-        return allowed_numa_nodes_list, logical_cpu_list
+        return True
 
     @classmethod
     def discover_numa_topology(cls) -> list[list[int]]:

vllm/platforms/interface.py

@@ -207,6 +207,15 @@ class Platform:
         """
         return cls.simple_compile_backend
 
+    @classmethod
+    def import_ir_kernels(cls) -> None:
+        """
+        The default implementation imports ``vllm.kernels``, which registers
+        the built-in IR op implementations. Out-of-tree (OOT) platforms should
+        override this method to import their own kernel modules.
+        """
+        import vllm.kernels  # noqa: F401
+
     @classmethod
     def device_id_to_physical_device_id(cls, device_id: int):
         # Treat empty device control env var as unset. This is a valid

vllm/platforms/rocm.py

@@ -402,7 +402,6 @@ class RocmPlatform(Platform):
         "gguf",
         "quark",
         "mxfp4",
-        "petit_nvfp4",
         "torchao",
         "bitsandbytes",
     ]

vllm/transformers_utils/configs/parakeet.py

@@ -44,15 +44,19 @@ class ExtractorConfig:
     subsampling_factor: int
     subsampling_conv_kernel_size: int
     subsampling_conv_stride: int
+    hop_length: int = 160
+    """Default `160`: Matches HF default"""
     clip_duration_s: int = 30
     clip_min_duration_s: float = 0.1
 
     @staticmethod
     def from_hf_config(config: PretrainedConfig) -> "ExtractorConfig":
         assert isinstance(config, PretrainedConfig)
+        hop_length = int(getattr(config, "hop_length", ExtractorConfig.hop_length))
         return ExtractorConfig(
             feature_size=config.num_mel_bins,
             sampling_rate=config.sampling_rate,
+            hop_length=hop_length,
             subsampling_factor=config.subsampling_factor,
             subsampling_conv_kernel_size=config.subsampling_conv_kernel_size,
             subsampling_conv_stride=config.subsampling_conv_stride,

vllm/transformers_utils/processors/nano_nemotron_vl.py

@@ -356,15 +356,6 @@ class DynamicResolutionImageTiler:
                 feature_sizes.append(param.num_embeddings)
         return images, feature_sizes
 
-    feature_size_cache: dict[Image.Image, int] = {}
-
-    @classmethod
-    def get_cached_feature_size(cls, image: Image.Image) -> int:
-        feature_size = cls.feature_size_cache[id(image)]
-        # hard assert that we only use the feature size once
-        del cls.feature_size_cache[id(image)]
-        return feature_size
-
     @dataclass
     class DynamicResolutionParams:
         media: Image.Image
@@ -519,7 +510,6 @@ class DynamicResolutionImageTiler:
                 param, token_count = self.process_media(media, tokens_for_media)
                 params.append(param)
                 token_counts.append(token_count)
-                self.feature_size_cache[id(param.media)] = param.num_embeddings
 
             # Step 2: Check if total tokens is within budget
             total_tokens = sum(token_counts)
@@ -857,13 +847,12 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
 
     @property
     def supports_video(self) -> bool:
-        return self.video_token_id is not None
+        return True
 
     @property
-    def video_token_id(self) -> int | None:
-        if self.video_token is None:
-            return None
-        return self.tokenizer.get_vocab().get(self.video_token, None)
+    def video_token_id(self) -> int:
+        assert self.video_token is not None
+        return self.tokenizer.get_vocab()[self.video_token]
 
     @property
     def image_token_id(self) -> int:
@@ -1055,6 +1044,13 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
         text_inputs = self.tokenizer(text, add_special_tokens=False)
 
         combined_inputs = {**text_inputs, **video_inputs, **audio_inputs}
+        frames_indices = combined_inputs.get("frames_indices")
+        ragged_frames_indices = (
+            isinstance(frames_indices, list)
+            and len({len(frame_indices) for frame_indices in frames_indices}) > 1
+        )
+        if ragged_frames_indices:
+            combined_inputs.pop("frames_indices")
 
         if self.dynamic_tiler is None:
             batch = BatchFeature(
@@ -1066,6 +1062,12 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
             # allow images to be exempt from the BatchFeature validation:
             # We will .stack() them in _parse_and_validate_image_input
             batch.update(image_inputs)
+        if ragged_frames_indices:
+            assert isinstance(frames_indices, list)
+            batch["frames_indices"] = [
+                torch.as_tensor(frame_indices, dtype=torch.int64)
+                for frame_indices in frames_indices
+            ]
         return batch
 
     def get_image_repl(

vllm/utils/ompmultiprocessing.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OMP Aware Multiprocessing manager for running multiprocessing.Process()
+Copyright (c) 2026 Red Hat Inc
+Copyright (c) 2026 Cambridge Greys Ltd
+"""
+
+import json
+import os
+import platform
+import subprocess
+
+
+def _int(arg):
+    """Relaxed parsing of ints which handles a - instead of a number.
+    The lscpu json may contain that for nodes in some cases. If that
+    is the case we parse it to zero
+    """
+    try:
+        if int(arg) >= 0:
+            return int(arg)
+    except ValueError:
+        pass
+    return 0
+
+
+def parse_mask(mask):
+    """Expand a X-Y,Z list"""
+    result = []
+    for token in mask.split(","):
+        try:
+            start, finish = token.split("-")
+            if int(start) > int(finish):
+                raise IndexError("Invalid Indexes for cpu ranges")
+            for cpu in range(int(start), int(finish) + 1):
+                result.append(cpu)
+        except ValueError:
+            result.append(int(token))
+    return set(result)
+
+
+def _get_default_affinity() -> set[int]:
+    """Get the set of CPUs the process is allowed to run on."""
+    if hasattr(os, "sched_getaffinity"):
+        return os.sched_getaffinity(0)
+    # macOS does not support sched_getaffinity; fall back to cpu_count
+    cpu_count = os.cpu_count() or 1
+    return set(range(cpu_count))
+
+
+def _get_cpu_topology_json() -> bytes:
+    """Get CPU topology as JSON.
+
+    On Linux this uses ``lscpu -Je``.  On other platforms (e.g. macOS) we
+    synthesize a simple topology where every logical CPU is its own core
+    on NUMA node 0, which is sufficient for the OMP place-list builder.
+    """
+    if platform.system() == "Linux":
+        return subprocess.run(["lscpu", "-Je"], check=True, capture_output=True).stdout
+
+    # Fallback for non-Linux (macOS, etc.)
+    cpu_count = os.cpu_count() or 1
+    cpus = []
+    for i in range(cpu_count):
+        cpus.append({"cpu": str(i), "core": str(i), "node": "0"})
+    return json.dumps({"cpus": cpus}).encode()
+
+
+def enumerate_resources(resource_map, mask=None, allowed=None):
+    """Enumerate system resources"""
+    if allowed is None:
+        allowed = _get_default_affinity()
+    if mask is not None:
+        allowed = allowed & mask
+
+    try:
+        allowed_nodes = parse_mask(os.environ["CPU_VISIBLE_MEMORY_NODES"])
+    except KeyError:
+        allowed_nodes = None
+
+    lscpu: dict[str, dict] = {"cpus": {}, "cores": {}, "nodes": {}}
+    for cpu in resource_map["cpus"]:
+        cpunum = int(cpu["cpu"])
+        if (
+            cpunum in allowed
+            and cpunum >= 0
+            and (allowed_nodes is None or _int(cpu["node"]) in allowed_nodes)
+        ):
+            lscpu["cpus"][cpunum] = [cpu]
+            core = _int(cpu["core"])
+            if lscpu["cores"].get(core, None) is None:
+                lscpu["cores"][core] = [cpu]
+            else:
+                lscpu["cores"][core].append(cpu)
+            node = _int(cpu["node"])
+            if lscpu["nodes"].get(node, None) is None:
+                lscpu["nodes"][node] = [cpu]
+            else:
+                lscpu["nodes"][node].append(cpu)
+    return lscpu
+
+
+def produce_cpu_list(cpus, smt=1):
+    """Produce a CPU list with/without SMT pairs - main cpu list case"""
+    mask: list[int] = []
+    for key, value in cpus.items():
+        exists = 0
+        for cpu in mask:
+            if cpu == value[0]["core"]:
+                exists += 1
+                break
+        if exists < smt:
+            mask.append(int(key))
+    return {"mask": set(mask), "available": True}
+
+
+def produce_cpu_sublist(scpus, smt=1):
+    """Produce a CPU list with/without SMT pairs - resource leaf case"""
+    cpu_list: list[dict] = []
+    for value in scpus:
+        exists = 0
+        for cpu in cpu_list:
+            if int(cpu["core"]) == int(value["core"]):
+                exists += 1
+                break
+        if exists < smt:
+            cpu_list.append(value)
+    mask = []
+    for cpu in cpu_list:
+        mask.append(int(cpu["cpu"]))
+
+    return {"mask": set(mask), "available": True}
+
+
+def create_omp_places(resources, strategy, smt=True):
+    """Parse CPU topology and generate possible CPU masks"""
+    omp_places = []
+    if strategy == "all":
+        omp_places.append(produce_cpu_list(resources["cpus"], smt))
+    elif strategy == "cores":
+        for value in resources["cores"].values():
+            omp_places.append(produce_cpu_sublist(value, smt))
+    elif strategy == "nodes":
+        for value in resources["nodes"].values():
+            omp_places.append(produce_cpu_sublist(value, smt))
+    else:
+        raise NotImplementedError("Unknown strategy")
+
+    return omp_places
+
+
+# pylint: disable=too-few-public-methods
+class OMPProcessManager:
+    """OMP aware wrapper to run mp Process()"""
+
+    def __init__(self, strategy="nodes", smt=1, mock=None, affinity=None):
+        self.strategy = strategy
+        self.smt = smt
+        self.omp_places = []
+        vllm_mask = os.environ.get("VLLM_CPU_OMP_THREADS_BIND", None)
+        self.setup_omp = vllm_mask != "nobind"
+        if self.setup_omp:
+            omp_places = []
+            if vllm_mask is not None:
+                masks = []
+                for spec in vllm_mask.split("|"):
+                    masks.append(parse_mask(spec))
+            else:
+                masks = [None]
+            if mock is None:
+                data = _get_cpu_topology_json()
+            else:
+                with open(mock, mode="rb") as jf:
+                    data = jf.read()
+            lscpu = json.loads(data)
+            for mask in masks:
+                resources = enumerate_resources(lscpu, mask, affinity)
+                omp_places.extend(create_omp_places(resources, strategy, smt))
+            self.omp_places = sorted(
+                omp_places,
+                key=lambda p: "{:04d}-{:04d}".format(len(p["mask"]), max(p["mask"])),
+                reverse=True,
+            )
+
+    def run(self, what, *args, **kwargs):
+        """Run arg with correct OMP environment"""
+        if self.setup_omp:
+            for place in self.omp_places:
+                if place["available"]:
+                    reserve = int(os.environ.get("VLLM_CPU_NUM_OF_RESERVED_CPU", 0))
+                    place["available"] = False
+                    # pylint: disable=consider-using-f-string
+                    os.environ["OMP_PLACES"] = "{}".format(place["mask"])
+                    os.environ["OMP_NUM_THREADS"] = "{}".format(
+                        len(place["mask"]) - reserve
+                    )
+                    os.environ["OMP_PROC_BIND"] = "TRUE"
+                    return what(*args, **kwargs)
+            raise IndexError("Out of OMP places")
+        return what(*args, **kwargs)

vllm/v1/attention/backends/gdn_attn.py

@@ -63,6 +63,10 @@ class GDNAttentionMetadata:
 
     num_accepted_tokens: torch.Tensor | None = None  # shape: [batch,]
 
+    # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices)
+    chunk_indices: torch.Tensor | None = None
+    chunk_offsets: torch.Tensor | None = None
+
     # The following attributes are for triton implementation of causal_conv1d
     nums_dict: dict | None = None
     batch_ptr: torch.Tensor | None = None
@@ -305,6 +309,26 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
             assert num_accepted_tokens is not None
             num_accepted_tokens = num_accepted_tokens[spec_sequence_masks]
 
+        chunk_indices: torch.Tensor | None = None
+        chunk_offsets: torch.Tensor | None = None
+        if num_prefills > 0:
+            # Only prefill batches use FLA chunk ops.
+            # Pre-compute on CPU and async-copy to GPU to avoid
+            # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
+            from vllm.model_executor.layers.fla.ops.index import (
+                prepare_chunk_indices,
+                prepare_chunk_offsets,
+            )
+            from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE
+
+            gpu_device = query_start_loc.device
+            chunk_indices = prepare_chunk_indices(
+                non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
+            ).to(device=gpu_device, non_blocking=True)
+            chunk_offsets = prepare_chunk_offsets(
+                non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE
+            ).to(device=gpu_device, non_blocking=True)
+
         if num_prefills > 0:
             has_initial_state = context_lens_tensor > 0
             if spec_sequence_masks is not None:
@@ -405,6 +429,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
             num_spec_decode_tokens=num_spec_decode_tokens,
             num_actual_tokens=m.num_actual_tokens,
             has_initial_state=has_initial_state,
+            chunk_indices=chunk_indices,
+            chunk_offsets=chunk_offsets,
             spec_query_start_loc=spec_query_start_loc,
             non_spec_query_start_loc=non_spec_query_start_loc,
             spec_state_indices_tensor=spec_state_indices_tensor,

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

@@ -129,9 +129,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
 
         from aiter import dtypes, get_mla_metadata_info_v1
 
-        self._num_attention_heads = vllm_config.model_config.get_num_attention_heads(
-            vllm_config.parallel_config
-        )
+        # For num_attention_heads < 16 (e.g. kimi-k2.5 head=8 with TP8),
+        # make sure get_mla_metadata_info_v1 / get_mla_metadata_v1 are consistent
+        # with the actual tensor shape passed to mla_decode_fwd.
+        self._num_attention_heads = max(16, self.num_heads)
         q_dtype = self.decode_attn_out_dtype
         kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
         if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):

vllm/v1/executor/multiproc_executor.py

@@ -119,7 +119,6 @@ class MultiprocExecutor(Executor):
             f"_parallel_size ({pcp_size}). "
         )
 
-        # Set multiprocessing envs
         set_multiprocessing_worker_envs()
 
         # use the loopback address get_loopback_ip() for communication.
@@ -172,16 +171,31 @@ class MultiprocExecutor(Executor):
             for local_rank in range(self.local_world_size):
                 global_rank = global_start_rank + local_rank
                 is_driver_worker = self._is_driver_worker(global_rank)
-                unready_worker_handle = WorkerProc.make_worker_process(
-                    vllm_config=self.vllm_config,
-                    local_rank=local_rank,
-                    rank=global_rank,
-                    distributed_init_method=distributed_init_method,
-                    input_shm_handle=scheduler_output_handle,
-                    shared_worker_lock=shared_worker_lock,
-                    is_driver_worker=is_driver_worker,
-                    inherited_fds=inherited_fds,
-                )
+                if current_platform.is_cpu():
+                    om = current_platform.get_omp_manager()
+                    logger.info("Configured OMP PLACES %s", str(om.omp_places))
+                    unready_worker_handle = om.run(
+                        WorkerProc.make_worker_process,
+                        vllm_config=self.vllm_config,
+                        local_rank=local_rank,
+                        rank=global_rank,
+                        distributed_init_method=distributed_init_method,
+                        input_shm_handle=scheduler_output_handle,
+                        shared_worker_lock=shared_worker_lock,
+                        is_driver_worker=is_driver_worker,
+                        inherited_fds=inherited_fds,
+                    )
+                else:
+                    unready_worker_handle = WorkerProc.make_worker_process(
+                        vllm_config=self.vllm_config,
+                        local_rank=local_rank,
+                        rank=global_rank,
+                        distributed_init_method=distributed_init_method,
+                        input_shm_handle=scheduler_output_handle,
+                        shared_worker_lock=shared_worker_lock,
+                        is_driver_worker=is_driver_worker,
+                        inherited_fds=inherited_fds,
+                    )
                 unready_workers.append(unready_worker_handle)
                 if inherited_fds is not None:
                     inherited_fds.append(unready_worker_handle.death_writer.fileno())
@@ -1000,24 +1014,26 @@ def set_multiprocessing_worker_envs():
 
     _maybe_force_spawn()
 
-    # Configure thread parallelism if OMP_NUM_THREADS isn't set
-    #
-    # Helps to avoid CPU contention. The default of spawning a thread per
-    # core combined with multiprocessing for each GPU can have a negative
-    # impact on performance. The contention is amplified when running in a
-    # container where CPU limits can cause throttling.
-    default_omp_num_threads = 1
-    if (
-        "OMP_NUM_THREADS" not in os.environ
-        and (current_parallelism := torch.get_num_threads()) > default_omp_num_threads
-    ):
-        logger.warning_once(
-            "Reducing Torch parallelism from %d threads to %d to avoid "
-            "unnecessary CPU contention. Set OMP_NUM_THREADS in the "
-            "external environment to tune this value as needed.",
-            current_parallelism,
-            default_omp_num_threads,
-            scope="local",
-        )
-        os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads)
-        torch.set_num_threads(default_omp_num_threads)
+    if not current_platform.is_cpu():
+        # Configure thread parallelism if OMP_NUM_THREADS isn't set
+        #
+        # Helps to avoid CPU contention. The default of spawning a thread per
+        # core combined with multiprocessing for each GPU can have a negative
+        # impact on performance. The contention is amplified when running in a
+        # container where CPU limits can cause throttling.
+        default_omp_num_threads = 1
+        if (
+            "OMP_NUM_THREADS" not in os.environ
+            and (current_parallelism := torch.get_num_threads())
+            > default_omp_num_threads
+        ):
+            logger.warning_once(
+                "Reducing Torch parallelism from %d threads to %d to avoid "
+                "unnecessary CPU contention. Set OMP_NUM_THREADS in the "
+                "external environment to tune this value as needed.",
+                current_parallelism,
+                default_omp_num_threads,
+                scope="local",
+            )
+            os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads)
+            torch.set_num_threads(default_omp_num_threads)

vllm/v1/worker/cpu_worker.py

@@ -1,18 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
-import platform
 import sys
-from collections.abc import Callable
 from typing import Any
 
 import torch
 
-from vllm import envs
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.platforms import CpuArchEnum, current_platform
-from vllm.platforms.cpu import CpuPlatform, LogicalCPUInfo
 from vllm.profiler.wrapper import TorchProfilerWrapper
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.cpu_model_runner import CPUModelRunner
@@ -71,44 +67,6 @@ class CPUWorker(Worker):
             if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
                 check_preloaded_libs("libiomp")
 
-        # Setup OpenMP threads affinity.
-        omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND
-        # Under numa binding some cores reserved for kv transfer in nixl_connector.py
-        if omp_cpuids == "auto" and platform.system() == "Linux":
-            cpu_arch = current_platform.get_cpu_architecture()
-            if cpu_arch in (CpuArchEnum.POWERPC, CpuArchEnum.S390X):
-                # For S390X/POWERPC SMT-8/4/2
-                self.local_omp_cpuid = self._get_autobind_cpu_ids(
-                    lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4]
-                )
-            elif cpu_arch == CpuArchEnum.X86:
-                # For x86 SMT-2, use 1 CPU per core
-                self.local_omp_cpuid = self._get_autobind_cpu_ids(
-                    lambda cpus: cpus[-1:]
-                )
-            elif cpu_arch == CpuArchEnum.ARM:
-                # For AArch64, no SMT
-                self.local_omp_cpuid = self._get_autobind_cpu_ids(lambda cpus: cpus)
-            else:
-                self.local_omp_cpuid = "nobind"
-        elif omp_cpuids == "nobind":
-            self.local_omp_cpuid = "nobind"
-        else:
-            local_dp_rank = self.parallel_config.data_parallel_rank_local
-            omp_cpuids_list = omp_cpuids.split("|")
-            if local_dp_rank is not None:
-                world_size = self.parallel_config.world_size
-                omp_cpuids_list = omp_cpuids_list[
-                    local_dp_rank * world_size : (local_dp_rank + 1) * world_size
-                ]
-            self.local_omp_cpuid = omp_cpuids_list[self.rank]
-
-        if self.local_omp_cpuid != "nobind":
-            ret = torch.ops._C.init_cpu_threads_env(self.local_omp_cpuid)
-            if ret:
-                logger.info(ret)
-
-        # After the thread binding, changing thread num is not allowed
         def skip_set_num_threads(x: int):
             logger.warning(
                 "CPU backend doesn't allow to use "
@@ -153,92 +111,6 @@ class CPUWorker(Worker):
         self.model_runner.warming_up_model()
         return self.compilation_config.compilation_time
 
-    def _get_autobind_cpu_ids(
-        self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]]
-    ) -> str:
-        """
-        Return CPU ids to bind based on NUMA nodes.
-        Currently for rank N, only CPU ids on the N-th node in available NUMA
-        node list will be selected.
-        Args:
-            cpu_selector: a callable object to select CPUs from a CPU list
-            of a physical core. The input is a LogicalCPUInfo list, sorted by
-            the LogicalCPUInfo.id. A selected LogicalCPUInfo list should be
-            returned.
-        """
-        # simulate multiple numa nodes, for testing
-        sim_multi_numa_nodes = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0"
-
-        allowed_numa_nodes, logical_cpu_list = (
-            CpuPlatform.get_allowed_cpu_core_node_list()
-        )
-        local_world_size = self.parallel_config.local_world_size
-        assert len(allowed_numa_nodes) >= local_world_size or sim_multi_numa_nodes, (
-            f"Not enough allowed NUMA nodes to bind threads of "
-            f"{local_world_size} local CPUWorkers. "
-            f"Allowed NUMA nodes are {allowed_numa_nodes}. "
-            "Please try to bind threads manually."
-        )
-
-        if not sim_multi_numa_nodes:
-            # Get CPUs on NUMA node `allowed_numa_nodes[local_rank]`
-            selected_numa_node = allowed_numa_nodes[self.local_rank]  # type: ignore
-            logical_cpu_list = [
-                x for x in logical_cpu_list if x.numa_node == selected_numa_node
-            ]
-        else:
-            # This is a bit tricky because the internal DP size
-            # is always 1 for non-MoE models
-            world_size_across_dp = (
-                self.parallel_config.world_size
-                * self.parallel_config._api_process_count
-            )
-            assert len(logical_cpu_list) >= world_size_across_dp
-            logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.numa_node)
-            sim_cpu_num_per_node = len(logical_cpu_list) // world_size_across_dp
-            assert self.parallel_config.data_parallel_rank_local is not None
-            start_idx = (
-                self.local_rank
-                + self.parallel_config.world_size
-                * self.parallel_config.data_parallel_rank_local
-            ) * sim_cpu_num_per_node
-            logical_cpu_list = logical_cpu_list[
-                start_idx : (start_idx + sim_cpu_num_per_node)
-            ]
-
-        # Select CPUs from each physical core via cpu_selector
-        core_to_cpus: dict[int, list[LogicalCPUInfo]] = {}
-        for cpu_info in logical_cpu_list:
-            if cpu_info.physical_core not in core_to_cpus:
-                core_to_cpus[cpu_info.physical_core] = []
-            core_to_cpus[cpu_info.physical_core].append(cpu_info)
-        logical_cpu_list = []
-        for cpu_list in core_to_cpus.values():
-            cpu_list = sorted(cpu_list, key=lambda x: x.id)
-            logical_cpu_list.extend(cpu_selector(cpu_list))
-        logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.id)
-
-        # Reserve CPUs for other processes
-        reserve_cpu_num = envs.VLLM_CPU_NUM_OF_RESERVED_CPU
-        if reserve_cpu_num is None:
-            need_reserve = (
-                self.parallel_config.world_size > 1
-                or self.parallel_config.data_parallel_size_local > 1
-            )
-            reserve_cpu_num = 1 if need_reserve else 0
-        assert len(logical_cpu_list) > reserve_cpu_num, (
-            f"VLLM_CPU_NUM_OF_RESERVED_CPU ({reserve_cpu_num}) "
-            f"should less than {len(logical_cpu_list)}."
-        )
-        if reserve_cpu_num != 0:
-            logical_cpu_list = logical_cpu_list[:-reserve_cpu_num]
-
-        logger.info(
-            "auto thread-binding list (id, physical core): %s",
-            [(x.id, x.physical_core) for x in logical_cpu_list],
-        )
-        return ",".join([str(x.id) for x in logical_cpu_list])
-
     def profile(self, is_start: bool = True, profile_prefix: str | None = None):
         if self.profiler is None:
             raise RuntimeError("Profiler is not enabled.")

vllm/v1/worker/gpu_model_runner.py

@@ -4192,6 +4192,7 @@ class GPUModelRunner(
         spec_config = self.speculative_config
         propose_drafts_after_bookkeeping = False
         if spec_config is not None:
+            # Decide whether to run the drafter or zero out draft tokens.
             input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
                 spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
                 <= self.effective_drafter_max_model_len
@@ -4227,10 +4228,6 @@ class GPUModelRunner(
                     self._copy_valid_sampled_token_count(
                         next_token_ids, valid_sampled_tokens_count
                     )
-                    self._draft_token_ids = torch.zeros(
-                        1, device=self.device, dtype=torch.int32
-                    ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
-                    self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
             elif (
                 spec_config.use_ngram_gpu()
                 and not spec_config.disable_padded_drafter_batch
@@ -4253,15 +4250,20 @@ class GPUModelRunner(
                     self._copy_valid_sampled_token_count(
                         next_token_ids, valid_sampled_tokens_count
                     )
-                    # Since we couldn't run the drafter,
-                    # just use zeros for the draft tokens.
-                    self._draft_token_ids = torch.zeros(
-                        1, device=self.device, dtype=torch.int32
-                    ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
-                    self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
             else:
                 propose_drafts_after_bookkeeping = input_fits_in_drafter
 
+            if not input_fits_in_drafter:
+                # Zero out draft tokens so the scheduler doesn't schedule
+                # stale drafts from the previous step.
+                # For Nemotron-H: it is necessary to zero out the draft tokens,
+                # otherwise the stale tokens will corrupt Mamba recurrent
+                # state and logprobs for sequences near max_model_len.
+                self._draft_token_ids = torch.zeros(
+                    1, device=self.device, dtype=torch.int32
+                ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
+                self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
+
         with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
             (
                 num_nans_in_logits,
@@ -6844,7 +6846,7 @@ class GPUModelRunner(
             # group
             self.drafter.validate_same_kv_cache_group(kv_cache_config)
 
-        if has_kv_transfer_group():
+        if has_kv_transfer_group() and not is_profiling:
             kv_transfer_group = get_kv_transfer_group()
             if self.cross_layers_kv_cache is not None:
                 assert self.cross_layers_attn_backend is not None

vllm/v1/worker/workspace.py

@@ -31,7 +31,7 @@ _manager: "WorkspaceManager | None" = None
 class WorkspaceManager:
     """Manager for workspace allocation.
 
-    Manages workspace buffers for DBO (Dual Batch Overlap) execution.
+    Manages one workspace buffer per active ubatch slot.
     Can be locked to prevent further growth during execution.
     """
 
@@ -39,7 +39,9 @@ class WorkspaceManager:
         self._device = device
         # Cache num ubatches at init based on configuration (default to 1)
         self._num_ubatches = num_ubatches if num_ubatches is not None else 1
-        self._current_workspaces: list[torch.Tensor | None] = [None, None]
+        self._current_workspaces: list[torch.Tensor | None] = [
+            None
+        ] * self._num_ubatches
         self._locked: bool = False
 
     @staticmethod
@@ -224,7 +226,7 @@ def init_workspace_manager(
 
     Args:
         device: The device to allocate workspace on.
-        num_ubatches: Number of micro-batches. Defaults to 1.
+        num_ubatches: Number of workspace ubatch slots. Defaults to 1.
     """
     global _manager
     if _manager is not None: