Merge branch 'main' into wentao-fix-flashinfer-layout

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

.buildkite/hardware_tests/cpu.yaml

@@ -54,6 +54,20 @@ steps:
       pytest -x -v -s tests/models/language/generation -m cpu_model
       pytest -x -v -s tests/models/language/pooling -m cpu_model"
 
+- label: CPU-ModelRunnerV2 Tests
+  depends_on: []
+  device: intel_cpu
+  no_plugin: true
+  soft_fail: true
+  source_file_dependencies:
+  - vllm/v1/worker/cpu/
+  - vllm/v1/worker/gpu/
+  commands:
+    - |
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
+      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
+
 - label: CPU-Quantization Model Tests
   depends_on: []
   device: intel_cpu

.buildkite/test-amd.yaml

@@ -2703,19 +2703,35 @@ steps:
   optional: true
   working_dir: "/vllm-workspace/"
   source_file_dependencies:
+  - csrc/custom_quickreduce.cu
+  - csrc/ops.h
+  - csrc/torch_bindings.cpp
   - vllm/distributed/
-  - vllm/v1/distributed/
+  - vllm/model_executor/layers/
+  - vllm/entrypoints/llm.py
+  - vllm/config/parallel.py
   - vllm/model_executor/layers/fused_moe/
+  - vllm/v1/engine/
+  - vllm/v1/executor/
+  - vllm/v1/worker/
+  - vllm/v1/distributed/
   - vllm/v1/attention/backends/
   - vllm/v1/attention/selector.py
-  - tests/distributed/test_context_parallel.py
-  - tests/v1/distributed/test_dbo.py
-  - examples/features/data_parallel/data_parallel_offline.py
   - vllm/_aiter_ops.py
+  - vllm/_custom_ops.py
   - vllm/platforms/rocm.py
+  - vllm/envs.py
+  - examples/offline_inference/data_parallel.py
+  - tests/distributed/test_context_parallel.py
+  - tests/distributed/test_rocm_quick_reduce.py
+  - tests/distributed/test_quick_all_reduce.py
+  - tests/v1/distributed/test_dbo.py
+  - tests/utils.py
   commands:
   - pytest -v -s tests/distributed/test_context_parallel.py
   - pytest -v -s tests/v1/distributed/test_dbo.py
+  - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
+  - pytest -v -s tests/distributed/test_quick_all_reduce.py
 
 #--------------------------------------------------------  mi355 · entrypoints  --------------------------------------------------------#
 

AGENTS.md

@@ -101,6 +101,8 @@ pre-commit run ruff-check --all-files
 pre-commit run mypy-3.10 --all-files --hook-stage manual
 ```
 
+The line length limit for Python code is 88 characters. If you are not sure, use pre-commit to check.
+
 ### Commit messages
 
 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:

csrc/cpu/cpu_attn_impl.hpp

@@ -408,9 +408,19 @@ class AttentionScheduler {
     const int64_t cache_size = cpu_utils::get_available_l2_size();
     const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
     const int32_t kv_len_alignment = input.kv_block_alignment;
+    bool has_decode_request = false;
+    bool decode_only_batch = true;
+    for (int32_t req_id = 0; req_id < input.num_reqs; ++req_id) {
+      const int32_t q_token_num =
+          input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
+      has_decode_request = has_decode_request || (q_token_num == 1);
+      decode_only_batch = decode_only_batch && (q_token_num == 1);
+    }
     int32_t q_head_per_kv = input.num_heads_q / input.num_heads_kv;
-    const bool use_gqa = (max_num_q_per_iter % q_head_per_kv == 0);
-    if (!use_gqa) {
+    const bool supports_gqa = q_head_per_kv <= max_num_q_per_iter;
+    const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
+    const bool use_gqa_scratchpad = supports_gqa && has_decode_request;
+    if (!use_gqa_scratchpad) {
       q_head_per_kv = 1;  // fallback to MHA
     }
     const int32_t min_split_kv_len =
@@ -680,7 +690,7 @@ class AttentionScheduler {
         metadata_ptr->attention_scratchpad_size_per_thread *
             metadata_ptr->thread_num +
         metadata_ptr->reduction_scratchpad_size_per_kv_head *
-            (use_gqa ? input.num_heads_kv : input.num_heads_q);
+            (use_gqa_fast_path ? input.num_heads_kv : input.num_heads_q);
     cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(
         scratchpad_size);
 
@@ -1409,13 +1419,24 @@ class AttentionMainLoop {
       const int32_t q_head_num = input->num_heads;
       const int32_t kv_head_num = input->num_kv_heads;
       const int32_t q_heads_per_kv = q_head_num / kv_head_num;
-      const bool use_gqa =
-          (max_q_head_num_per_iter % q_heads_per_kv == 0) ? true : false;
-      const int32_t actual_kv_head_num = use_gqa ? kv_head_num : q_head_num;
-      const int32_t actual_q_heads_per_kv = use_gqa ? q_heads_per_kv : 1;
+      AttentionWorkItemGroup* const workitem_groups =
+          metadata.workitem_groups_ptr;
+      const int32_t* cu_workitem_num_per_thread =
+          metadata.cu_workitem_num_per_thread;
+      ReductionWorkItemGroup* const reduction_items =
+          metadata.reduction_items_ptr;
+      const bool supports_gqa = q_heads_per_kv <= max_q_head_num_per_iter;
+      bool decode_only_batch = true;
+      for (int32_t i = 0; i < metadata.workitem_group_num; ++i) {
+        decode_only_batch =
+            decode_only_batch && (workitem_groups[i].q_token_num == 1);
+      }
+      const bool use_gqa_fast_path = supports_gqa && decode_only_batch;
+      const int32_t actual_kv_head_num =
+          use_gqa_fast_path ? kv_head_num : q_head_num;
+      const int32_t actual_q_heads_per_kv =
+          use_gqa_fast_path ? q_heads_per_kv : 1;
       TORCH_CHECK_LE(actual_q_heads_per_kv, max_q_head_num_per_iter);
-      const int32_t max_q_token_num_per_iter =
-          max_q_head_num_per_iter / actual_q_heads_per_kv;
       const int64_t q_token_num_stride = input->query_num_tokens_stride;
       const int64_t q_head_num_stride = input->query_num_heads_stride;
       const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride;
@@ -1461,15 +1482,6 @@ class AttentionMainLoop {
               sizeof(q_buffer_t), sizeof(logits_buffer_t),
               sizeof(partial_output_buffer_t), max_q_head_num_per_iter,
               max_q_head_num_per_iter);
-      const int32_t default_q_tile_token_num =
-          default_tile_size / actual_q_heads_per_kv;
-
-      AttentionWorkItemGroup* const workitem_groups =
-          metadata.workitem_groups_ptr;
-      const int32_t* cu_workitem_num_per_thread =
-          metadata.cu_workitem_num_per_thread;
-      ReductionWorkItemGroup* const reduction_items =
-          metadata.reduction_items_ptr;
 
       const int32_t effective_thread_num = metadata.effective_thread_num;
       const int32_t reduction_item_num = metadata.reduction_item_num;
@@ -1513,8 +1525,6 @@ class AttentionMainLoop {
               cu_workitem_num_per_thread[thread_offset + 1] -
               cu_workitem_num_per_thread[thread_offset];
 
-          const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
-
           for (int32_t workitem_group_idx = 0;
                workitem_group_idx < curr_workitem_groups_num;
                ++workitem_group_idx) {
@@ -1529,6 +1539,21 @@ class AttentionMainLoop {
             const int32_t q_token_id_start =
                 current_workitem_group->q_token_id_start;
             const int32_t q_token_num = current_workitem_group->q_token_num;
+            const bool curr_use_gqa =
+                use_gqa_fast_path || (supports_gqa && q_token_num == 1);
+            if (!use_gqa_fast_path && curr_use_gqa &&
+                kv_head_idx % q_heads_per_kv != 0) {
+              continue;
+            }
+            const int32_t curr_q_heads_per_kv =
+                curr_use_gqa ? q_heads_per_kv : 1;
+            const int32_t curr_max_q_token_num_per_iter =
+                max_q_head_num_per_iter / curr_q_heads_per_kv;
+            const int32_t curr_default_q_tile_token_num =
+                default_tile_size / curr_q_heads_per_kv;
+            const int32_t q_head_start_idx =
+                use_gqa_fast_path ? (kv_head_idx * q_heads_per_kv)
+                                  : kv_head_idx;
 
             // taskgroup general information
             const int32_t q_end = input->query_start_loc[current_group_idx + 1];
@@ -1542,7 +1567,7 @@ class AttentionMainLoop {
                              current_workitem_group->local_split_id == 0);
 
             for (int32_t q_token_offset = 0; q_token_offset < q_token_num;
-                 q_token_offset += default_q_tile_token_num) {
+                 q_token_offset += curr_default_q_tile_token_num) {
               bool first_iter_flag[AttentionScheduler::MaxQTileIterNum];
               for (int32_t i = 0; i < AttentionScheduler::MaxQTileIterNum;
                    ++i) {
@@ -1552,9 +1577,9 @@ class AttentionMainLoop {
               const int32_t q_token_start_idx =
                   q_start + q_token_offset + q_token_id_start;
               const int32_t actual_q_token_num = std::min(
-                  default_q_tile_token_num, q_token_num - q_token_offset);
+                  curr_default_q_tile_token_num, q_token_num - q_token_offset);
               const int32_t q_head_tile_size =
-                  actual_q_token_num * actual_q_heads_per_kv;
+                  actual_q_token_num * curr_q_heads_per_kv;
               const int32_t rounded_q_head_tile_size =
                   ((q_head_tile_size + max_q_head_num_per_iter - 1) /
                    max_q_head_num_per_iter) *
@@ -1591,10 +1616,9 @@ class AttentionMainLoop {
                   AttentionScheduler::align_kv_tile_pos(
                       kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment);
 
-              int32_t curr_kv_head_idx =
-                  use_gqa ? kv_head_idx
-                          : (kv_head_idx /
-                             q_heads_per_kv);  // for GQA disabled case
+              const int32_t curr_kv_head_idx =
+                  use_gqa_fast_path ? kv_head_idx
+                                    : (kv_head_idx / q_heads_per_kv);
 
               // std::printf("thread_id: %d, req_id: %d, q_token_start: %d,
               // q_token_end: %d, q_head_start: %d, q_head_end: %d, kv_head_idx:
@@ -1629,12 +1653,12 @@ class AttentionMainLoop {
                   (s_aux != nullptr ? s_aux + q_head_start_idx : nullptr);
 
               // copy the Q tile to q_buffer, the logical layout of q_buffer is
-              // [actual_q_token_num, actual_q_heads_per_kv, head_dim]
+              // [actual_q_token_num, curr_q_heads_per_kv, head_dim]
               {
                 attn_impl.copy_q_heads_tile(
                     q_tile_ptr, q_buffer, actual_q_token_num,
-                    actual_q_heads_per_kv, q_token_num_stride,
-                    q_head_num_stride, scale);
+                    curr_q_heads_per_kv, q_token_num_stride, q_head_num_stride,
+                    scale);
               }
 
               if (use_sink) {
@@ -1648,29 +1672,29 @@ class AttentionMainLoop {
                 float* __restrict__ curr_max_buffer = max_buffer;
                 for (int32_t token_idx = 0; token_idx < actual_q_token_num;
                      ++token_idx) {
-                  for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
+                  for (int32_t head_idx = 0; head_idx < curr_q_heads_per_kv;
                        ++head_idx) {
                     curr_sum_buffer[head_idx] = 1.0f;
                     curr_max_buffer[head_idx] = s_aux_fp32[head_idx];
                   }
 
-                  curr_sum_buffer += actual_q_heads_per_kv;
-                  curr_max_buffer += actual_q_heads_per_kv;
+                  curr_sum_buffer += curr_q_heads_per_kv;
+                  curr_max_buffer += curr_q_heads_per_kv;
                 }
               } else {
                 float* __restrict__ curr_sum_buffer = sum_buffer;
                 float* __restrict__ curr_max_buffer = max_buffer;
                 for (int32_t token_idx = 0; token_idx < actual_q_token_num;
                      ++token_idx) {
-                  for (int32_t head_idx = 0; head_idx < actual_q_heads_per_kv;
+                  for (int32_t head_idx = 0; head_idx < curr_q_heads_per_kv;
                        ++head_idx) {
                     curr_sum_buffer[head_idx] = 0.0f;
                     curr_max_buffer[head_idx] =
                         std::numeric_limits<float>::lowest();
                   }
 
-                  curr_sum_buffer += actual_q_heads_per_kv;
-                  curr_max_buffer += actual_q_heads_per_kv;
+                  curr_sum_buffer += curr_q_heads_per_kv;
+                  curr_max_buffer += curr_q_heads_per_kv;
                 }
               }
 
@@ -1683,16 +1707,17 @@ class AttentionMainLoop {
                     kv_tile_pos_left + kv_tile_size, rounded_kv_tile_end_pos);
                 for (int32_t q_head_tile_token_offset = 0;
                      q_head_tile_token_offset < actual_q_token_num;
-                     q_head_tile_token_offset += max_q_token_num_per_iter) {
+                     q_head_tile_token_offset +=
+                     curr_max_q_token_num_per_iter) {
                   const int32_t q_tile_pos_left =
                       q_tile_start_pos + q_head_tile_token_offset;
                   const int32_t q_tile_token_num =
-                      std::min(max_q_token_num_per_iter,
+                      std::min(curr_max_q_token_num_per_iter,
                                actual_q_token_num - q_head_tile_token_offset);
                   const int32_t q_tile_head_offset =
-                      q_head_tile_token_offset * actual_q_heads_per_kv;
+                      q_head_tile_token_offset * curr_q_heads_per_kv;
                   const int32_t q_tile_head_num =
-                      q_tile_token_num * actual_q_heads_per_kv;
+                      q_tile_token_num * curr_q_heads_per_kv;
                   const int32_t q_tile_pos_right =
                       q_tile_pos_left + q_tile_token_num;
                   const auto [actual_kv_tile_pos_left,
@@ -1702,7 +1727,7 @@ class AttentionMainLoop {
                           q_tile_pos_right, sliding_window_left,
                           sliding_window_right);
                   const int32_t q_iter_idx =
-                      q_head_tile_token_offset / max_q_token_num_per_iter;
+                      q_head_tile_token_offset / curr_max_q_token_num_per_iter;
 
                   if (actual_kv_tile_pos_right <= actual_kv_tile_pos_left) {
                     continue;
@@ -1768,7 +1793,7 @@ class AttentionMainLoop {
                       aligned_actual_kv_tile_pos_left,
                       aligned_actual_kv_tile_pos_right, actual_kv_token_num,
                       kv_cache_block_num_stride, q_tile_head_num,
-                      q_tile_token_num, q_tile_pos_left, actual_q_heads_per_kv,
+                      q_tile_token_num, q_tile_pos_left, curr_q_heads_per_kv,
                       block_size, sliding_window_left, sliding_window_right,
                       scale, softcap_scale, curr_alibi_slopes,
                       first_iter_flag[q_iter_idx], use_sink, debug_info);
@@ -1782,11 +1807,11 @@ class AttentionMainLoop {
                   final_output(partial_q_buffer,
                                reinterpret_cast<query_t*>(input->output) +
                                    output_buffer_offset,
-                               sum_buffer, actual_q_heads_per_kv,
+                               sum_buffer, curr_q_heads_per_kv,
                                actual_q_token_num, q_head_num, output_v_scale);
                 } else {
                   const int32_t stride =
-                      actual_q_heads_per_kv * split_kv_q_token_num_threshold;
+                      curr_q_heads_per_kv * split_kv_q_token_num_threshold;
                   buffer_manager.update(kv_head_idx, total_reduction_split_num,
                                         head_dim, stride, sizeof(float));
                   volatile bool* split_flag_buffer =
@@ -1822,18 +1847,26 @@ class AttentionMainLoop {
           const int32_t curr_split_id = curr_workitem_groups->split_start_id;
           const int32_t curr_split_num = curr_workitem_groups->split_num;
           const int32_t current_group_idx = curr_workitem_groups->req_id;
+          const bool curr_use_gqa =
+              use_gqa_fast_path || (supports_gqa && curr_output_token_num == 1);
+          if (!use_gqa_fast_path && curr_use_gqa &&
+              kv_head_idx % q_heads_per_kv != 0) {
+            continue;
+          }
+          const int32_t curr_q_heads_per_kv = curr_use_gqa ? q_heads_per_kv : 1;
           const int32_t curr_output_head_num =
-              curr_output_token_num * actual_q_heads_per_kv;
+              curr_output_token_num * curr_q_heads_per_kv;
 
           const int32_t q_start = input->query_start_loc[current_group_idx];
           const int32_t q_token_start_idx = q_start + curr_output_token_idx;
-          const int32_t q_head_start_idx = kv_head_idx * actual_q_heads_per_kv;
+          const int32_t q_head_start_idx =
+              use_gqa_fast_path ? (kv_head_idx * q_heads_per_kv) : kv_head_idx;
           size_t output_buffer_offset =
               q_token_start_idx * q_head_num * head_dim +
               q_head_start_idx * head_dim;
 
           const int32_t stride =
-              actual_q_heads_per_kv * split_kv_q_token_num_threshold;
+              curr_q_heads_per_kv * split_kv_q_token_num_threshold;
           buffer_manager.update(kv_head_idx, total_reduction_split_num,
                                 head_dim, stride, sizeof(float));
           volatile bool* split_flag_buffer =
@@ -1852,7 +1885,7 @@ class AttentionMainLoop {
           final_output(
               split_output_buffer,
               reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
-              split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
+              split_sum_buffer, curr_q_heads_per_kv, curr_output_token_num,
               q_head_num, output_v_scale);
         }
       }

docker/Dockerfile.cpu

@@ -27,11 +27,14 @@ WORKDIR /workspace
 ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
 
+ARG max_jobs=32
+ENV MAX_JOBS=${max_jobs}
+
 # Install minimal dependencies and uv
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt-get update -y \
-    && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates \
+    && apt-get install -y --no-install-recommends sudo ccache git curl wget ca-certificates zlib1g-dev \
     gcc-12 g++-12 libtcmalloc-minimal4 libnuma-dev ffmpeg libsm6 libxext6 libgl1 jq lsof make xz-utils \
     && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
     && curl -LsSf https://astral.sh/uv/install.sh | sh
@@ -123,9 +126,6 @@ RUN --mount=type=cache,target=/root/.cargo/registry \
 ######################### BUILD IMAGE #########################
 FROM base AS vllm-build
 
-ARG max_jobs=32
-ENV MAX_JOBS=${max_jobs}
-
 ARG GIT_REPO_CHECK=0
 # Support for cross-compilation with x86 ISA including AVX2 and AVX512: docker build --build-arg VLLM_CPU_X86="true" ...
 ARG VLLM_CPU_X86=0
@@ -257,8 +257,7 @@ WORKDIR /vllm-workspace
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
-    uv pip install dist/*.whl && \
-    uv pip install "vllm[audio]"
+    uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
 
 # Add labels to document build configuration
 LABEL org.opencontainers.image.title="vLLM CPU"

docs/benchmarking/cli.md

@@ -231,13 +231,21 @@ vllm bench serve \
 
 #### Custom Image Dataset
 
-If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
+If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and can use "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
 
 ```json
 {"prompt": "How many animals are present in the given image?", "image_files": ["/path/to/image/folder/horsepony.jpg"]}
 {"prompt": "What colour is the bird shown in the image?", "image_files": ["/path/to/image/folder/flycatcher.jpeg"]}
 ```
 
+Every image listed in "image_files" is added to the request in the listed order after the prompt text. To preserve an interleaved order of text and images, use a "content" field with OpenAI-compatible content parts:
+
+```json
+{"content": [{"type": "text", "text": "Compare "}, {"type": "image", "image": "/path/to/image/folder/chart_a.png"}, {"type": "text", "text": " with "}, {"type": "image_url", "image_url": {"url": "/path/to/image/folder/chart_b.png"}}]}
+```
+
+The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string.
+
 ```bash
 # need a model with vision capability here
 vllm serve Qwen/Qwen2-VL-7B-Instruct

docs/features/disagg_prefill.md

@@ -19,25 +19,25 @@ Two main reasons:
 
 Please refer to [examples/disaggregated/disaggregated_prefill.sh](../../examples/disaggregated/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
 
-Now supports 6 types of connectors:
+Now supports 9 types of connectors:
 
 - **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
 - **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
-- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md).
+- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as:
+
+  ```bash
+  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
+  ```
+
 - **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
 - **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
+- **MoRIIOConnector** (ROCm only): see [MoRI-IO Usage Guide](moriio_connector_usage.md) for example usage and detailed documentation.
 - **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:
 
   ```bash
   --kv-transfer-config '{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both"},{"kv_connector":"ExampleConnector","kv_role":"kv_both","kv_connector_extra_config":{"shared_storage_path":"local_storage"}}]}}'
   ```
 
-For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
-
-  ```bash
-  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both", "kv_buffer_device":"cuda", "kv_connector_extra_config":{"backends":["UCX", "GDS"]}}'
-  ```
-
 - **OffloadingConnector**: enable offloading of KV data to CPU memory, customizing the CPU block size (in tokens) and total CPU memory bytes to allocate:
 
   ```bash

docs/features/moriio_connector_usage.md

@@ -0,0 +1,264 @@
+# MoRIIOConnector Usage Guide
+
+`MoRIIOConnector` is a high-performance KV connector used for KV cache transfer in PD disaggregated deployments, built on ROCm's [MoRI-IO](https://github.com/rocm/mori) communication library for point-to-point communication with ultra-low overhead.
+
+## Prerequisites
+
+### Installation
+
+**Docker:** MoRI is shipped with the official ROCm vLLM image: `vllm/vllm-openai-rocm:nightly`.
+
+**Manual installation:** MoRI wheel can be installed with
+
+```bash
+pip install amd_mori
+```
+
+Refer to the [Dockerfile.rocm_base](../../docker/Dockerfile.rocm_base) for more information, or [official MoRI repository](https://github.com/rocm/mori) for instructions on how to build MoRI from source.
+
+For instructions on installing appropriate NIC userspace libraries, see [Installing NIC userspace libraries](#appendix-installing-nic-userspace-libraries).
+
+## Basic usage (single host)
+
+Start the proxy first; the producer and consumer instances will retry registration until the proxy is reachable.
+
+### Producer (prefiller) configuration
+
+Start a prefiller instance that produces KV caches
+
+```bash
+# Prefill instance (GPU 0-3) 
+export VLLM_ROCM_USE_AITER=1
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+export HIP_VISIBLE_DEVICES=0,1,2,3
+ 
+vllm serve Qwen/Qwen3-235B-A22B-FP8 \
+  -tp 4 \
+  --port 20005 \
+  --gpu-memory-utilization 0.9 \
+  --kv-transfer-config '{
+    "kv_connector": "MoRIIOConnector",
+    "kv_role": "kv_producer",
+    "kv_connector_extra_config": {
+      "proxy_ip": "127.0.0.1",
+      "proxy_ping_port": "36367",
+      "http_port": "20005",
+      "handshake_port": "6301",
+      "notify_port": "6105"
+    }
+  }'
+```
+
+### Consumer (decoder) configuration
+
+Start a decoder instance that consumes KV caches:
+
+```bash
+# Decode instance (GPU 4-7)
+export VLLM_ROCM_USE_AITER=1
+export CUDA_VISIBLE_DEVICES=4,5,6,7
+export HIP_VISIBLE_DEVICES=4,5,6,7
+
+vllm serve Qwen/Qwen3-235B-A22B-FP8 \
+  -tp 4 \
+  --port 40005 \
+  --gpu-memory-utilization 0.9 \
+  --kv-transfer-config '{
+    "kv_connector": "MoRIIOConnector",
+    "kv_role": "kv_consumer",
+    "kv_connector_extra_config": {
+      "proxy_ip": "127.0.0.1",
+      "http_port": "40005",
+      "proxy_ping_port": "36367",
+      "handshake_port": "7301",
+      "notify_port": "7501"
+    }
+  }'
+```
+
+### Proxy server
+
+The proxy fronts the producer and consumer instances and routes incoming requests to them. `vllm-router` is the recommended proxy; it can be installed manually or run as a Docker container. Note that the port `36367` below is the `proxy_ping_port` configured on each vLLM instance.
+
+**Docker:**
+
+```bash
+docker run \
+  --network host \
+  vllm/vllm-router:nightly \
+  vllm-router \
+  --vllm-pd-disaggregation \
+  --kv-connector moriio \
+  --vllm-discovery-address "0.0.0.0:36367"
+```
+
+**Manual install:**
+
+```bash
+pip install vllm-router
+vllm-router \
+  --vllm-pd-disaggregation \
+  --kv-connector moriio \
+  --vllm-discovery-address "0.0.0.0:36367"
+```
+
+Alternatively, you can use the reference implementation proxy shipped with vLLM:
+
+```bash
+cd <path_to>/vllm
+pip install quart aiohttp msgpack
+python examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+```
+
+## Configuration
+
+The connector is configured at two levels: the application level and the transport level.
+
+### Application-level configuration
+
+**Modes:** MoRI has two modes of operation: WRITE and READ mode.
+
+- In WRITE mode, the producer actively pushes computed KV blocks after every layer into the consumer's memory.
+- In READ mode, the consumer pulls the KV blocks from the producer all at once, as soon as it has been notified those blocks are ready.
+
+WRITE mode is used by default. READ mode can be configured by setting `--kv-transfer-config.kv_connector_extra_config.read_mode true`.
+
+**Control-plane configuration:** MoRI moves KV bytes over RDMA/xGMI, but producers and consumers also need out-of-band TCP channels for handshake, block id exchange, liveness, and completion signaling. These keys live under `kv_connector_extra_config`:
+
+- `proxy_ip`: IP address of the disaggregation proxy/router that fronts the prefiller and decoder. Each vLLM instance uses it to register itself and to send heartbeats so the proxy knows where to route incoming requests.
+- `proxy_ping_port`: TCP port on `proxy_ip` where the proxy listens for instance heartbeats and registration messages. Used to detect dead vLLM instances and keep routing tables fresh.
+- `http_port`: HTTP port that this vLLM instance exposes its OpenAI-compatible API on. The proxy registers this port, and forwards user requests to this port once it has picked an instance.
+- `handshake_port`: TCP port used for the one-time MoRI engine handshake between a prefiller and a decoder. The two sides exchange RDMA engine descriptors here before any KV transfer can happen.
+- `notify_port`: TCP port used for control and synchronization messages between prefiller and decoder. Used differently in the two modes:
+    - WRITE mode: **Block allocation:** the decoder notifies the prefiller about its block ids, so the prefiller can push its computed KV blocks into the correct place on the decoder instance. **Completion:** once all blocks have been transferred, the prefiller notifies the decoder that it's safe to use its blocks.
+    - READ mode: **Completion:** once the decoder has read all blocks from the prefiller, it notifies the prefiller so it can free its KV cache blocks.
+
+!!! note
+    `notify_port` is used as a *base* port: each (DP rank, TP rank) pair within an instance uses `notify_port + offset` where the offset is based on the rank. Make sure the range starting at `notify_port` is free on the host.
+
+### Transport configuration
+
+MoRI has two transport backends: RDMA and xGMI. You can select backend using `--kv-transfer-config.kv_connector_extra_config.backend $BACKEND`, with `$BACKEND` being `rdma` or `xgmi`. RDMA is the default backend and should be used in multi-node deployments.
+
+The configuration options for each backend are as follows.
+
+#### RDMA backend
+
+- `qp_per_transfer`: number of RDMA Queue Pairs (QPs) used per transfer. More QPs let a single transfer be striped over multiple QPs to increase NIC concurrency, at the cost of more RDMA resources.
+- `post_batch_size`: how many RDMA Work Requests (WR) are batched into one `ibv_post_send` doorbell. Defaults to -1, meaning the backend default. Larger batches reduce the posting overhead per WR.
+- `num_workers`: number of worker threads MoRI uses to post and poll transfer completions.
+
+Advanced users can also configure MoRI itself using environment variables such as `MORI_IO_QP_MAX_SEND_WR`, `MORI_IO_QP_MAX_CQE`, etc. These are MoRI library variables and are separate from vLLM's own `VLLM_MORIIO_*` settings. Refer to the [MoRI repository](https://github.com/rocm/mori) for more information.
+
+#### xGMI backend
+
+Use xGMI when the prefiller and decoder run on the same physical host so transfers go over the AMD GPU fabric and skip the NIC entirely. Currently only configured using MoRI-specific environment variables; see the [MoRI repository](https://github.com/rocm/mori).
+
+## Multi-node deployment
+
+The example below shows how to run a 1P1D deployment on two nodes. We run the proxy on the same node as the prefill instance.
+
+### On both nodes
+
+```bash
+# Set on both nodes before running any command
+export PREFILL_IP=<node1-ip>
+export DECODE_IP=<node2-ip>
+```
+
+### On node 1
+
+Start the proxy first as described in [Proxy server](#proxy-server), then start the prefill instance:
+
+```bash
+docker run \
+  --name moriio-prefill \
+  --init --network host --ipc host --privileged \
+  --security-opt seccomp=unconfined \
+  --ulimit memlock=-1 --ulimit stack=67108864 --shm-size 256G \
+  --group-add video --group-add render \
+  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+  -e VLLM_ROCM_USE_AITER=1 \
+  vllm/vllm-openai-rocm:nightly \
+  deepseek-ai/DeepSeek-R1-0528 \
+    --port 8100 \
+    --tensor-parallel-size 8 \
+    --enable-expert-parallel \
+    --gpu-memory-utilization 0.8 \
+    --trust-remote-code \
+    --kv-transfer-config '{
+      "kv_connector": "MoRIIOConnector",
+      "kv_role": "kv_producer",
+      "kv_connector_extra_config": {
+        "proxy_ip": "'"${PREFILL_IP}"'",
+        "proxy_ping_port": "36367",
+        "http_port": "8100",
+        "handshake_port": "6301",
+        "notify_port": "61005"
+      }
+    }'
+```
+
+### On node 2
+
+Decode instance:
+
+```bash
+docker run \
+  --name moriio-decode \
+  --init --network host --ipc host --privileged \
+  --security-opt seccomp=unconfined \
+  --ulimit memlock=-1 --ulimit stack=67108864 --shm-size 256G \
+  --group-add video --group-add render \
+  --device /dev/kfd --device /dev/dri --device /dev/infiniband \
+  -e VLLM_ROCM_USE_AITER=1 \
+  vllm/vllm-openai-rocm:nightly \
+  deepseek-ai/DeepSeek-R1-0528 \
+    --port 8200 \
+    --tensor-parallel-size 8 \
+    --gpu-memory-utilization 0.8 \
+    --trust-remote-code \
+    --enable-expert-parallel \
+    --kv-transfer-config '{
+      "kv_connector": "MoRIIOConnector",
+      "kv_role": "kv_consumer",
+      "kv_connector_extra_config": {
+        "proxy_ip": "'"${PREFILL_IP}"'",
+        "proxy_ping_port": "36367",
+        "http_port": "8200",
+        "handshake_port": "6301",
+        "notify_port": "61005"
+      }
+    }'
+```
+
+## Troubleshooting
+
+### `availDevices.size() > 0` assertion failure
+
+**Problem:** vLLM fails to launch with the following log:
+
+```bash
+libibverbs: Warning: Driver bnxt_re does not support the kernel ABI of 6 (supports 1 to 1) for device /sys/class/infiniband/rdma4
+...
+ker: /app/mori/src/io/rdma/backend_impl.cpp: mori::io::RdmaManager::RdmaManager(const RdmaBackendConfig, application::RdmaContext *): Assertion `availDevices.size() > 0' failed.
+```
+
+**Fix:** The installed RDMA userspace libraries do not match the driver and firmware version installed on the host. You must install NIC userspace libraries corresponding to your RDMA kernel module and firmware version. See [Installing NIC userspace
+libraries](#appendix-installing-nic-userspace-libraries) for more information.
+
+## Appendix: installing NIC userspace libraries
+
+To run MoRI with RDMA, your environment must have the necessary RDMA userspace libraries installed that match the associated kernel module and firmware version.
+
+The official image `vllm/vllm-openai-rocm:nightly` comes pre-installed with userspace libraries for the following NICs and kernel module versions:
+
+- AINIC (AMD Pensando Pollara): version `1.117.3-hydra`, tested with `ioinic-dkms=25.11.1.001`
+- Thor2 (Broadcom): version `235.2.86.0`, tested with `bnxt-en-dkms=1.10.3.235.2.86.0`, `bnxt-re-dkms=235.2.86.0`
+
+Refer to [Dockerfile.rocm](../../docker/Dockerfile.rocm) for more details. For users with NICs, kernel modules, and/or FW other than those stated above we refer to
+the vendors' own installation instructions.
+
+## Further reading
+
+- [Next-Level Inference: Why Your Single-Node vLLM Setup Needs Prefill-Decode Disaggregation](https://vllm.ai/blog/2026-04-07-moriio-kv-connector).

docs/features/reasoning_outputs.md

@@ -15,7 +15,6 @@ vLLM currently supports the following reasoning models:
 | ------------ | ----------- | ---------------- | ----------- |
 | [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
 | [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
-| [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ |
 | [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
 | [ERNIE-4.5-VL series](https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-PT) | `ernie45` | `json`, `regex` | ❌ |
 | [ERNIE-4.5-21B-A3B-Thinking](https://huggingface.co/baidu/ERNIE-4.5-21B-A3B-Thinking) | `ernie45` | `json`, `regex` | ✅ |
@@ -30,7 +29,6 @@ vLLM currently supports the following reasoning models:
 !!! note
     IBM Granite 3.2 and DeepSeek-V3.1 reasoning is disabled by default; to enable it, you must also pass `thinking=True` in your `chat_template_kwargs`.
     The reasoning feature for the Qwen3 series is enabled by default. To disable it, you must pass `enable_thinking=False` in your `chat_template_kwargs`.
-    Gemma 4 reasoning is disabled by default; to enable it, pass `enable_thinking=True` in your `chat_template_kwargs` or set `reasoning_effort` (which enables it automatically).
     DeepSeek-V3.1 tool calling is supported in non-thinking mode.
     Holo2 reasoning is enabled by default. To disable it, you must also pass `thinking=False` in your `chat_template_kwargs`.
 
@@ -316,44 +314,9 @@ for output in outputs:
     print("text:", output.outputs[0].text)
 ```
 
-## Automatic `enable_thinking` Activation
-
-Some models (such as Gemma 4, DeepSeek-V4-Pro and IBM Granite 3.2) require `enable_thinking: true` in their chat template kwargs to activate thinking mode — without it, reasoning tokens are never generated regardless of other settings.
-
-When you set `reasoning_effort` in a Chat Completions request (or `reasoning.effort` in a Responses API request), vLLM automatically injects `enable_thinking` into the chat template kwargs:
-
-- `reasoning_effort` = `"low"`, `"medium"`, or `"high"` → `enable_thinking = true`
-- `reasoning_effort` = `"none"` → `enable_thinking = false`
-- `reasoning_effort` not set → `enable_thinking` is not injected (preserves existing behavior)
-
-This means you no longer need to manually pass `chat_template_kwargs: {"enable_thinking": true}` when using `reasoning_effort` — it is handled automatically.
-
-!!! note
-    If you explicitly set `enable_thinking` in `chat_template_kwargs`, your value takes priority over the automatic injection. This allows you to override the behavior if needed.
-
-    For models whose templates don't declare `enable_thinking` (e.g., DeepSeek R1), the injected kwarg is harmlessly filtered out by `resolve_chat_template_kwargs`.
-
-### Example
-
-```python
-from openai import OpenAI
-
-client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")
-
-# reasoning_effort automatically enables thinking for models that need it
-response = client.chat.completions.create(
-    model="google/gemma-4-26B-A4B-it",
-    messages=[{"role": "user", "content": "What is 15 * 37?"}],
-    reasoning_effort="high",  # Automatically sets enable_thinking=true
-)
-
-print(response.choices[0].message.reasoning)
-print(response.choices[0].message.content)
-```
-
 ## Limitations
 
-- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`).
+- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`).
 
 ## How to support a new reasoning model
 

requirements/build/rocm.txt

@@ -16,4 +16,3 @@ wheel
 jinja2>=3.1.6
 amdsmi==7.0.2
 timm>=1.0.17
-tilelang==0.1.10

requirements/cuda.txt

@@ -21,7 +21,7 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
 fastsafetensors >= 0.2.2
 
 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
-nvidia-cutlass-dsl[cu13]==4.5.2
+nvidia-cutlass-dsl[cu13]==4.5.0
 quack-kernels>=0.3.3
 
 # Tokenspeed_MLA for faster mla with spec decode

requirements/rocm.txt

@@ -22,4 +22,3 @@ timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm 
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
-tilelang==0.1.10

requirements/test/rocm.in

@@ -43,7 +43,6 @@ schemathesis>=3.39.15 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9
-tilelang==0.1.10
 
 genai_perf>=0.0.8
 tritonclient>=2.51.0

requirements/test/rocm.txt

@@ -43,9 +43,7 @@ anyio==4.13.0
     #   starlette
     #   watchfiles
 apache-tvm-ffi==0.1.10
-    # via
-    #   tilelang
-    #   xgrammar
+    # via xgrammar
 arctic-inference==0.1.1
     # via -r requirements/test/rocm.in
 argcomplete==3.6.3
@@ -131,9 +129,7 @@ click==8.3.1
     #   typer
     #   uvicorn
 cloudpickle==3.1.2
-    # via
-    #   -r requirements/test/../common.txt
-    #   tilelang
+    # via -r requirements/test/../common.txt
 colorama==0.4.6
     # via
     #   perceptron
@@ -515,8 +511,6 @@ mistral-common==1.11.2
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   -r requirements/test/rocm.in
-ml-dtypes==0.5.4
-    # via tilelang
 model-hosting-container-standards==0.1.14
     # via
     #   -c requirements/common.txt
@@ -593,7 +587,6 @@ numpy==2.2.6
     #   lm-eval
     #   matplotlib
     #   mistral-common
-    #   ml-dtypes
     #   mteb
     #   numba
     #   opencv-python-headless
@@ -617,7 +610,6 @@ numpy==2.2.6
     #   statsmodels
     #   tensorizer
     #   tifffile
-    #   tilelang
     #   torchvision
     #   transformers
     #   tritonclient
@@ -819,7 +811,6 @@ psutil==7.2.2
     #   accelerate
     #   peft
     #   tensorizer
-    #   tilelang
 py==1.11.0
     # via pytest-forked
 py-cpuinfo==9.0.0
@@ -1201,10 +1192,6 @@ tiktoken==0.12.0
     #   gpt-oss
     #   lm-eval
     #   mistral-common
-tilelang==0.1.10
-    # via
-    #   -c requirements/rocm.txt
-    #   -r requirements/test/rocm.in
 timm==1.0.17
     # via
     #   -c requirements/rocm.txt
@@ -1221,8 +1208,6 @@ tomli==2.4.0
     # via schemathesis
 tomli-w==1.2.0
     # via schemathesis
-torch-c-dlpack-ext==0.1.5
-    # via tilelang
 tqdm==4.67.3
     # via
     #   -r requirements/test/../common.txt
@@ -1240,7 +1225,6 @@ tqdm==4.67.3
     #   pqdm
     #   segmentation-models-pytorch
     #   sentence-transformers
-    #   tilelang
     #   transformers
 transformers==5.5.3
     # via
@@ -1309,7 +1293,6 @@ typing-extensions==4.15.0
     #   sentence-transformers
     #   sqlalchemy
     #   starlette
-    #   tilelang
     #   torch
     #   typeguard
     #   typing-inspection
@@ -1376,8 +1359,6 @@ yarl==1.23.0
     # via
     #   aiohttp
     #   schemathesis
-z3-solver==4.15.4.0
-    # via tilelang
 zipp==3.23.0
     # via importlib-metadata
 

requirements/tpu.txt

@@ -12,4 +12,4 @@ ray[data]
 setuptools==78.1.0
 setuptools-rust>=1.9.0
 nixl==0.3.0
-tpu-inference==0.19.0
+tpu-inference==0.20.0

setup.py

@@ -1195,6 +1195,11 @@ setup(
             "opentelemetry-exporter-otlp>=1.26.0",
             "opentelemetry-semantic-conventions-ai>=0.4.1",
         ],
+        "triton-cpu": [
+            "triton @ "
+            "git+https://github.com/triton-lang/triton-cpu.git@270e696d ; "
+            "platform_machine == 'x86_64'",
+        ],  # Remove after stable release
     },
     cmdclass=cmdclass,
     package_data=package_data,

tests/benchmarks/test_custom_image_dataset.py

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

tests/distributed/test_quick_all_reduce.py

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

tests/distributed/test_rocm_quick_reduce.py

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

tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py

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

tests/entrypoints/openai/test_reasoning_enable_thinking.py

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

tests/entrypoints/test_api_server_process_manager.py

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

tests/kernels/moe/test_cpu_int4_moe.py

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

tests/kernels/moe/test_moe.py

@@ -1656,7 +1656,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
         layer.routing_method_type = RoutingMethodType.Renormalize
         layer.expert_map = None
         layer.apply_router_weight_on_input = False
-        layer.routed_scaling_factor = 2.446
+        layer.routed_scaling_factor = None
         layer.shared_experts = None
         layer._expert_routing_tables = lambda: None
 
@@ -1678,10 +1678,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
         # Compute torch baseline
         w1_original = w1.clone()
         w2_original = w2.clone()
-        baseline_output = (
-            torch_moe(a, w1_original, w2_original, router_logits, topk)
-            * layer.routed_scaling_factor
-        )
+        baseline_output = torch_moe(a, w1_original, w2_original, router_logits, topk)
 
     close = torch.isclose(trtllm_output, baseline_output, atol=1e-1, rtol=0.85)
     assert close.float().mean() > 0.925

tests/kernels/test_mhc_kernels.py

@@ -4,12 +4,7 @@ import pytest
 import torch
 
 import vllm.model_executor.kernels.mhc  # noqa: F401
-from vllm.model_executor.kernels.mhc.tilelang import (
-    _tilelang_hc_prenorm_gemm,
-    _torch_hc_prenorm_gemm,
-)
 from vllm.platforms import current_platform
-from vllm.utils.import_utils import has_tilelang
 from vllm.utils.torch_utils import set_random_seed
 
 DEVICE = current_platform.device_type
@@ -97,128 +92,8 @@ def hc_head_ref(
 
 
 @pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
-)
-@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
-@pytest.mark.parametrize("hidden_size", [4096, 7168])
-@pytest.mark.parametrize("hc_mult", [4])
-def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult):
-    torch.set_default_device(DEVICE)
-    set_random_seed(0)
-
-    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
-    hc_mult2 = hc_mult * hc_mult
-    hc_mult3 = 2 * hc_mult + hc_mult2
-    fn = (
-        torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float)
-        * 1e-4
-        * (1 + torch.arange(hc_mult).mul(0.01).view(1, -1, 1))
-    ).flatten(1, 2)
-    hc_scale = torch.randn((3,), dtype=torch.float) * 0.1
-    hc_base = torch.randn((hc_mult3,), dtype=torch.float) * 0.1
-
-    hc_sinkhorn_eps = hc_pre_eps = rms_eps = 1e-6
-    sinkhorn_repeat = 20
-    hc_post_alpha = 1.0
-
-    ref = mhc_pre_ref(
-        residual,
-        fn,
-        hc_scale,
-        hc_base,
-        rms_eps,
-        hc_pre_eps,
-        hc_sinkhorn_eps,
-        hc_post_alpha,
-        sinkhorn_repeat,
-    )
-    out = torch.ops.vllm.mhc_pre_tilelang(
-        residual,
-        fn,
-        hc_scale,
-        hc_base,
-        rms_eps,
-        hc_pre_eps,
-        hc_sinkhorn_eps,
-        hc_post_alpha,
-        sinkhorn_repeat,
-    )
-
-    for actual, expected in zip(out, ref, strict=True):
-        torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2)
-
-
-@pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
-)
-@pytest.mark.parametrize(
-    ("num_tokens", "hidden_size"),
-    [
-        (1, 1280),
-        (512, 1280),
-        (2048, 1280),
-        (1, 4096),
-        (64, 4096),
-        (512, 4096),
-        (2048, 4096),
-        (1, 7168),
-        (64, 7168),
-        (512, 7168),
-        (2048, 7168),
-    ],
-)
-def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size):
-    torch.set_default_device(DEVICE)
-    set_random_seed(0)
-
-    hc_mult = 4
-    hc_mult3 = 2 * hc_mult + hc_mult * hc_mult
-    x = torch.randn((num_tokens, hc_mult * hidden_size), dtype=torch.bfloat16)
-    fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
-    out_ref = torch.empty((1, num_tokens, hc_mult3), dtype=torch.float32)
-    sqrsum_ref = torch.empty((1, num_tokens), dtype=torch.float32)
-    out = torch.empty_like(out_ref)
-    sqrsum = torch.empty_like(sqrsum_ref)
-
-    _torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref)
-    _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult)
-
-    torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4)
-    torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4)
-
-
-@pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
-)
-@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
-@pytest.mark.parametrize("hidden_size", [4096, 7168])
-@pytest.mark.parametrize("hc_mult", [4])
-def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult):
-    torch.set_default_device(DEVICE)
-    set_random_seed(0)
-
-    x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
-    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
-    post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
-    comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
-
-    ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix)
-    out = torch.ops.vllm.mhc_post_tilelang(
-        x,
-        residual,
-        post_layer_mix,
-        comb_res_mix,
-    )
-
-    torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2)
-
-
-@pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
+    not current_platform.is_cuda(),
+    reason="CUDA required",
 )
 @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
 @pytest.mark.parametrize("hidden_size", [4096, 7168])
@@ -321,42 +196,3 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult):
 
     out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
     torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)
-
-
-@pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() and has_tilelang()),
-    reason="CUDA or ROCm and tilelang required",
-)
-@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128])
-@pytest.mark.parametrize("hidden_size", [4096, 7168])
-@pytest.mark.parametrize("hc_mult", [4])
-def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult):
-    torch.set_default_device(DEVICE)
-    set_random_seed(0)
-
-    residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
-    fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
-    hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
-    hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
-    rms_eps = hc_eps = 1e-6
-
-    out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16)
-    out.fill_(float("nan"))
-
-    result = torch.ops.vllm.hc_head_fused_kernel_tilelang(
-        residual,
-        fn,
-        hc_scale,
-        hc_base,
-        out,
-        hidden_size,
-        rms_eps,
-        hc_eps,
-        hc_mult,
-    )
-
-    assert result is None
-    assert not torch.isnan(out).any()
-
-    out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
-    torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2)

tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py

@@ -6,7 +6,6 @@ import tempfile
 
 import huggingface_hub.constants
 import torch
-from safetensors.torch import save_file
 
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
@@ -15,27 +14,6 @@ from vllm.model_executor.model_loader.weight_utils import (
 )
 
 
-def test_runai_safetensors_weights_iterator_clones_reused_buffers(
-    tmp_path, monkeypatch
-):
-    monkeypatch.setenv("RUNAI_STREAMER_MEMORY_LIMIT", "0")
-    weights_file = tmp_path / "model.safetensors"
-    expected_tensors = {
-        "first": torch.tensor([1.0, 2.0]),
-        "second": torch.tensor([3.0, 4.0]),
-    }
-    save_file(expected_tensors, weights_file)
-
-    actual_tensors = dict(
-        runai_safetensors_weights_iterator([str(weights_file)], False)
-    )
-
-    assert actual_tensors.keys() == expected_tensors.keys()
-    assert actual_tensors["first"].data_ptr() != actual_tensors["second"].data_ptr()
-    for name, expected_tensor in expected_tensors.items():
-        assert torch.equal(actual_tensors[name], expected_tensor)
-
-
 def test_runai_model_loader():
     with tempfile.TemporaryDirectory() as tmpdir:
         huggingface_hub.constants.HF_HUB_OFFLINE = False

tests/model_executor/test_nemotron_h_quantization.py

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

tests/model_executor/test_qwen3_5_quantization.py

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

tests/quantization/test_modelopt.py

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

tests/quantization/test_register_quantization_config.py

@@ -108,14 +108,14 @@ def test_register_quantization_config(caplog_vllm):
     assert get_quantization_config("custom_quant") == CustomQuantConfig
 
     # The quantization method `custom_quant` is already exists,
-    # should raise a warning when re-registering it.
-    with caplog_vllm.at_level(logging.WARNING):
+    # should raise a debug message when re-registering it.
+    with caplog_vllm.at_level(logging.DEBUG, logger="vllm"):
         register_quantization_config("custom_quant")(CustomQuantConfig)
 
     assert any(
         "The quantization method 'custom_quant' already exists" in message
         for message in caplog_vllm.messages
-    ), "Expected a warning when re-registering custom_quant"
+    ), "Expected a debug message when re-registering custom_quant"
 
 
 @pytest.mark.parametrize(

tests/utils.py

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

tests/v1/kv_connector/unit/test_hma_auto_config.py

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

tests/v1/kv_connector/unit/test_moriio_connector.py

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

tests/v1/kv_connector/unit/test_multi_connector.py

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

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

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

tests/v1/kv_connector/unit/test_offloading_connector.py

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

tests/v1/logits_processors/test_correctness.py

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

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

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

vllm/_custom_ops.py

@@ -12,6 +12,7 @@ from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType
 from vllm.utils.flashinfer import (
     flashinfer_quant_nvfp4_8x4_sf_layout,
+    flashinfer_trtllm_fp4_8x4_is_safe,
 )
 from vllm.utils.math_utils import cdiv
 
@@ -1642,7 +1643,9 @@ def scaled_fp4_quant(
             f"padded_n has to be a multiple of {block_size}, but got {padded_n}."
         )
 
-    use_8x4_sf_layout = True if "trtllm" in backend and m <= 32 else False  # noqa: SIM210
+    use_8x4_sf_layout = (
+        "trtllm" in backend and m <= 32 and flashinfer_trtllm_fp4_8x4_is_safe()
+    )
     if use_8x4_sf_layout and padded_n is not None and padded_n != n:
         # TODO: support this case
         raise ValueError("padded_n is not supported with TRTLLM 8x4 scale layout.")

vllm/_tilelang_ops.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import math
 from functools import cache
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import torch
 
@@ -10,9 +10,8 @@ from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_tilelang
 from vllm.utils.math_utils import cdiv
 
-# TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so
-# registering the Python wrapper modules does not require TileLang everywhere.
-if TYPE_CHECKING or current_platform.is_cuda_alike():
+# tilelang is only available on CUDA platforms
+if TYPE_CHECKING or current_platform.is_cuda():
     if not has_tilelang():
         raise ImportError(
             "tilelang is required for mhc but is not installed. Install it with "
@@ -24,8 +23,6 @@ else:
     tilelang = None  # type: ignore[assignment]
     T = None  # type: ignore[assignment]
 
-ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()
-
 
 @cache
 def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
@@ -40,17 +37,12 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
     return split_k
 
 
-pass_configs: dict[tilelang.PassConfigKey, Any] = {
-    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
-    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
-}
-
-if current_platform.is_cuda():
-    pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10
-
-
 @tilelang.jit(
-    pass_configs=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def mhc_pre_big_fuse_tilelang(
     gemm_out_mul,
@@ -86,8 +78,7 @@ def mhc_pre_big_fuse_tilelang(
     layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
 
     with T.Kernel(num_tokens, threads=96) as i:
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
         ##################################################################
         # _pre_norm_fn_fwd_norm
         rms = T.alloc_fragment(1, T.float32)
@@ -183,16 +174,18 @@ def mhc_pre_big_fuse_tilelang(
                         ol[i1_h] += pre * xl[i_hc, i1_h]
 
                 T.copy(ol, layer_input[i, i0_h * hidden_block])
-
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()
 
 
 # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478
 
 
 @tilelang.jit(
-    pass_configs=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def mhc_pre_big_fuse_with_norm_tilelang(
     gemm_out_mul,
@@ -237,8 +230,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
         T.clear(mixes)
         rms[0] = 0
 
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
 
         for i_split in T.serial(n_splits):
             rms[0] += gemm_out_sqrsum[i_split, i]
@@ -349,12 +341,15 @@ def mhc_pre_big_fuse_with_norm_tilelang(
 
                 T.copy(ol, layer_input[i, i0_h * hidden_block])
 
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()
 
 
 @tilelang.jit(
-    pass_configs=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def mhc_fused_tilelang(
     comb_mix,
@@ -395,8 +390,8 @@ def mhc_fused_tilelang(
 
     with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):
         tid = T.get_thread_binding()
-        warp_id = tid // 32
-        lane = tid % 32
+        warp_id = T.get_warp_idx()
+        lane = T.get_lane_idx()
 
         s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
         s_post = T.alloc_shared((hc,), T.float32)
@@ -412,8 +407,7 @@ def mhc_fused_tilelang(
         T.clear(sqr)
         h_split_start = i_ks * h_per_split
 
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
 
         T.copy(post_mix[i_n, 0], s_post)
         T.copy(comb_mix[i_n, 0, 0], s_comb)
@@ -472,12 +466,15 @@ def mhc_fused_tilelang(
                     v2 += s_warp[w, tile_n]
                 rp_out[i_ks, i_n] = v2
 
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()
 
 
 @tilelang.jit(
-    pass_configs=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def mhc_post_tilelang(
     a,
@@ -510,8 +507,7 @@ def mhc_post_tilelang(
 
         a_local = T.alloc_fragment((hc, hc), T.float32)
         c_local = T.alloc_fragment(hc, T.float32)
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
         T.copy(a[i_n, 0, 0], a_local)
         T.copy(c[i_n, 0], c_local)
 
@@ -527,193 +523,15 @@ def mhc_post_tilelang(
                     x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]
 
             T.copy(x_local, x[i_n, 0, i0_h * h_blk])
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()
 
 
 @tilelang.jit(
-    pass_configs=pass_configs,
-)
-def hc_prenorm_gemm_tilelang(
-    x,
-    fn,
-    out,
-    sqrsum,
-    hidden_size: int,
-    hc_mult: int = 4,
-    n_out: int = 24,
-    n_thr: int = 512,
-    tile_n: int = 12,
-    n_splits: int = 1,
-) -> tilelang.JITKernel:
-    num_tokens = T.dynamic("num_tokens")
-    hc_hidden_size = hc_mult * hidden_size
-    k_per_split = hc_hidden_size // n_splits
-    k_iters = k_per_split // n_thr
-    n_tiles = T.ceildiv(n_out, tile_n)
-
-    x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16)  # type: ignore[no-redef, valid-type]
-    fn: T.Tensor((n_out, hc_hidden_size), T.float32)  # type: ignore[no-redef, valid-type]
-    out: T.Tensor((n_splits, num_tokens, n_out), T.float32)  # type: ignore[no-redef, valid-type]
-    sqrsum: T.Tensor((n_splits, num_tokens), T.float32)  # type: ignore[no-redef, valid-type]
-
-    with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as (
-        i_n,
-        i_t,
-        i_s,
-    ):
-        tid = T.get_thread_binding()
-        acc = T.alloc_local((tile_n,), T.float32)
-        sqr = T.alloc_local((1,), T.float32)
-        T.clear(acc)
-        T.clear(sqr)
-
-        if ENABLE_PDL:
-            T.pdl_sync()
-
-        for it in T.serial(k_iters):
-            i_k = i_s * k_per_split + it * n_thr + tid
-            x_val = x[i_n, i_k]
-            for i_o in T.unroll(tile_n):
-                out_idx = i_t * tile_n + i_o
-                if out_idx < n_out:
-                    acc[i_o] += x_val * fn[out_idx, i_k]
-            if i_t == 0:
-                sqr[0] += x_val * x_val
-
-        for i_o in T.unroll(tile_n):
-            acc[i_o] = T.warp_reduce_sum(acc[i_o])
-        if i_t == 0:
-            sqr[0] = T.warp_reduce_sum(sqr[0])
-
-        lane = tid % 32
-        warp_id = tid // 32
-        num_warps = n_thr // 32
-        warp_acc = T.alloc_shared((num_warps, tile_n), T.float32)
-        warp_sqr = T.alloc_shared(num_warps, T.float32)
-
-        if lane == 0:
-            for i_o in T.unroll(tile_n):
-                warp_acc[warp_id, i_o] = acc[i_o]
-            if i_t == 0:
-                warp_sqr[warp_id] = sqr[0]
-        T.sync_threads()
-
-        if warp_id == 0:
-            if lane < tile_n:
-                reduced_acc = T.alloc_var(T.float32, init=0.0)
-                for i_w in T.unroll(num_warps):
-                    reduced_acc += warp_acc[i_w, lane]
-                out_idx = i_t * tile_n + lane
-                if out_idx < n_out:
-                    out[i_s, i_n, out_idx] = reduced_acc
-            if lane == 0 and i_t == 0:
-                reduced_sqr = T.alloc_var(T.float32, init=0.0)
-                for i_w in T.unroll(num_warps):
-                    reduced_sqr += warp_sqr[i_w]
-                sqrsum[i_s, i_n] = reduced_sqr
-
-        if ENABLE_PDL:
-            T.pdl_trigger()
-
-
-@tilelang.jit(
-    pass_configs=pass_configs,
-)
-def hc_prenorm_gemm_block_m_tilelang(
-    x,
-    fn,
-    out,
-    sqrsum,
-    hidden_size: int,
-    hc_mult: int = 4,
-    n_out: int = 24,
-    n_thr: int = 512,
-    tile_n: int = 12,
-    block_m: int = 2,
-) -> tilelang.JITKernel:
-    num_tokens = T.dynamic("num_tokens")
-    hc_hidden_size = hc_mult * hidden_size
-    k_iters = hc_hidden_size // n_thr
-    n_tiles = T.ceildiv(n_out, tile_n)
-    m_tiles = T.ceildiv(num_tokens, block_m)
-
-    x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16)  # type: ignore[no-redef, valid-type]
-    fn: T.Tensor((n_out, hc_hidden_size), T.float32)  # type: ignore[no-redef, valid-type]
-    out: T.Tensor((1, num_tokens, n_out), T.float32)  # type: ignore[no-redef, valid-type]
-    sqrsum: T.Tensor((1, num_tokens), T.float32)  # type: ignore[no-redef, valid-type]
-
-    with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t):
-        tid = T.get_thread_binding()
-        acc = T.alloc_local((block_m, tile_n), T.float32)
-        sqr = T.alloc_local((block_m,), T.float32)
-        T.clear(acc)
-        T.clear(sqr)
-
-        if ENABLE_PDL:
-            T.pdl_sync()
-
-        for it in T.serial(k_iters):
-            i_k = it * n_thr + tid
-            fn_val = T.alloc_local((tile_n,), T.float32)
-            for i_o in T.unroll(tile_n):
-                out_idx = i_t * tile_n + i_o
-                if out_idx < n_out:
-                    fn_val[i_o] = fn[out_idx, i_k]
-                else:
-                    fn_val[i_o] = 0.0
-            for i_m in T.unroll(block_m):
-                token_idx = i_mt * block_m + i_m
-                if token_idx < num_tokens:
-                    x_val = x[token_idx, i_k]
-                    for i_o in T.unroll(tile_n):
-                        acc[i_m, i_o] += x_val * fn_val[i_o]
-                    if i_t == 0:
-                        sqr[i_m] += x_val * x_val
-
-        for i_m in T.unroll(block_m):
-            for i_o in T.unroll(tile_n):
-                acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o])
-            if i_t == 0:
-                sqr[i_m] = T.warp_reduce_sum(sqr[i_m])
-
-        lane = tid % 32
-        warp_id = tid // 32
-        num_warps = n_thr // 32
-        warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32)
-        warp_sqr = T.alloc_shared((num_warps, block_m), T.float32)
-
-        if lane == 0:
-            for i_m in T.unroll(block_m):
-                for i_o in T.unroll(tile_n):
-                    warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o]
-                if i_t == 0:
-                    warp_sqr[warp_id, i_m] = sqr[i_m]
-        T.sync_threads()
-
-        if warp_id == 0:
-            for i_m in T.unroll(block_m):
-                token_idx = i_mt * block_m + i_m
-                if token_idx < num_tokens:
-                    if lane < tile_n:
-                        reduced_acc = T.alloc_var(T.float32, init=0.0)
-                        for i_w in T.unroll(num_warps):
-                            reduced_acc += warp_acc[i_w, i_m, lane]
-                        out_idx = i_t * tile_n + lane
-                        if out_idx < n_out:
-                            out[0, token_idx, out_idx] = reduced_acc
-                    if lane == 0 and i_t == 0:
-                        reduced_sqr = T.alloc_var(T.float32, init=0.0)
-                        for i_w in T.unroll(num_warps):
-                            reduced_sqr += warp_sqr[i_w, i_m]
-                        sqrsum[0, token_idx] = reduced_sqr
-
-        if ENABLE_PDL:
-            T.pdl_trigger()
-
-
-@tilelang.jit(
-    pass_configs=pass_configs,
+    pass_configs={
+        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+        tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
+    },
 )
 def hc_head_fuse_tilelang(
     residual,
@@ -748,8 +566,7 @@ def hc_head_fuse_tilelang(
     out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef,valid-type]
 
     with T.Kernel(num_tokens, threads=n_thr) as i:
-        if ENABLE_PDL:
-            T.pdl_sync()
+        T.pdl_sync()
 
         # ------------------------------------------------------------------
         # Pass 1 – for each residual channel m_c and h_block:
@@ -807,5 +624,4 @@ def hc_head_fuse_tilelang(
 
             T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
 
-        if ENABLE_PDL:
-            T.pdl_trigger()
+        T.pdl_trigger()

vllm/benchmarks/datasets/datasets.py

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

vllm/benchmarks/lib/endpoint_request_func.py

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

vllm/config/vllm.py

@@ -1406,7 +1406,7 @@ class VllmConfig:
         # Hybrid KV cache manager (HMA) runtime rules:
         # - Explicit enable (--no-disable-kv-cache-manager): error if runtime
         #   disables it
-        # - No preference: auto-disable for unsupported features (e.g. kv connector)
+        # - No preference: auto-disable for unsupported features or connector configs
         # - Explicit disable (--disable-kv-cache-manager): always respect it
         need_disable_hybrid_kv_cache_manager = False
         # logger should only print warning message for hybrid models. As we
@@ -1438,43 +1438,25 @@ class VllmConfig:
                 need_disable_hybrid_kv_cache_manager = True
 
         if self.scheduler_config.disable_hybrid_kv_cache_manager is None:
-            # Default to disable HMA, but only if the user didn't express a preference.
+            # Auto-disable HMA only when the connector config does not support it.
             if self.kv_transfer_config is not None:
-                from vllm.config.kv_transfer import KVTransferConfig
                 from vllm.distributed.kv_transfer.kv_connector.factory import (
                     KVConnectorFactory,
                 )
-                from vllm.distributed.kv_transfer.kv_connector.v1.base import (
-                    supports_hma,
-                )
 
-                connector_cls = KVConnectorFactory.get_connector_class(
-                    self.kv_transfer_config
-                )
-                all_support_hma = supports_hma(connector_cls)
-                # MultiConnector subclasses SupportsHMA; only effectively
-                # supports HMA when every sub-connector does.
-                if all_support_hma and connector_cls.__name__ == "MultiConnector":
-                    sub_ktcs = self.kv_transfer_config.kv_connector_extra_config.get(
-                        "connectors", []
-                    )
-                    all_support_hma = all(
-                        supports_hma(
-                            KVConnectorFactory.get_connector_class(
-                                KVTransferConfig(**sub)
-                            )
-                        )
-                        for sub in sub_ktcs
-                    )
-                if not all_support_hma:
+                if not KVConnectorFactory.supports_hma_config(self.kv_transfer_config):
                     need_disable_hybrid_kv_cache_manager = True
                     logger.warning(
                         "Turning off hybrid kv cache manager because "
-                        "connector %s does not subclass `SupportsHMA`. "
-                        "This will reduce performance on models with "
-                        "sliding window or Mamba attention. See "
-                        "kv_connector/v1/base.py for details.",
-                        connector_cls.__name__,
+                        "`--kv-transfer-config` selects a KV connector that "
+                        "does not support it. Impact: hybrid SSM models "
+                        "(e.g. Jamba, Bamba) require HMA and will fail at "
+                        "startup without it; models with sliding window "
+                        "attention will run with reduced performance. "
+                        "To add HMA support to a KV connector, subclass "
+                        "`SupportsHMA` defined in kv_connector/v1/base.py "
+                        "(for MultiConnector, all child connectors must "
+                        "support HMA)."
                     )
             self.scheduler_config.disable_hybrid_kv_cache_manager = (
                 need_disable_hybrid_kv_cache_manager

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -5,6 +5,7 @@ import importlib
 from collections.abc import Callable
 from typing import TYPE_CHECKING, cast
 
+from vllm.config.kv_transfer import KVTransferConfig
 from vllm.distributed.kv_transfer.kv_connector.base import (
     KVConnectorBase,
     KVConnectorBaseType,
@@ -18,7 +19,6 @@ from vllm.utils.func_utils import supports_kw
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
-    from vllm.config.kv_transfer import KVTransferConfig
     from vllm.v1.kv_cache_interface import KVCacheConfig
 
 logger = init_logger(__name__)
@@ -53,7 +53,7 @@ class KVConnectorFactory:
 
         # check if the connector supports HMA
         hma_enabled = not config.scheduler_config.disable_hybrid_kv_cache_manager
-        if hma_enabled and not supports_hma(connector_cls):
+        if hma_enabled and not cls.supports_hma_config(kv_transfer_config):
             raise ValueError(
                 f"Connector {connector_cls.__name__} does not support HMA but "
                 f"HMA is enabled. Please set `--disable-hybrid-kv-cache-manager`."
@@ -127,6 +127,23 @@ class KVConnectorFactory:
             raise ValueError(f"Unsupported connector type: {connector_name}")
         return connector_cls
 
+    @classmethod
+    def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool:
+        """Return whether this KV transfer config supports HMA.
+
+        MultiConnector is a special case: the wrapper class implements
+        SupportsHMA, but effective support depends on every configured child.
+        """
+        connector_cls = cls.get_connector_class(kv_transfer_config)
+        if kv_transfer_config.kv_connector != "MultiConnector":
+            return supports_hma(connector_cls)
+
+        from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
+            MultiConnector,
+        )
+
+        return MultiConnector.all_children_support_hma(kv_transfer_config)
+
 
 # Register various connectors here.
 # The registration should not be done in each individual file, as we want to

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

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import contextlib
+import os
 import threading
 import time
 from collections.abc import Iterator
@@ -12,8 +13,7 @@ import regex as re
 import torch
 import zmq
 
-from vllm import envs
-from vllm.config import VllmConfig
+from vllm.config import KVTransferConfig, VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorMetadata,
 )
@@ -162,8 +162,10 @@ class TransferError(MoRIIOError):
     pass
 
 
-def get_moriio_mode() -> MoRIIOMode:
-    read_mode = envs.VLLM_MORIIO_CONNECTOR_READ_MODE
+def get_moriio_mode(kv_transfer_config: KVTransferConfig) -> MoRIIOMode:
+    read_mode = str(
+        kv_transfer_config.kv_connector_extra_config.get("read_mode", "false")
+    ).lower().strip() in ("true", "1")
     logger.debug("MoRIIO Connector read_mode: %s", read_mode)
     if read_mode:
         return MoRIIOMode.READ
@@ -175,6 +177,26 @@ def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int:
     return (dp_rank) * tp_size + tp_rank
 
 
+_DEPRECATED_ENV_VARS: dict[str, str] = {
+    "VLLM_MORIIO_CONNECTOR_READ_MODE": "read_mode",
+    "VLLM_MORIIO_QP_PER_TRANSFER": "qp_per_transfer",
+    "VLLM_MORIIO_POST_BATCH_SIZE": "post_batch_size",
+    "VLLM_MORIIO_NUM_WORKERS": "num_workers",
+}
+
+
+def _warn_deprecated_env_vars() -> None:
+    for env_var, new_key in _DEPRECATED_ENV_VARS.items():
+        if env_var in os.environ:
+            logger.warning_once(
+                "The environment variable %s is deprecated and ignored. "
+                "Set %r inside kv_transfer_config.kv_connector_extra_config "
+                "instead.",
+                env_var,
+                new_key,
+            )
+
+
 @dataclass
 class MoRIIOConfig:
     local_ip: str
@@ -189,6 +211,10 @@ class MoRIIOConfig:
     dp_rank: int
     dp_size: int
     tp_size: int
+    read_mode: bool = False
+    qp_per_transfer: int = 1
+    post_batch_size: int = -1
+    num_workers: int = 1
     backend: str = "rdma"
 
     @classmethod
@@ -201,11 +227,24 @@ class MoRIIOConfig:
         # notify_port       -> For synchronizing stages between prefill and decode
         # handshake_port    -> For initial handshake between mori engine
 
+        # Optional tuning knobs
+        # read_mode        -> If true, run the connector in READ mode (consumer
+        #                     pulls KV from producer) instead of the default
+        #                     WRITE mode.
+
+        # Knobs for RDMA transfers, ignored if on xgmi backend
+        # qp_per_transfer  -> Number of RDMA Queue Pairs per KV transfer.
+        # post_batch_size  -> Batch size for posting transfer work requests
+        #                     (-1 lets the MoRI backend choose).
+        # num_workers      -> Number of background worker threads the MoRI
+        #                     engine uses for transfer processing.
+
         # TODO : merge notify_port and handshake_port to simplify port management
         #        supports non-contiguous ports
         assert vllm_config.kv_transfer_config is not None, (
             "kv_transfer_config must be set for MoRIIOConnector"
         )
+        _warn_deprecated_env_vars()
         kv_transfer_config = vllm_config.kv_transfer_config
         extra_config = kv_transfer_config.kv_connector_extra_config
         tp_rank = get_tensor_model_parallel_rank()
@@ -234,6 +273,10 @@ class MoRIIOConfig:
             dp_rank=dp_rank,
             dp_size=dp_size,
             tp_size=tp_size,
+            read_mode=get_moriio_mode(kv_transfer_config) == MoRIIOMode.READ,
+            qp_per_transfer=int(extra_config.get("qp_per_transfer", 1)),
+            post_batch_size=int(extra_config.get("post_batch_size", -1)),
+            num_workers=int(extra_config.get("num_workers", 1)),
             backend=backend,
         )
 

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

@@ -108,7 +108,7 @@ class MoRIIOConnector(KVConnectorBase_V1):
             + ":"
             + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
         )
-        self.mode = get_moriio_mode()
+        self.mode = get_moriio_mode(self.kv_transfer_config)
         if role == KVConnectorRole.SCHEDULER:
             self.connector_scheduler: MoRIIOConnectorScheduler | None = (
                 MoRIIOConnectorScheduler(vllm_config, self.engine_id)
@@ -250,7 +250,7 @@ class MoRIIOConnectorScheduler:
         self.kv_transfer_config = vllm_config.kv_transfer_config
         self.block_size = vllm_config.cache_config.block_size
         self.engine_id: EngineId = engine_id
-        self.mode = get_moriio_mode()
+        self.mode = get_moriio_mode(self.kv_transfer_config)
         self.host_ip = get_ip()
         self.handshake_port = self.kv_transfer_config.kv_connector_extra_config[
             "handshake_port"
@@ -615,8 +615,11 @@ class MoRIIOConnectorWorker:
                 "is installed and properly configured."
             )
 
+        assert vllm_config.kv_transfer_config is not None
         self.moriio_config = MoRIIOConfig.from_vllm_config(vllm_config)
-        self.mode = get_moriio_mode()
+        self.mode = (
+            MoRIIOMode.READ if self.moriio_config.read_mode else MoRIIOMode.WRITE
+        )
 
         logger.info("Initializing MoRIIO worker %s", engine_id)
 
@@ -700,7 +703,12 @@ class MoRIIOConnectorWorker:
             if self.moriio_config.backend == "xgmi"
             else BackendType.RDMA
         )
-        self.moriio_wrapper.set_backend_type(backend)
+        self.moriio_wrapper.set_backend_type(
+            backend,
+            qp_per_transfer=self.moriio_config.qp_per_transfer,
+            post_batch_size=self.moriio_config.post_batch_size,
+            num_workers=self.moriio_config.num_workers,
+        )
         self.moriio_wrapper.notify_port = self.moriio_config.notify_port
         self.local_kv_cache_metadata: list[bytes] = []
         self.local_kv_cache_size: list[int] = []

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

@@ -8,7 +8,6 @@ import msgpack
 import torch
 import zmq
 
-from vllm import envs
 from vllm.logger import init_logger
 from vllm.utils.network_utils import (
     make_zmq_path,
@@ -16,7 +15,7 @@ from vllm.utils.network_utils import (
 )
 
 if TYPE_CHECKING:
-    pass
+    from mori.io import BackendType
 
 from queue import Empty, Queue
 
@@ -376,7 +375,13 @@ class MoRIIOWrapper:
         )
         self.moriio_engine = moriio_engine
 
-    def set_backend_type(self, backend_type):
+    def set_backend_type(
+        self,
+        backend_type: "BackendType",
+        qp_per_transfer: int = 1,
+        post_batch_size: int = -1,
+        num_workers: int = 1,
+    ) -> None:
         assert self.moriio_engine is not None, "MoRIIO engine must be set first"
         if backend_type == BackendType.XGMI:
             logger.info("Using MoRIIO backend: XGMI")
@@ -385,14 +390,14 @@ class MoRIIOWrapper:
             logger.info(
                 "Using MoRIIO backend: RDMA "
                 "(qp_per_transfer=%d, post_batch_size=%d, num_workers=%d)",
-                envs.VLLM_MORIIO_QP_PER_TRANSFER,
-                envs.VLLM_MORIIO_POST_BATCH_SIZE,
-                envs.VLLM_MORIIO_NUM_WORKERS,
+                qp_per_transfer,
+                post_batch_size,
+                num_workers,
             )
             rdma_cfg = RdmaBackendConfig(
-                envs.VLLM_MORIIO_QP_PER_TRANSFER,
-                envs.VLLM_MORIIO_POST_BATCH_SIZE,
-                envs.VLLM_MORIIO_NUM_WORKERS,
+                qp_per_transfer,
+                post_batch_size,
+                num_workers,
                 PollCqMode.POLLING,
             )
             self.moriio_engine.create_backend(backend_type, rdma_cfg)

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

@@ -19,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorRole,
     KVConnectorWorkerMetadata,
     SupportsHMA,
-    supports_hma,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
     KVConnectorPromMetrics,
@@ -151,6 +150,22 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
                 return True
         return False
 
+    @classmethod
+    def all_children_support_hma(cls, kv_transfer_config: "KVTransferConfig") -> bool:
+        """Return True only if every configured child connector supports HMA."""
+        connectors_config = kv_transfer_config.kv_connector_extra_config.get(
+            "connectors", []
+        )
+        if not connectors_config:
+            return False
+        for conn_config in connectors_config:
+            child_config = KVTransferConfig(
+                **{"engine_id": kv_transfer_config.engine_id, **conn_config}
+            )
+            if not KVConnectorFactory.supports_hma_config(child_config):
+                return False
+        return True
+
     def __init__(
         self,
         vllm_config: "VllmConfig",
@@ -169,7 +184,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
             self._connectors.append(connector_cls(temp_config, role, kv_cache_config))
             self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)
 
-        self._all_support_hma = all(supports_hma(c) for c in self._connectors)
+        assert vllm_config.kv_transfer_config is not None
+        self._all_support_hma = MultiConnector.all_children_support_hma(
+            vllm_config.kv_transfer_config
+        )
         assert (
             vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
             or self._all_support_hma

vllm/entrypoints/cli/serve.py

@@ -308,14 +308,7 @@ def run_multi_api_server(args: argparse.Namespace):
 
     from vllm.v1.engine.utils import get_engine_zmq_addresses
 
-    # Per-API-server ports are picked by the kernel at each child's bind()
-    # to avoid parent-probe vs child-bind TOCTOU; Rust front-end opts out
-    # because it has no port-report-back channel.
-    addresses = get_engine_zmq_addresses(
-        vllm_config,
-        num_api_servers,
-        defer_api_server_ports=not rust_frontend_path,
-    )
+    addresses = get_engine_zmq_addresses(vllm_config, num_api_servers)
 
     with launch_core_engines(
         vllm_config, executor_class, log_stats, addresses, num_api_servers
@@ -348,12 +341,6 @@ def run_multi_api_server(args: argparse.Namespace):
                 tensor_queue=tensor_queue,
             )
 
-            # Forward each child's bound endpoints to the engine handshake
-            # (runs on ``with`` exit).
-            actual_inputs, actual_outputs = api_server_manager.gather_actual_addresses()
-            addresses.inputs = actual_inputs
-            addresses.outputs = actual_outputs
-
     # Wait for API servers.
     try:
         wait_for_completion_or_failure(

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -41,6 +41,7 @@ from vllm.sampling_params import (
     RequestOutputKind,
     SamplingParams,
     StructuredOutputsParams,
+    ThinkingTokenBudget,
 )
 from vllm.utils import random_uuid
 
@@ -225,7 +226,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
             "part of the standard OpenAI API specification."
         ),
     )
-    thinking_token_budget: int | None = None
+    thinking_token_budget: ThinkingTokenBudget = None
     include_reasoning: bool = True
     parallel_tool_calls: bool | None = True
 
@@ -472,27 +473,17 @@ class ChatCompletionRequest(OpenAIBaseModel):
         default_template: str | None,
         default_template_content_format: ChatTemplateContentFormatOption,
     ) -> ChatParams:
-        extra_kwargs: dict[str, Any] = dict(
-            add_generation_prompt=self.add_generation_prompt,
-            continue_final_message=self.continue_final_message,
-            documents=self.documents,
-            reasoning_effort=self.reasoning_effort,
-        )
-
-        # When reasoning is requested, activate thinking for models whose
-        # chat templates require explicit opt-in (e.g., Gemma4 defaults
-        # enable_thinking to false). For templates that don't declare the
-        # variable, resolve_chat_template_kwargs filters it out harmlessly.
-        user_kwargs = self.chat_template_kwargs or {}
-        if self.reasoning_effort is not None and "enable_thinking" not in user_kwargs:
-            extra_kwargs["enable_thinking"] = self.reasoning_effort != "none"
-
         return ChatParams(
             chat_template=self.chat_template or default_template,
             chat_template_content_format=default_template_content_format,
             chat_template_kwargs=merge_kwargs(
                 self.chat_template_kwargs,
-                extra_kwargs,
+                dict(
+                    add_generation_prompt=self.add_generation_prompt,
+                    continue_final_message=self.continue_final_message,
+                    documents=self.documents,
+                    reasoning_effort=self.reasoning_effort,
+                ),
             ),
             media_io_kwargs=self.media_io_kwargs,
         )

vllm/entrypoints/openai/completion/protocol.py

@@ -29,6 +29,7 @@ from vllm.sampling_params import (
     RequestOutputKind,
     SamplingParams,
     StructuredOutputsParams,
+    ThinkingTokenBudget,
 )
 from vllm.utils import random_uuid
 
@@ -185,11 +186,12 @@ class CompletionRequest(OpenAIBaseModel):
         "can detect such behavior and terminate early, saving time and tokens.",
     )
 
-    thinking_token_budget: int | None = Field(
+    thinking_token_budget: ThinkingTokenBudget = Field(
         default=None,
         description=(
             "Maximum number of tokens allowed for thinking operations "
-            "(reasoning models). -1 = unlimited."
+            "(reasoning models). Non-negative integer sets the limit; "
+            "-1 means unlimited (treated as unset)."
         ),
     )
 

vllm/entrypoints/openai/responses/protocol.py

@@ -298,28 +298,17 @@ class ResponsesRequest(OpenAIBaseModel):
         continue_final = should_continue_final_message(self.input)
 
         reasoning = self.reasoning
-        reasoning_effort = None if reasoning is None else reasoning.effort
-
-        extra_kwargs: dict[str, Any] = dict(
-            add_generation_prompt=not continue_final,
-            continue_final_message=continue_final,
-            reasoning_effort=reasoning_effort,
-        )
-
-        # When reasoning is requested, activate thinking for models whose
-        # chat templates require explicit opt-in (e.g., Gemma4 defaults
-        # enable_thinking to false). For templates that don't declare the
-        # variable, resolve_chat_template_kwargs filters it out harmlessly.
-        user_kwargs = self.chat_template_kwargs or {}
-        if reasoning_effort is not None and "enable_thinking" not in user_kwargs:
-            extra_kwargs["enable_thinking"] = reasoning_effort != "none"
 
         return ChatParams(
             chat_template=default_template,
             chat_template_content_format=default_template_content_format,
-            chat_template_kwargs=merge_kwargs(
+            chat_template_kwargs=merge_kwargs(  # To remove unset values
                 self.chat_template_kwargs,
-                extra_kwargs,
+                dict(
+                    add_generation_prompt=not continue_final,
+                    continue_final_message=continue_final,
+                    reasoning_effort=None if reasoning is None else reasoning.effort,
+                ),
             ),
             media_io_kwargs=self.media_io_kwargs,
         )

vllm/envs.py

@@ -216,10 +216,6 @@ if TYPE_CHECKING:
     VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
     VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB: int | None = None
     VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB: int | None = None
-    VLLM_MORIIO_CONNECTOR_READ_MODE: bool = False
-    VLLM_MORIIO_QP_PER_TRANSFER: int = 1
-    VLLM_MORIIO_POST_BATCH_SIZE: int = -1
-    VLLM_MORIIO_NUM_WORKERS: int = 1
     VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480
     VLLM_ENABLE_CUDAGRAPH_GC: bool = False
     VLLM_LOOPBACK_IP: str = ""
@@ -1642,20 +1638,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
         "Use --linear-backend emulation.",
         lambda: bool(int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0"))),
     ),
-    # Controls the read mode for the Mori-IO connector
-    "VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: (
-        os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() in ("true", "1")
-    ),
-    # Controls the QP (Queue Pair) per transfer configuration for the Mori-IO connector
-    "VLLM_MORIIO_QP_PER_TRANSFER": lambda: int(
-        os.getenv("VLLM_MORIIO_QP_PER_TRANSFER", "1")
-    ),
-    # Controls the post-processing batch size for the Mori-IO connector
-    "VLLM_MORIIO_POST_BATCH_SIZE": lambda: int(
-        os.getenv("VLLM_MORIIO_POST_BATCH_SIZE", "-1")
-    ),
-    # Controls the number of workers for Mori operations for the Mori-IO connector
-    "VLLM_MORIIO_NUM_WORKERS": lambda: int(os.getenv("VLLM_MORIIO_NUM_WORKERS", "1")),
     # Timeout (in seconds) for MooncakeConnector in PD disaggregated setup.
     "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int(
         os.getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480")

vllm/model_executor/kernels/mhc/tilelang.py

@@ -5,88 +5,6 @@ import torch
 from vllm.utils.torch_utils import direct_register_custom_op
 
 
-def _torch_hc_prenorm_gemm(
-    x: torch.Tensor,
-    fn: torch.Tensor,
-    out: torch.Tensor,
-    sqrsum: torch.Tensor,
-) -> None:
-    assert out.shape[0] == 1
-    assert sqrsum.shape[0] == 1
-    x_float = x.float()
-    out[0].copy_(x_float @ fn.t())
-    sqrsum[0].copy_(x_float.square().sum(dim=-1))
-
-
-def _tilelang_hc_prenorm_gemm(
-    x: torch.Tensor,
-    fn: torch.Tensor,
-    out: torch.Tensor,
-    sqrsum: torch.Tensor,
-    hidden_size: int,
-    hc_mult: int,
-    tile_n: int = 12,
-    n_thr: int = 512,
-    n_splits: int = 1,
-) -> None:
-    from vllm._tilelang_ops import (
-        hc_prenorm_gemm_block_m_tilelang,
-        hc_prenorm_gemm_tilelang,
-    )
-
-    assert out.shape[0] == n_splits
-    assert sqrsum.shape[0] == n_splits
-    assert x.shape[1] == hc_mult * hidden_size
-    assert x.shape[1] % n_splits == 0
-    assert (x.shape[1] // n_splits) % n_thr == 0
-    use_default_config = tile_n == 12 and n_thr == 512
-    if n_splits == 1 and use_default_config and x.shape[0] >= 1024:
-        hc_prenorm_gemm_block_m_tilelang(
-            x,
-            fn,
-            out,
-            sqrsum,
-            hidden_size,
-            hc_mult,
-            fn.shape[0],
-            n_thr,
-            tile_n,
-            2,
-        )
-        return
-    if (
-        n_splits == 1
-        and use_default_config
-        and x.shape[0] < 128
-        and x.shape[1] % 1024 == 0
-    ):
-        hc_prenorm_gemm_tilelang(
-            x,
-            fn,
-            out,
-            sqrsum,
-            hidden_size,
-            hc_mult,
-            fn.shape[0],
-            1024,
-            4,
-            n_splits,
-        )
-        return
-    hc_prenorm_gemm_tilelang(
-        x,
-        fn,
-        out,
-        sqrsum,
-        hidden_size,
-        hc_mult,
-        fn.shape[0],
-        n_thr,
-        tile_n,
-        n_splits,
-    )
-
-
 def mhc_pre_tilelang(
     residual: torch.Tensor,
     fn: torch.Tensor,
@@ -162,16 +80,10 @@ def mhc_pre_tilelang(
     residual_flat = residual.view(-1, hc_mult, hidden_size)
     num_tokens = residual_flat.shape[0]
 
-    from vllm.utils.deep_gemm import is_deep_gemm_supported
-
-    use_deep_gemm = is_deep_gemm_supported()
-    if use_deep_gemm:
-        # these numbers are from deepgemm kernel impl
-        block_k = 64
-        block_m = 64
-        n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
-    else:
-        n_splits = 1
+    # these numbers are from deepgemm kernel impl
+    block_k = 64
+    block_m = 64
+    n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
 
     post_mix = torch.empty(
         num_tokens, hc_mult, dtype=torch.float32, device=residual.device
@@ -190,24 +102,13 @@ def mhc_pre_tilelang(
         n_splits, num_tokens, dtype=torch.float32, device=residual.device
     )
 
-    residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size)
-    if use_deep_gemm:
-        tf32_hc_prenorm_gemm(
-            residual_2d,
-            fn,
-            gemm_out_mul,
-            gemm_out_sqrsum,
-            n_splits,
-        )
-    else:
-        _tilelang_hc_prenorm_gemm(
-            residual_2d,
-            fn,
-            gemm_out_mul,
-            gemm_out_sqrsum,
-            hidden_size,
-            hc_mult,
-        )
+    tf32_hc_prenorm_gemm(
+        residual_flat.view(num_tokens, hc_mult * hidden_size),
+        fn,
+        gemm_out_mul,
+        gemm_out_sqrsum,
+        n_splits,
+    )
 
     if norm_weight is None:
         mhc_pre_big_fuse_tilelang(
@@ -403,24 +304,16 @@ def mhc_fused_post_pre_tilelang(
     post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
     comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
 
-    from vllm.utils.deep_gemm import is_deep_gemm_supported
-
-    use_deep_gemm = is_deep_gemm_supported()
-    use_small_fma = num_tokens <= 16
-    if use_small_fma:
+    fma_token_threshold = 16
+    if num_tokens <= fma_token_threshold:
         # TODO(gnovack): investigate autotuning these heuristics
         tile_n = 2 if num_tokens < 8 else 3
         n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
     else:
-        if use_deep_gemm:
-            # these number are from deepgemm kernel impl
-            block_k = 64
-            block_m = 64
-            n_splits = compute_num_split(
-                block_k, hc_hidden_size, cdiv(num_tokens, block_m)
-            )
-        else:
-            n_splits = 1
+        # these number are from deepgemm kernel impl
+        block_k = 64
+        block_m = 64
+        n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
 
     gemm_out_mul = torch.empty(
         n_splits,
@@ -455,7 +348,7 @@ def mhc_fused_post_pre_tilelang(
         device=residual.device,
     )
 
-    if use_small_fma:
+    if num_tokens <= fma_token_threshold:
         mhc_fused_tilelang(
             comb_res_mix_flat,
             residual_flat,
@@ -482,26 +375,15 @@ def mhc_fused_post_pre_tilelang(
             residual.shape[-1],
         )
 
-        residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size)
-        if use_deep_gemm:
-            from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
+        from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
 
-            tf32_hc_prenorm_gemm(
-                residual_cur_2d,
-                fn,
-                gemm_out_mul,
-                gemm_out_sqrsum,
-                n_splits,
-            )
-        else:
-            _tilelang_hc_prenorm_gemm(
-                residual_cur_2d,
-                fn,
-                gemm_out_mul,
-                gemm_out_sqrsum,
-                hidden_size,
-                hc_mult,
-            )
+        tf32_hc_prenorm_gemm(
+            residual_cur.view(num_tokens, hc_mult * hidden_size),
+            fn,
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            n_splits,
+        )
 
     if norm_weight is None:
         mhc_pre_big_fuse_tilelang(

vllm/model_executor/layers/fused_moe/config.py

@@ -888,20 +888,33 @@ def int4_w4a16_moe_quant_config(
 def fp8_w8a16_moe_quant_config(
     w1_scale: torch.Tensor,
     w2_scale: torch.Tensor,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
     block_shape: list[int] | None = None,
 ) -> FusedMoEQuantConfig:
     """
     Construct a quant config for 16-bit float activations and fp8 weights.
     """
     group_shape = GroupShape(*block_shape) if block_shape is not None else None
+    fp8_dtype = current_platform.fp8_dtype()
     return FusedMoEQuantConfig(
         _a1=FusedMoEQuantDesc(),
         _a2=FusedMoEQuantDesc(),
         _w1=FusedMoEQuantDesc(
-            current_platform.fp8_dtype(), group_shape, w1_scale, None, None
+            fp8_dtype,
+            group_shape,
+            w1_scale,
+            None,
+            None,
+            w1_bias,
         ),
         _w2=FusedMoEQuantDesc(
-            current_platform.fp8_dtype(), group_shape, w2_scale, None, None
+            fp8_dtype,
+            group_shape,
+            w2_scale,
+            None,
+            None,
+            w2_bias,
         ),
     )
 
@@ -911,6 +924,8 @@ def int8_w8a16_moe_quant_config(
     w2_scale: torch.Tensor,
     w1_zp: torch.Tensor | None,
     w2_zp: torch.Tensor | None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
     block_shape: list[int] | None = None,
 ) -> FusedMoEQuantConfig:
     """
@@ -920,8 +935,8 @@ def int8_w8a16_moe_quant_config(
     return FusedMoEQuantConfig(
         _a1=FusedMoEQuantDesc(shape=group_shape),
         _a2=FusedMoEQuantDesc(shape=group_shape),
-        _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp),
-        _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp),
+        _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias),
+        _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias),
     )
 
 

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

@@ -0,0 +1,229 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU INT4 W4A8 dynamic quantized fused MoE experts."""
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+    RoutingMethodType,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kInt4W4A8StaticGroup32Sym,
+    kInt4W4A8StaticGroup64Sym,
+    kInt4W4A8StaticGroup128Sym,
+    kInt4W4A8StaticGroupSym,
+)
+from vllm.platforms import current_platform
+
+
+class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
+    """CPU INT4 W4A8 dynamic quantized monolithic MoE experts.
+
+    Uses the dynamic_4bit_int_moe kernel for efficient 4-bit weight,
+    8-bit activation MoE inference on CPU.
+    """
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ):
+        super().__init__(
+            moe_config,
+            quant_config,
+        )
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        """Expects unquantized inputs (quantization happens in kernel)."""
+        return True
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        return current_platform.is_cpu()
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        """Does not support no_act_and_mul (requires SwiGLU or SiLU)."""
+        return False
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        """Supports SiLU and SwiGLU variants."""
+        return activation in (
+            MoEActivation.SILU,
+            MoEActivation.SWIGLUOAI,
+            MoEActivation.SWIGLUSTEP,
+        )
+
+    @staticmethod
+    def _supports_parallel_config(
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> bool:
+        """Currently does not support expert parallelism."""
+        # Based on compressed_tensors implementation check
+        return moe_parallel_config.ep_size == 1
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        """Supports INT4 weights with INT8 dynamic activations.
+
+        This is W4A8 with:
+        - Weights: 4-bit integer (stored as int8, packed to uint8 nibbles)
+          Can be channel-wise or group-wise quantization
+        - Activations: dynamic per-token 8-bit integer quantization
+        """
+        # group size must be multiple of 32
+        SUPPORTED_W_A = [
+            (kInt4W4A8StaticGroup128Sym, None),
+            (kInt4W4A8StaticGroup64Sym, None),
+            (kInt4W4A8StaticGroup32Sym, None),
+            (kInt4W4A8StaticGroupSym, None),
+        ]
+        return (weight_key, activation_key) in SUPPORTED_W_A
+
+    @staticmethod
+    def _supports_routing_method(
+        routing_method: RoutingMethodType,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        """Supports standard routing methods."""
+        return routing_method in [
+            RoutingMethodType.Default,
+            RoutingMethodType.Renormalize,
+            RoutingMethodType.RenormalizeNaive,
+        ]
+
+    @staticmethod
+    def _supports_router_logits_dtype(
+        router_logits_dtype: torch.dtype | None,
+        routing_method: RoutingMethodType,
+    ) -> bool:
+        """Accepts any router logits dtype."""
+        return True
+
+    def supports_expert_map(self) -> bool:
+        """Expert parallelism not yet supported."""
+        return False
+
+    def _activation_kind(self, activation: MoEActivation) -> int:
+        """Convert MoEActivation to kernel activation kind integer.
+
+        Returns:
+            0 = SwiGLU_Gu (SiLU(g)*u)
+            1 = SwiGLU_Ug (SiLU(u)*g)
+            2 = SiLU
+        """
+        if activation == MoEActivation.SWIGLUSTEP:
+            return 0
+        if activation == MoEActivation.SWIGLUOAI:
+            return 1
+        if activation == MoEActivation.SILU:
+            return 2
+        raise ValueError(f"Unsupported activation '{activation}'")
+
+    def apply(
+        self,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,  # w13_weight_packed
+        w2: torch.Tensor,  # w2_weight_packed
+        router_logits: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        # grouped topk + fused topk bias parameters
+        num_expert_group: int | None = None,
+        e_score_correction_bias: torch.Tensor | None = None,
+        routed_scaling_factor: float | None = None,
+        topk_group: int | None = None,
+    ) -> torch.Tensor:
+        """Apply the monolithic 4-bit INT MoE forward pass.
+
+        Args:
+            hidden_states: Input tensor [num_tokens, hidden_size]
+            w1: Packed w13 weights (w1+w3 gated weights)
+            w2: Packed w2 weights (down projection)
+            router_logits: Router output logits [num_tokens, num_experts]
+            activation: Activation function type
+            global_num_experts: Total number of experts
+            expert_map: Expert mapping for EP (not supported)
+            a1q_scale: Activation quantization scale (not used, dynamic)
+            apply_router_weight_on_input: Whether to apply routing on input
+            num_expert_group: For grouped topk
+            e_score_correction_bias: Bias for expert scores
+            routed_scaling_factor: Scaling factor for routing
+            topk_group: Group size for topk
+
+        Returns:
+            Output tensor after MoE computation
+        """
+        from vllm.model_executor.layers.fused_moe.cpu_fused_moe import (
+            select_experts,
+        )
+
+        renormalize = self.moe_config.routing_method in (
+            RoutingMethodType.Renormalize,
+            RoutingMethodType.RenormalizeNaive,
+        )
+
+        # TODO(bnell): this could be factored into a CPURouter class and
+        # turn this into a modular kernel
+        # Perform topk selection
+        topk_weights, topk_ids = select_experts(
+            hidden_states=hidden_states,
+            router_logits=router_logits,
+            use_grouped_topk=num_expert_group is not None,
+            top_k=self.moe_config.experts_per_token,
+            renormalize=renormalize,
+            topk_group=topk_group,
+            num_expert_group=num_expert_group,
+            scoring_func="softmax",
+            routed_scaling_factor=(
+                routed_scaling_factor if routed_scaling_factor is not None else 1.0
+            ),
+            e_score_correction_bias=e_score_correction_bias,
+        )
+
+        # Extract dimensions from weight tensors
+        # w1 is w13_packed: [num_experts, packed_data...]
+        # w2 is w2_packed: [num_experts, packed_data...]
+        # These dimensions should be available from the layer
+        # For now, we'll extract from moe_config
+        K = self.moe_config.hidden_dim
+        N = self.moe_config.intermediate_size_per_partition
+        assert self.quant_config.block_shape is not None
+        if self.quant_config.is_per_act_token:
+            group_size = -1
+        else:
+            group_size = self.quant_config.block_shape[1]
+
+        # Call the dynamic 4-bit int MoE kernel
+        return torch.ops._C.dynamic_4bit_int_moe(
+            hidden_states,
+            topk_ids.to(torch.long),
+            topk_weights,
+            w1,  # w13_weight_packed
+            w2,  # w2_weight_packed
+            K,  # hidden_size (w2_out_features)
+            N,  # intermediate_size (w2_in_features)
+            N * 2,  # 2*intermediate_size (w13_out_features)
+            group_size,
+            apply_router_weight_on_input,
+            self._activation_kind(activation),
+        )

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

@@ -99,6 +99,9 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
     ) -> bool:
         return True
 
+    def supports_chunking(self) -> bool:
+        return False
+
     def supports_expert_map(self) -> bool:
         return False
 
@@ -137,6 +140,5 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
             intermediate_size=self.intermediate_size_per_partition,
             local_expert_offset=self.ep_rank * self.local_num_experts,
             local_num_experts=self.local_num_experts,
-            routed_scaling_factor=routed_scaling_factor,
             routing_method_type=self.routing_method_type,
         )

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

@@ -88,6 +88,9 @@ class TrtLlmFp8ExpertsBase:
             or moe_parallel_config.use_ag_rs_all2all_kernels
         ) and not moe_parallel_config.enable_eplb
 
+    def supports_chunking(self) -> bool:
+        return False
+
     def supports_expert_map(self) -> bool:
         return False
 

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

@@ -113,6 +113,9 @@ class TrtLlmMxfp4ExpertsBase:
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
+    def supports_chunking(self) -> bool:
+        return False
+
     def supports_expert_map(self) -> bool:
         return False
 

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

@@ -157,27 +157,8 @@ class TrtLlmNvFp4ExpertsBase:
     def activation_format() -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
 
-    def _get_chunk_size(self) -> int:
-        MAX_GRID_Y = 65535
-        MAX_TILE_TOKENS_DIM = 128
-
-        def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int:
-            """Calculates the max number of supported tokens, so the CUDA grid.Y limit
-            won't be reached.
-            Based on getMaxNumCtasInBatchDim function in flashinfer's TRTLLM MoE runner:
-            https://github.com/flashinfer-ai/flashinfer/blob/719ee23fd82cb220d51ad118ca60198718f6c9d1/include/flashinfer/trtllm/fused_moe/runner.h#L97
-            Which given numTokens, topK, numExperts, tileTokensDim calculates maxNumCtas
-            which is used as the CUDA grid.Y dimension, which we want to
-            be <= MAX_GRID_Y. Solving for numTokens gives the formula below.
-            """
-            return (
-                num_experts + (MAX_GRID_Y - num_experts + 1) * MAX_TILE_TOKENS_DIM - 1
-            ) // top_k
-
-        # Using 305k or more causes IMA error in the kernel, so limit to 300k.
-        return min(
-            300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts)
-        )
+    def supports_chunking(self) -> bool:
+        return False
 
     def supports_expert_map(self) -> bool:
         return False
@@ -218,7 +199,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceNoOP()
 
-    def _invoke_kernel(
+    def apply(
         self,
         output: torch.Tensor,
         hidden_states: torch.Tensor,
@@ -228,10 +209,18 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
         topk_ids: torch.Tensor,
         activation: MoEActivation,
         global_num_experts: int,
-        a1q_scale: torch.Tensor,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
     ):
         import flashinfer
 
+        assert self._supports_activation(activation)
+        assert a1q_scale is not None
         assert self.quant_config.w1_scale is not None
         assert self.quant_config.w2_scale is not None
 
@@ -273,57 +262,6 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
             output=output,
         )
 
-    def apply(
-        self,
-        output: torch.Tensor,
-        hidden_states: torch.Tensor,
-        w1: torch.Tensor,
-        w2: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        activation: MoEActivation,
-        global_num_experts: int,
-        expert_map: torch.Tensor | None,
-        a1q_scale: torch.Tensor | None,
-        a2_scale: torch.Tensor | None,
-        workspace13: torch.Tensor,
-        workspace2: torch.Tensor,
-        expert_tokens_meta: mk.ExpertTokensMetadata | None,
-        apply_router_weight_on_input: bool,
-    ):
-        assert self._supports_activation(activation)
-        assert a1q_scale is not None
-
-        M = hidden_states.shape[0]
-        chunk_size = self._get_chunk_size()
-
-        if chunk_size >= M:
-            self._invoke_kernel(
-                output,
-                hidden_states,
-                w1,
-                w2,
-                topk_weights,
-                topk_ids,
-                activation,
-                global_num_experts,
-                a1q_scale,
-            )
-        else:
-            for start in range(0, M, chunk_size):
-                end = min(start + chunk_size, M)
-                self._invoke_kernel(
-                    output[start:end],
-                    hidden_states[start:end],
-                    w1,
-                    w2,
-                    topk_weights[start:end],
-                    topk_ids[start:end],
-                    activation,
-                    global_num_experts,
-                    a1q_scale[start:end],
-                )
-
 
 class TrtLlmNvFp4ExpertsMonolithic(
     TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -421,6 +421,7 @@ def select_fp8_moe_backend(
 
 def convert_to_fp8_moe_kernel_format(
     fp8_backend: Fp8MoeBackend,
+    # TODO(bnell): replace layer with weight_block_size
     layer: torch.nn.Module,
     w13: torch.Tensor,
     w2: torch.Tensor,
@@ -508,6 +509,8 @@ def make_fp8_moe_quant_config(
     w2_scale: torch.Tensor,
     a1_scale: torch.Tensor | None,
     a2_scale: torch.Tensor | None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
     block_shape: list[int] | None = None,
     per_act_token_quant: bool = False,
     per_out_ch_quant: bool = False,
@@ -526,19 +529,13 @@ def make_fp8_moe_quant_config(
     a method of the modular kernel itself.
     """
 
-    # MARLIN is mixed precision W8A16 config.
-    if fp8_backend == Fp8MoeBackend.MARLIN:
-        return fp8_w8a16_moe_quant_config(
-            w1_scale=w1_scale,
-            w2_scale=w2_scale,
-            block_shape=block_shape,
-        )
-
-    # CPU is mixed precision W8A16 config.
-    if fp8_backend == Fp8MoeBackend.CPU:
+    # MARLIN and CPU are mixed precision W8A16 config.
+    if fp8_backend == Fp8MoeBackend.MARLIN or fp8_backend == Fp8MoeBackend.CPU:
         return fp8_w8a16_moe_quant_config(
             w1_scale=w1_scale,
             w2_scale=w2_scale,
+            w1_bias=w1_bias,
+            w2_bias=w2_bias,
             block_shape=block_shape,
         )
 
@@ -549,6 +546,8 @@ def make_fp8_moe_quant_config(
         return fp8_w8a8_moe_quant_config(
             w1_scale=w1_scale,
             w2_scale=w2_scale,
+            w1_bias=w1_bias,
+            w2_bias=w2_bias,
             a1_scale=a1_scale,
             a2_scale=a2_scale,
             a1_gscale=(1.0 / a1_scale),
@@ -566,6 +565,8 @@ def make_fp8_moe_quant_config(
             "mxfp8",
             w1_scale=w1_scale,
             w2_scale=w2_scale,
+            w1_bias=w1_bias,
+            w2_bias=w2_bias,
             a1_scale=a1_scale,
             a2_scale=a2_scale,
             block_shape=block_shape,
@@ -577,6 +578,8 @@ def make_fp8_moe_quant_config(
     return fp8_w8a8_moe_quant_config(
         w1_scale=w1_scale,
         w2_scale=w2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
         a1_scale=a1_scale,
         a2_scale=a2_scale,
         block_shape=block_shape,

vllm/model_executor/layers/fused_moe/oracle/int8.py

@@ -147,6 +147,8 @@ def make_int8_moe_quant_config(
     w2_scale: torch.Tensor,
     a1_scale: torch.Tensor | None = None,
     a2_scale: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
     per_act_token_quant: bool = False,
 ) -> FusedMoEQuantConfig:
     assert (a1_scale is None and a2_scale is None) or (
@@ -159,6 +161,8 @@ def make_int8_moe_quant_config(
             w2_scale=w2_scale,
             w1_zp=None,
             w2_zp=None,
+            w1_bias=w1_bias,
+            w2_bias=w2_bias,
         )
 
     return int8_w8a8_moe_quant_config(
@@ -166,6 +170,8 @@ def make_int8_moe_quant_config(
         w2_scale=w2_scale,
         a1_scale=a1_scale,
         a2_scale=a2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
         per_act_token_quant=per_act_token_quant,
     )
 

vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from enum import Enum
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.config.kernel import MoEBackend
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
+    maybe_make_prepare_finalize,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+    FusedMoEQuantDesc,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    QuantKey,
+)
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+class W4A8Int8MoeBackend(Enum):
+    CPU_INT4 = "CPU_INT4"
+
+
+def _get_priority_backends(
+    moe_config: FusedMoEConfig,
+) -> list[W4A8Int8MoeBackend]:
+    """
+    Get available backends in priority order based on platform and config.
+
+    Currently only CPU INT4 backend is available for W4A8 INT8 MoE.
+    """
+    if current_platform.is_cpu():
+        return [W4A8Int8MoeBackend.CPU_INT4]
+    return []
+
+
+def backend_to_kernel_cls(
+    backend: W4A8Int8MoeBackend,
+) -> list[type[mk.FusedMoEExperts]]:
+    """Map W4A8Int8MoeBackend to kernel class."""
+    if backend == W4A8Int8MoeBackend.CPU_INT4:
+        from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import (
+            CPUExpertsInt4,
+        )
+
+        return [CPUExpertsInt4]
+
+    else:
+        raise ValueError(f"Unknown W4A8 Int8 MoE backend: {backend.value}")
+
+
+def map_w4a8_int8_backend(runner_backend: MoEBackend) -> W4A8Int8MoeBackend:
+    """Map user's MoEBackend to W4A8Int8MoeBackend."""
+    mapping = {
+        "cpu": W4A8Int8MoeBackend.CPU_INT4,
+    }
+    if backend := mapping.get(runner_backend):
+        return backend
+    raise ValueError(
+        f"moe_backend='{runner_backend}' is not supported for W4A8 Int8 MoE. "
+        f"Expected one of {list(mapping.keys())}."
+    )
+
+
+def select_w4a8_int8_moe_backend(
+    config: FusedMoEConfig,
+    weight_key: QuantKey | None,
+    activation_key: QuantKey | None,
+) -> tuple[W4A8Int8MoeBackend, type[mk.FusedMoEExperts]]:
+    """
+    Select the primary W4A8 Int8 MoE backend.
+
+    Args:
+        config: MoE configuration
+        weight_key: Weight quantization key (should be one of kInt4W4A8Static*)
+        activation_key: Activation quantization key (currently unused for W4A8)
+
+    Returns:
+        Tuple of (backend, kernel_class)
+    """
+
+    AVAILABLE_BACKENDS = _get_priority_backends(config)
+
+    if not AVAILABLE_BACKENDS:
+        raise NotImplementedError("W4A8 Int8 MoE is only supported on CPU platforms")
+
+    activation_format = (
+        mk.FusedMoEActivationFormat.BatchedExperts
+        if config.moe_parallel_config.use_batched_activation_format
+        else mk.FusedMoEActivationFormat.Standard
+    )
+
+    def _make_log_backend(backend: W4A8Int8MoeBackend) -> str:
+        available_backend_strs = [b.value for b in AVAILABLE_BACKENDS]
+        return (
+            f"Using {backend.value} W4A8 Int8 MoE backend out "
+            f"of potential backends: {available_backend_strs}."
+        )
+
+    def _make_log_unsupported(backend: W4A8Int8MoeBackend, reason: str | None) -> str:
+        if reason:
+            return (
+                f"W4A8 Int8 MoE backend {backend.value} does not support the "
+                f"deployment configuration since {reason}."
+            )
+        else:
+            return (
+                f"W4A8 Int8 MoE backend '{backend.value}' does not support the "
+                "deployment configuration."
+            )
+
+    def _return_or_raise(
+        backend: W4A8Int8MoeBackend,
+    ) -> tuple[W4A8Int8MoeBackend, type[mk.FusedMoEExperts]]:
+        reason = None
+        for k_cls in backend_to_kernel_cls(backend):
+            supported, reason = k_cls.is_supported_config(
+                k_cls, config, weight_key, activation_key, activation_format
+            )
+            if supported:
+                logger.info_once(_make_log_backend(backend))
+                return backend, k_cls
+        raise ValueError(_make_log_unsupported(backend, reason))
+
+    # Handle explicit moe_backend from user.
+    runner_backend = config.moe_backend
+    if runner_backend != "auto":
+        requested_backend = map_w4a8_int8_backend(runner_backend)
+        return _return_or_raise(requested_backend)
+
+    # Select kernels in order of backend.
+    for backend in AVAILABLE_BACKENDS:
+        for k_cls in backend_to_kernel_cls(backend):
+            supported, reason = k_cls.is_supported_config(
+                k_cls,
+                config,
+                weight_key,
+                activation_key,
+                activation_format,
+            )
+            if supported:
+                logger.info_once(_make_log_backend(backend))
+                return backend, k_cls
+            else:
+                logger.debug_once(_make_log_unsupported(backend, reason))
+
+    raise NotImplementedError(
+        "No W4A8 Int8 MoE backend supports the deployment configuration."
+    )
+
+
+def make_w4a8_int8_moe_quant_config(
+    block_shape: tuple[int, int] | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Create FusedMoEQuantConfig for W4A8 Int8 MoE.
+
+    Args:
+        block_shape: Quantization block shape (row, col).
+                    For channel-wise: (-1, 1) or None
+                    For group-wise: (1, group_size)
+
+    Returns:
+        FusedMoEQuantConfig with appropriate settings for W4A8 Int8
+    """
+    # W4A8 Int8 uses static weight quantization, dynamic activation quantization
+    # Weights are 4-bit (stored as int8, packed to uint8),
+    # activations are dynamically quantized to 8-bit in kernel
+
+    group_shape = GroupShape(*block_shape) if block_shape is not None else None
+
+    return FusedMoEQuantConfig(
+        # Activations: unquantized (FP/BF16), dynamically quantized in kernel
+        _a1=FusedMoEQuantDesc(shape=group_shape),
+        _a2=FusedMoEQuantDesc(shape=group_shape),
+        # Weights: INT8 (4-bit values), pre-packed with scales
+        # dtype=None means already quantized/packed
+        _w1=FusedMoEQuantDesc(dtype=None, shape=group_shape),
+        _w2=FusedMoEQuantDesc(dtype=None, shape=group_shape),
+    )
+
+
+def pack_int4_weights_for_kleidi(
+    int4_as_int8: torch.Tensor,
+    scales: torch.Tensor,
+    bias: torch.Tensor | None,
+    in_features: int,
+    out_features: int,
+    group_size: int,
+) -> torch.Tensor:
+    """
+    Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.
+
+    Args:
+        int4_as_int8: [out, in] int8 tensor with values in [-8, 7]
+        scales: [out, in//group_size] or [out, 1] for channel-wise
+        bias: [out] optional bias
+        in_features: Input dimension
+        out_features: Output dimension
+        group_size: Quantization group size (-1 for channel-wise)
+
+    Returns:
+        Packed weight tensor in KleidiAI format
+    """
+    # Shift to unsigned nibble [0, 15]
+    tmp = int4_as_int8.add(8)
+    # Pack pairs along input dimension
+    uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(torch.uint8)
+
+    # Determine scale dtype based on group_size
+    # KleidiAI groupwise kernels accept bfloat16 scales
+    # KleidiAI channelwise kernels accept float32 scales
+    scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
+    scales_typed = scales.to(scale_dtype)
+    bias_typed = None if bias is None else bias.to(torch.float32)
+
+    # Pack using KleidiAI op
+    actual_group_size = in_features if group_size == -1 else group_size
+    return torch.ops.aten._dyn_quant_pack_4bit_weight(
+        uint8_nibbles,
+        scales_typed,
+        bias_typed,
+        actual_group_size,
+        in_features,
+        out_features,
+    )
+
+
+def convert_to_w4a8_int8_moe_format(
+    w13_weight: torch.Tensor,
+    w2_weight: torch.Tensor,
+    w13_weight_scale: torch.Tensor,
+    w2_weight_scale: torch.Tensor,
+    group_size: int,
+    w13_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor | None,
+    torch.Tensor | None,
+]:
+    """
+    Pack INT4 MoE weights to KleidiAI format.
+
+    This function packs the INT4 weights (stored as int8 values) into
+    the format expected by the KleidiAI dynamic_4bit_int_moe kernel.
+
+    Args:
+        w13_weight: [E, 2*IN, H] int8 tensor (int4 values in [-8,7])
+        w2_weight: [E, H, IN] int8 tensor (int4 values in [-8,7])
+        w13_weight_scale: [E, 2*IN, H/g or 1] scale tensor
+        w2_weight_scale: [E, H, IN/g or 1] scale tensor
+        group_size: Quantization group size (-1 for channel-wise)
+        w13_bias: Optional [E, 2*IN] bias tensor
+        w2_bias: Optional [E, H] bias tensor
+
+    Returns:
+        Tuple of (w13_packed, w2_packed) tensors
+    """
+    # Derive dimensions from tensor shapes
+    E = w13_weight.shape[0]  # num_experts
+    I2 = w13_weight.shape[1]  # w13_out_features (2*IN)
+    H = w13_weight.shape[2]  # w13_in_features (hidden_size)
+    IN = w2_weight.shape[2]  # w2_in_features (intermediate_size)
+    w2_out_features = w2_weight.shape[1]  # Should equal H
+
+    # Pack per expert
+    w13_packed_list = []
+    w2_packed_list = []
+
+    for e in range(E):
+        w13_packed_list.append(
+            pack_int4_weights_for_kleidi(
+                w13_weight[e],  # [2I, H]
+                w13_weight_scale[e],  # [2I, H/g or 1]
+                w13_bias[e] if w13_bias is not None else None,  # [2I]
+                H,
+                I2,
+                group_size,
+            )
+        )
+        w2_packed_list.append(
+            pack_int4_weights_for_kleidi(
+                w2_weight[e],  # [H, IN]
+                w2_weight_scale[e],  # [H, IN/g or 1]
+                w2_bias[e] if w2_bias is not None else None,  # [H]
+                IN,
+                w2_out_features,  # in_features=IN, out_features=H
+                group_size,
+            )
+        )
+
+    # Stack all experts
+    w13_packed = torch.stack(w13_packed_list, dim=0)
+    w2_packed = torch.stack(w2_packed_list, dim=0)
+    empty = torch.empty(0)
+
+    return w13_packed, w2_packed, empty, empty, empty, empty
+
+
+def make_w4a8_int8_moe_kernel(
+    moe_quant_config: FusedMoEQuantConfig,
+    moe_config: FusedMoEConfig,
+    experts_cls: type[mk.FusedMoEExperts],
+    routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+) -> mk.FusedMoEKernel:
+    """
+    Create FusedMoEKernel for W4A8 Int8 MoE.
+
+    Args:
+        moe_quant_config: Quantization configuration
+        moe_config: MoE configuration
+        experts_cls: Expert kernel class (should be CPUExpertsInt4)
+        routing_tables: Optional routing tables for expert parallelism
+
+    Returns:
+        Configured FusedMoEKernel instance
+    """
+    # Create Prepare/Finalize.
+    prepare_finalize = maybe_make_prepare_finalize(
+        moe=moe_config,
+        quant_config=moe_quant_config,
+        routing_tables=routing_tables,
+        allow_new_interface=True,
+        use_monolithic=issubclass(experts_cls, mk.FusedMoEExpertsMonolithic),
+    )
+    assert prepare_finalize is not None
+
+    logger.info_once("Using %s", prepare_finalize.__class__.__name__)
+
+    # Create Experts.
+    # W4A8 Int8 currently only supports monolithic interface
+    if not issubclass(experts_cls, mk.FusedMoEExpertsMonolithic):
+        raise ValueError(
+            f"W4A8 Int8 MoE only supports monolithic experts, "
+            f"but got {experts_cls.__name__}"
+        )
+
+    experts = experts_cls(
+        moe_config=moe_config,
+        quant_config=moe_quant_config,
+    )
+
+    kernel = mk.FusedMoEKernel(
+        prepare_finalize,
+        experts,
+        inplace=not moe_config.disable_inplace,
+    )
+
+    return kernel

vllm/model_executor/layers/mhc.py

@@ -3,12 +3,8 @@
 import torch
 
 # this import will also register the custom ops
-# import vllm.model_executor.kernels.mhc  # noqa: F401
 import vllm.model_executor.kernels.mhc as mhc_kernels
 from vllm.model_executor.custom_op import CustomOp
-from vllm.utils.import_utils import has_tilelang
-
-HAS_TILELANG = has_tilelang()
 
 
 # --8<-- [start:mhc_pre]
@@ -89,52 +85,6 @@ class MHCPreOp(CustomOp):
         #         sinkhorn_repeat,
         #     )
         # else:
-        if HAS_TILELANG:
-            return torch.ops.vllm.mhc_pre_tilelang(
-                residual,
-                fn,
-                hc_scale,
-                hc_base,
-                rms_eps,
-                hc_pre_eps,
-                hc_sinkhorn_eps,
-                hc_post_mult_value,
-                sinkhorn_repeat,
-                n_splits,
-                norm_weight,
-                norm_eps,
-            )
-        else:
-            return self.forward_native(
-                residual,
-                fn,
-                hc_scale,
-                hc_base,
-                rms_eps,
-                hc_pre_eps,
-                hc_sinkhorn_eps,
-                hc_post_mult_value,
-                sinkhorn_repeat,
-                n_splits,
-                norm_weight,
-                norm_eps,
-            )
-
-    def forward_native(
-        self,
-        residual: torch.Tensor,
-        fn: torch.Tensor,
-        hc_scale: torch.Tensor,
-        hc_base: torch.Tensor,
-        rms_eps: float,
-        hc_pre_eps: float,
-        hc_sinkhorn_eps: float,
-        hc_post_mult_value: float,
-        sinkhorn_repeat: int,
-        n_splits: int = 1,
-        norm_weight: torch.Tensor | None = None,
-        norm_eps: float = 0.0,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         return mhc_kernels.mhc_pre_torch(
             residual,
             fn,
@@ -147,6 +97,9 @@ class MHCPreOp(CustomOp):
             sinkhorn_repeat,
         )
 
+    def forward_native(self, *args, **kwargs):
+        raise NotImplementedError("Native implementation of mhc_pre is not available")
+
 
 # --8<-- [start:mhc_post]
 @CustomOp.register("mhc_post")
@@ -194,20 +147,6 @@ class MHCPostOp(CustomOp):
         #         comb_res_mix,
         #     )
         # else:
-        if HAS_TILELANG:
-            return torch.ops.vllm.mhc_post_tilelang(
-                x, residual, post_layer_mix, comb_res_mix
-            )
-        else:
-            return self.forward_native(x, residual, post_layer_mix, comb_res_mix)
-
-    def forward_native(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor,
-        post_layer_mix: torch.Tensor,
-        comb_res_mix: torch.Tensor,
-    ) -> torch.Tensor:
         return mhc_kernels.mhc_post_torch(
             x,
             residual,
@@ -215,6 +154,9 @@ class MHCPostOp(CustomOp):
             comb_res_mix,
         )
 
+    def forward_native(self, *args, **kwargs):
+        raise NotImplementedError("Native implementation of mhc_post is not available")
+
 
 # --8<-- [start:hc_head]
 @CustomOp.register("hc_head")
@@ -278,32 +220,17 @@ class HCHeadOp(CustomOp):
         out = torch.empty(
             num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
         )
-
-        if HAS_TILELANG:
-            torch.ops.vllm.hc_head_fused_kernel_tilelang(
-                hs_flat,
-                hc_fn,
-                hc_scale,
-                hc_base,
-                out,
-                hidden_size,
-                rms_norm_eps,
-                hc_eps,
-                hc_mult,
-            )
-        else:
-            torch.ops.vllm.hc_head_triton(
-                hs_flat,
-                hc_fn,
-                hc_scale,
-                hc_base,
-                out,
-                hidden_size,
-                rms_norm_eps,
-                hc_eps,
-                hc_mult,
-            )
-
+        torch.ops.vllm.hc_head_triton(
+            hs_flat,
+            hc_fn,
+            hc_scale,
+            hc_base,
+            out,
+            hidden_size,
+            rms_norm_eps,
+            hc_eps,
+            hc_mult,
+        )
         return out.view(*outer_shape, hidden_size)
 
     def forward_native(self, *args, **kwargs):
@@ -363,42 +290,9 @@ class MHCFusedPostPreOp(CustomOp):
             norm_eps,
         )
 
-    def forward_hip(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor,
-        post_layer_mix: torch.Tensor,
-        comb_res_mix: torch.Tensor,
-        fn: torch.Tensor,
-        hc_scale: torch.Tensor,
-        hc_base: torch.Tensor,
-        rms_eps: float,
-        hc_pre_eps: float,
-        hc_sinkhorn_eps: float,
-        hc_post_mult_value: float,
-        sinkhorn_repeat: int,
-        n_splits: int = 1,
-        tile_n: int = 1,
-        norm_weight: torch.Tensor | None = None,
-        norm_eps: float = 0.0,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        return torch.ops.vllm.mhc_fused_post_pre_tilelang(
-            x,
-            residual,
-            post_layer_mix,
-            comb_res_mix,
-            fn,
-            hc_scale,
-            hc_base,
-            rms_eps,
-            hc_pre_eps,
-            hc_sinkhorn_eps,
-            hc_post_mult_value,
-            sinkhorn_repeat,
-            n_splits,
-            tile_n,
-            norm_weight,
-            norm_eps,
+    def forward_hip(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Hip implementation of mhc_fused_post_pre is not available"
         )
 
     def forward_native(self, *args, **kwargs):

vllm/model_executor/layers/quantization/__init__.py

@@ -84,7 +84,7 @@ def register_quantization_config(quantization: str):
 
     def _wrapper(quant_config_cls):
         if quantization in QUANTIZATION_METHODS:
-            logger.warning(
+            logger.debug(
                 "The quantization method '%s' already exists and will be "
                 "overwritten by the quantization config %s.",
                 quantization,

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

@@ -11,16 +11,26 @@ from compressed_tensors.quantization import (
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
     RoutedExperts,
+    SharedExperts,
 )
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEQuantConfig,
 )
-from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts
+from vllm.model_executor.layers.fused_moe.oracle.w4a8_int8 import (
+    convert_to_w4a8_int8_moe_format,
+    make_w4a8_int8_moe_kernel,
+    make_w4a8_int8_moe_quant_config,
+    select_w4a8_int8_moe_backend,
+)
 from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa E501
     CompressedTensorsMoEMethod,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    QuantKey,
+    ScaleDesc,
+)
 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
 from vllm.platforms import CpuArchEnum, current_platform
 
@@ -48,6 +58,11 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
         self.has_bias = self.moe.has_bias
         self.weight_quant = weight_quant
         self.input_quant = input_quant
+        self.static_input_scales = False  # always dynamic per token
+        # Weight can be channel-wise (group_size=None) or group-wise
+        self.group_size = (
+            weight_quant.group_size if (weight_quant.group_size is not None) else -1
+        )
 
         # Validate scheme: weights=W4 (channel or group),
         # activations=dynamic TOKEN (A8)
@@ -61,17 +76,9 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
                 "W4A8-int MoE needs dynamic per-token activation quantization."
             )
 
-        # Weight can be channel-wise (group_size=None) or group-wise
-        self.group_size = (
-            weight_quant.group_size if (weight_quant.group_size is not None) else -1
-        )
         if weight_quant.num_bits != 4:
             raise ValueError("This method only supports 4-bit weights (num_bits=4).")
 
-        # CPU only
-        if not current_platform.is_cpu():
-            raise ValueError("CompressedTensorsW4A8Int8MoEMethod is CPU-only.")
-
         # Arm: check _dyn ops availability
         if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
             try:
@@ -82,7 +89,26 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
                     f"""PyTorch {torch.__version__} lacks _dyn_quant_* 4bit ops;
                     install a newer build."""
                 ) from err
-        self.static_input_scales = False  # always dynamic per token
+
+        # Construct QuantKey for weights from QuantizationArgs
+        # W4A8 INT4: 4-bit weights (stored as int8), static quantization
+        if self.group_size == -1:
+            # Channel-wise quantization
+            group_shape = GroupShape(-1, 1)
+            scale_dtype = torch.float32
+        else:
+            # Group-wise quantization
+            group_shape = GroupShape(1, self.group_size)
+            scale_dtype = torch.bfloat16
+
+        weight_scale_desc = ScaleDesc(scale_dtype, static=True, group_shape=group_shape)
+        weight_key = QuantKey(torch.int8, weight_scale_desc, symmetric=True)
+
+        self.backend, self.experts_cls = select_w4a8_int8_moe_backend(
+            moe,
+            weight_key,
+            activation_key=None,  # unquantized inputs
+        )
 
     # ---- parameter creation ----
     def create_weights(
@@ -182,72 +208,20 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
 
     # post-load packing to dyn-4bit KleidiAI kernel's format
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        E = layer.w13_weight.shape[0]
-        H = layer.w13_in_features
-        I2 = layer.w13_out_features
-        IN = layer.w2_in_features
-        g = layer.group_size
-
-        def _pack_matrix(
-            int4_as_int8_2d: torch.Tensor,
-            scales_2d: torch.Tensor,
-            bias_1d: torch.Tensor | None,
-            in_features: int,
-            out_features: int,
-        ) -> torch.Tensor:
-            # int4 values are stored as int8 in [-8,7].
-            # Shift to unsigned nibble and pack pairs along input-dim.
-            tmp = int4_as_int8_2d.add(8)  # [out, in]
-            uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(
-                torch.uint8
-            )  # [out, in//2]
-
-            # KleidiAI groupwise kernels accepts float32 scales
-            # KleidiAI groupwise kernels accepts bfloat16 scales
-            scale_dtype = torch.float32 if g == -1 else torch.bfloat16
-            scales = scales_2d.to(scale_dtype)
-            bias = None if bias_1d is None else bias_1d.to(torch.float32)
-            return torch.ops.aten._dyn_quant_pack_4bit_weight(
-                uint8_nibbles,
-                scales,
-                bias,
-                g if g != -1 else in_features,
-                in_features,
-                out_features,
+        # Use oracle to pack weights.
+        w13_packed, w2_packed, w13_weight_scale, w2_weight_scale, w13_bias, w2_bias = (
+            convert_to_w4a8_int8_moe_format(
+                w13_weight=layer.w13_weight,
+                w2_weight=layer.w2_weight,
+                w13_weight_scale=layer.w13_weight_scale,
+                w2_weight_scale=layer.w2_weight_scale,
+                group_size=self.group_size,
+                w13_bias=layer.w13_bias if self.has_bias else None,
+                w2_bias=layer.w2_bias if self.has_bias else None,
             )
+        )
 
-        # Pack per expert
-        w13_packed_list = []
-        w2_packed_list = []
-
-        has_w13_bias = hasattr(layer, "w13_bias") and layer.w13_bias is not None
-        has_w2_bias = hasattr(layer, "w2_bias") and layer.w2_bias is not None
-
-        for e in range(E):
-            w13_packed_list.append(
-                _pack_matrix(
-                    layer.w13_weight[e],  # [2I, H]
-                    layer.w13_weight_scale[e],  # [2I, H/g or 1]
-                    layer.w13_bias[e] if has_w13_bias else None,  # [2I]
-                    H,
-                    I2,
-                )
-            )
-            w2_packed_list.append(
-                _pack_matrix(
-                    # w2 shape is [H, IN]; we need [out, in] == [H, IN].
-                    layer.w2_weight[e],  # [H, IN]
-                    layer.w2_weight_scale[e],  # [H, IN/g or 1]
-                    layer.w2_bias[e] if has_w2_bias else None,  # [H]
-                    IN,
-                    layer.w2_out_features,  # in_features=IN, out_features=H
-                )
-            )
-
-        # each packed tensor has identical shape per expert; stack on dim 0
-        w13_packed = torch.stack(w13_packed_list, dim=0)
-        w2_packed = torch.stack(w2_packed_list, dim=0)
-
+        # Register packed weights as parameters
         replace_parameter(
             layer,
             "w13_weight_packed",
@@ -259,7 +233,6 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
             torch.nn.Parameter(w2_packed, requires_grad=False),
         )
 
-        # free raw tensors/scales/bias now that they're packed into the payload.
         replace_parameter(
             layer, "w13_weight", torch.nn.Parameter(torch.empty(0), requires_grad=False)
         )
@@ -269,36 +242,46 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
         replace_parameter(
             layer,
             "w13_weight_scale",
-            torch.nn.Parameter(torch.empty(0), requires_grad=False),
+            torch.nn.Parameter(w13_weight_scale, requires_grad=False),
         )
         replace_parameter(
             layer,
             "w2_weight_scale",
-            torch.nn.Parameter(torch.empty(0), requires_grad=False),
+            torch.nn.Parameter(w2_weight_scale, requires_grad=False),
         )
-        if has_w13_bias:
+
+        if self.has_bias:
             replace_parameter(
                 layer,
                 "w13_bias",
-                torch.nn.Parameter(torch.empty(0), requires_grad=False),
+                torch.nn.Parameter(w13_bias, requires_grad=False),
             )
-        if has_w2_bias:
+        if self.has_bias:
             replace_parameter(
                 layer,
                 "w2_bias",
-                torch.nn.Parameter(torch.empty(0), requires_grad=False),
+                torch.nn.Parameter(w2_bias, requires_grad=False),
             )
 
+        quant_config = self.get_fused_moe_quant_config(layer)
+        assert quant_config is not None
+        assert self.experts_cls is not None
+        self.moe_kernel = make_w4a8_int8_moe_kernel(
+            moe_quant_config=quant_config,
+            moe_config=self.moe,
+            experts_cls=self.experts_cls,
+            routing_tables=layer._expert_routing_tables(),
+        )
+
     def get_fused_moe_quant_config(
         self, layer: torch.nn.Module
     ) -> FusedMoEQuantConfig | None:
-        # CPU dynamic 4-bit MoE path does not use modular kernels or
-        # fused_experts; quant config is not needed.
-        return None
+        # Determine block shape from group_size
+        # group_size=-1 means channel-wise: (-1, 1)
+        # group_size=N means group-wise: (1, N)
+        block_shape = (-1, 1) if self.group_size == -1 else (1, self.group_size)
 
-    @property
-    def is_monolithic(self) -> bool:
-        return True
+        return make_w4a8_int8_moe_quant_config(block_shape=block_shape)
 
     def apply_monolithic(
         self,
@@ -307,43 +290,43 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
         router_logits: torch.Tensor,
         input_ids: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        assert layer.activation in (
-            MoEActivation.SILU,
-            MoEActivation.SWIGLUOAI,
-            MoEActivation.SWIGLUSTEP,
-        ), "Only SiLU/SwiGLUGU/SwiGLUUG are supported."
-        assert layer.expert_map is None, """expert_map/EP not implemented
-        for CPU dyn-4bit MoE."""
-
-        def _act_kind(s: MoEActivation) -> int:
-            # 0 = SwiGLU_Gu (SiLU(g)*u), 1 = SwiGLU_Ug (SiLU(u)*g), 2 = SiLU
-            if s == MoEActivation.SWIGLUSTEP:
-                return 0
-            if s == MoEActivation.SWIGLUOAI:
-                return 1
-            if s == MoEActivation.SILU:
-                return 2
-            raise ValueError(f"Unknown activation '{s}'")
-
-        # Apply topk softmax on router output
-        topk_weights, topk_ids = select_experts(
-            hidden_states=x,
-            router_logits=router_logits,
-            top_k=layer.top_k,
-            use_grouped_topk=layer.use_grouped_topk,
-            renormalize=layer.renormalize,
-        )
-
-        return torch.ops._C.dynamic_4bit_int_moe(
+        assert self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply_monolithic(
             x,
-            topk_ids.to(torch.long),
-            topk_weights,
             layer.w13_weight_packed,
             layer.w2_weight_packed,
-            layer.w2_out_features,
-            layer.w2_in_features,
-            layer.w13_out_features,
-            layer.group_size,
-            layer.apply_router_weight_on_input,
-            int(_act_kind(layer.activation)),
+            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: RoutedExperts,
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts: SharedExperts | None,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor:
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply(
+            x,
+            layer.w13_weight_packed,
+            layer.w2_weight_packed,
+            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=shared_experts,
+            shared_experts_input=shared_experts_input,
         )

vllm/model_executor/layers/quantization/modelopt.py

@@ -23,13 +23,10 @@ from vllm.model_executor.layers.fused_moe import (
     FusedMoEMethodBase,
     FusedMoEQuantConfig,
     FusedMoeWeightScaleSupported,
-    MoEActivation,
     RoutedExperts,
-    RoutingMethodType,
     SharedExperts,
 )
 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
-    Fp8MoeBackend,
     convert_to_fp8_moe_kernel_format,
     make_fp8_moe_kernel,
     make_fp8_moe_quant_config,
@@ -70,7 +67,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
     MXFP8_BLOCK_SIZE,
     MXFP8_SCALE_DTYPE,
     MXFP8_VALUE_DTYPE,
-    mxfp8_e4m3_quantize,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
@@ -85,6 +81,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
     requantize_with_max_scale,
 )
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.parameter import (
     BlockQuantScaleParameter,
     ChannelQuantScaleParameter,
@@ -93,7 +90,6 @@ from vllm.model_executor.parameter import (
     PerTensorScaleParameter,
 )
 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
-from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe
 
 if TYPE_CHECKING:
     from vllm.model_executor.models.utils import WeightsMapper
@@ -187,7 +183,7 @@ class ModelOptQuantConfigBase(QuantizationConfig):
 
         # handle exclusion
         if self.is_layer_excluded(prefix):
-            if isinstance(layer, LinearBase):
+            if isinstance(layer, (LinearBase, ParallelLMHead)):
                 return UnquantizedLinearMethod()
             return None
 
@@ -200,7 +196,7 @@ class ModelOptQuantConfigBase(QuantizationConfig):
             return UnquantizedLinearMethod()
 
         # now, the layer is quantized, handle it here
-        if isinstance(layer, LinearBase):
+        if isinstance(layer, (LinearBase, ParallelLMHead)):
             quant_method = self.LinearMethodCls(self)
             if getattr(quant_method, "backend", "") == "marlin":
                 quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
@@ -1860,10 +1856,11 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
         moe_config: FusedMoEConfig,
     ) -> None:
         super().__init__(moe_config)
+        self.weight_block_size = [1, MXFP8_BLOCK_SIZE]
         self.quant_config = quant_config
         assert self.quant_config.is_checkpoint_mxfp8_serialized
 
-        self.mxfp8_backend, _ = select_mxfp8_moe_backend(self.moe)
+        self.mxfp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)
 
     def create_weights(
         self,
@@ -2059,12 +2056,41 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
         )
 
     def process_weights_after_loading(self, layer: RoutedExperts) -> None:
+        # TODO(bnell): why is this required only for mxfp8?
         if getattr(layer, "_already_called_process_weights_after_loading", False):
             return
+        layer._already_called_process_weights_after_loading = True
 
         self._check_weight_dtypes(layer)
-        self._shuffle_weights_for_trtllm(layer)
-        layer._already_called_process_weights_after_loading = True
+
+        layer.weight_block_size = self.weight_block_size
+
+        w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
+            fp8_backend=self.mxfp8_backend,
+            layer=layer,
+            w13=layer.w13_weight,
+            w2=layer.w2_weight,
+            w13_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            w13_input_scale=None,
+            w2_input_scale=None,
+        )
+
+        replace_parameter(layer, "w13_weight", w13)
+        replace_parameter(layer, "w2_weight", w2)
+        replace_parameter(layer, "w13_weight_scale", w13_scale)
+        replace_parameter(layer, "w2_weight_scale", w2_scale)
+
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        assert self.moe_quant_config is not None
+        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.mxfp8_backend,
+            experts_cls=self.experts_cls,
+            routing_tables=layer._expert_routing_tables(),
+        )
 
     def maybe_make_prepare_finalize(
         self,
@@ -2088,12 +2114,14 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
     def get_fused_moe_quant_config(
         self, layer: RoutedExperts
     ) -> FusedMoEQuantConfig | None:
-        # TRTLLM MXFP8 path is monolithic and does not use modular kernel config.
-        return None
-
-    @property
-    def is_monolithic(self) -> bool:
-        return self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
+        return make_fp8_moe_quant_config(
+            fp8_backend=self.mxfp8_backend,
+            w1_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            a1_scale=None,
+            a2_scale=None,
+            block_shape=self.weight_block_size,
+        )
 
     def apply_monolithic(
         self,
@@ -2102,83 +2130,23 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
         router_logits: torch.Tensor,
         input_ids: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        from flashinfer.fused_moe.core import (
-            ActivationType,
-            Fp8QuantizationType,
-        )
-
-        assert self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM
-
-        if layer.eplb_state is not None:
-            raise NotImplementedError(
-                "EPLB is not supported for FlashInfer TRTLLM MXFP8 MoE backend."
-            )
-
-        supported_activations = [MoEActivation.SILU]
-        if layer.activation not in supported_activations:
-            raise NotImplementedError(
-                "FlashInfer TRTLLM MXFP8 MoE supports only "
-                f"{supported_activations}, got {layer.activation}."
-            )
-
-        # Map vLLM MoEActivation to FlashInfer ActivationType.
-        activation_map = {
-            MoEActivation.SILU: ActivationType.Swiglu,
-            MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
-        }
-        fi_activation_type: ActivationType = activation_map[layer.activation]
-
-        # DeepSeekV3 routing requires float32 logits; others expect bfloat16.
-        if layer.routing_method_type == RoutingMethodType.DeepSeekV3:
-            assert router_logits.dtype == torch.float32, (
-                "DeepSeekV3 routing requires float32 router_logits, "
-                f"got {router_logits.dtype}."
-            )
-        else:
-            router_logits = router_logits.to(torch.bfloat16)
-
-        # Treat 0 as "unset" for compatibility with ungrouped routing configs.
-        n_group = layer.num_expert_group or None
-        topk_group = layer.topk_group or None
-
-        hidden_states_mxfp8, hidden_states_scale = mxfp8_e4m3_quantize(
+        assert self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply_monolithic(
             x,
-            is_sf_swizzled_layout=False,
-        )
-
-        kwargs: dict = dict(
-            routing_logits=router_logits,
-            routing_bias=layer.e_score_correction_bias,
-            hidden_states=hidden_states_mxfp8,
-            hidden_states_scale=hidden_states_scale,
-            gemm1_weights=layer.w13_weight,
-            gemm1_weights_scale=layer.w13_weight_scale,
-            gemm2_weights=layer.w2_weight,
-            gemm2_weights_scale=layer.w2_weight_scale,
-            num_experts=layer.global_num_experts,
-            top_k=layer.top_k,
-            # Keep Optional semantics: FlashInfer expects None for non-grouped
-            # routing (e.g. Qwen3 Renormalize), not 0.
-            n_group=n_group,
-            topk_group=topk_group,
-            intermediate_size=layer.intermediate_size_per_partition,
-            local_expert_offset=layer.ep_rank * layer.local_num_experts,
-            local_num_experts=layer.local_num_experts,
+            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,
-            routing_method_type=layer.routing_method_type,
-            use_shuffled_weight=True,
-            weight_layout=0,
-            fp8_quantization_type=Fp8QuantizationType.MxFp8,
         )
 
-        if fi_activation_type != ActivationType.Swiglu:
-            raise NotImplementedError(
-                "FlashInfer TRTLLM MXFP8 MoE supports only Swiglu activation, "
-                f"got {fi_activation_type}."
-            )
-
-        return flashinfer_trtllm_fp8_block_scale_moe(**kwargs)
-
     def apply(
         self,
         layer: RoutedExperts,
@@ -2189,8 +2157,19 @@ class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
         shared_experts_input: torch.Tensor | None,
     ) -> torch.Tensor:
         assert not self.is_monolithic
-        raise NotImplementedError(
-            "Non-monolithic MXFP8 MoE path is not yet implemented."
+        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=shared_experts,
+            shared_experts_input=shared_experts_input,
         )
 
 
@@ -2393,13 +2372,13 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
 
         # Excluded layers
         if self.is_layer_excluded(prefix):
-            if isinstance(layer, LinearBase):
+            if isinstance(layer, (LinearBase, ParallelLMHead)):
                 return UnquantizedLinearMethod()
             return None
 
         quant_algo = self._resolve_quant_algo(prefix)
 
-        if isinstance(layer, LinearBase):
+        if isinstance(layer, (LinearBase, ParallelLMHead)):
             if quant_algo == "FP8":
                 return ModelOptFp8LinearMethod(self.fp8_config)
             if quant_algo == "NVFP4":

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

@@ -371,19 +371,18 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
         a1_scale = layer.w13_input_scale
         a2_scale = layer.w2_input_scale
 
-        quant_config = make_fp8_moe_quant_config(
+        return 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,
+            w1_bias=getattr(layer, "w13_bias", None),
+            w2_bias=getattr(layer, "w2_bias", None),
             block_shape=self.weight_block_size,
             swiglu_limit=getattr(layer, "swiglu_limit", None),
         )
 
-        self._maybe_inject_biases(quant_config, layer)
-        return quant_config
-
 
 class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
     """Online tensorwise FP8 MoE quantization.

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

@@ -105,9 +105,9 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
     def get_fused_moe_quant_config(
         self, layer: torch.nn.Module
     ) -> "FusedMoEQuantConfig | None":
-        quant_config = make_int8_moe_quant_config(
+        return make_int8_moe_quant_config(
             w1_scale=layer.w13_scale,
             w2_scale=layer.w2_scale,
+            w1_bias=getattr(layer, "w13_bias", None),
+            w2_bias=getattr(layer, "w2_bias", None),
         )
-        self._maybe_inject_biases(quant_config, layer)
-        return quant_config

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

@@ -8,7 +8,6 @@ import torch
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.model_executor.layers.fused_moe import (
     FusedMoEMethodBase,
-    FusedMoEQuantConfig,
     RoutedExperts,
     SharedExperts,
 )
@@ -101,21 +100,6 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         pass
 
-    def _maybe_inject_biases(
-        self,
-        quant_config: FusedMoEQuantConfig,
-        layer: torch.nn.Module,
-    ) -> None:
-        """Inject biases into the quant config if the model has them
-        (e.g. GPT-OSS biased MoE)."""
-        if 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
-
     def maybe_make_prepare_finalize(
         self,
         routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

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

@@ -214,19 +214,18 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
         a1_scale = layer.w13_input_scale
         a2_scale = layer.w2_input_scale
 
-        quant_config = make_fp8_moe_quant_config(
+        return 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,
+            w1_bias=getattr(layer, "w13_bias", None),
+            w2_bias=getattr(layer, "w2_bias", None),
             block_shape=self.weight_block_size,
             swiglu_limit=getattr(layer, "swiglu_limit", None),
         )
 
-        self._maybe_inject_biases(quant_config, layer)
-        return quant_config
-
     def process_weights_after_loading(self, layer: Module) -> None:
         if getattr(layer, "_already_called_process_weights_after_loading", False):
             return

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

@@ -181,6 +181,31 @@ kInt8Static = QuantKey(INT8_DTYPE, scale=kInt8StaticGroupScale, symmetric=True)
 kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)
 kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)
 
+# INT4 W4A8 quantization keys
+
+# For group-wise quantization (e.g., group_size=128)
+# Note: group_size will be specified at runtime, this is a generic group scale
+kInt4W4A8StaticGroupScale128 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 128))
+kInt4W4A8StaticGroup128Sym = QuantKey(
+    torch.int8, kInt4W4A8StaticGroupScale128, symmetric=True
+)
+
+kInt4W4A8StaticGroupScale64 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 64))
+kInt4W4A8StaticGroup64Sym = QuantKey(
+    torch.int8, kInt4W4A8StaticGroupScale64, symmetric=True
+)
+
+kInt4W4A8StaticGroupScale32 = ScaleDesc(torch.bfloat16, True, GroupShape(1, 32))
+kInt4W4A8StaticGroup32Sym = QuantKey(
+    torch.int8, kInt4W4A8StaticGroupScale32, symmetric=True
+)
+
+# Generic group-wise with flexible group size (per-token groups)
+kInt4W4A8StaticGroupScale = ScaleDesc(torch.bfloat16, True, GroupShape(1, -1))
+kInt4W4A8StaticGroupSym = QuantKey(
+    torch.int8, kInt4W4A8StaticGroupScale, symmetric=True
+)
+
 
 def create_fp8_quant_key(
     static: bool,

vllm/model_executor/layers/vocab_parallel_embedding.py

@@ -290,6 +290,7 @@ class VocabParallelEmbedding(PluggableLayer):
 
         if params_dtype is None:
             params_dtype = torch.get_default_dtype()
+        self.params_dtype = params_dtype
         # Divide the weight matrix along the vocabulary dimension.
         self.num_added_embeddings = self.num_embeddings - self.org_vocab_size
         self.num_embeddings_per_partition = divide(
@@ -438,6 +439,12 @@ class VocabParallelEmbedding(PluggableLayer):
         # If parameter does not have output dim, then it should
         # be copied onto all gpus (e.g. g_idx for act_order gptq).
         if output_dim is None:
+            if (
+                loaded_weight.ndim == 0
+                and param.data.ndim == 1
+                and param.data.numel() == 1
+            ):
+                loaded_weight = loaded_weight.reshape(1)
             assert param.data.shape == loaded_weight.shape
             param.data.copy_(loaded_weight)
             return

vllm/model_executor/model_loader/weight_utils.py

@@ -1090,8 +1090,7 @@ def runai_safetensors_weights_iterator(
             mininterval=2,
         )
 
-        for name, tensor in tensor_iter:
-            yield name, tensor.clone()
+        yield from tensor_iter
 
 
 def _init_fastsafetensors_loader(

vllm/model_executor/models/deepseek_eagle3.py

@@ -407,6 +407,12 @@ class Eagle3DeepseekV2ForCausalLM(DeepseekV2ForCausalLM):
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
         # Combine multiple auxiliary hidden states returned by Eagle3
+        if self.model.fc_norm is not None:
+            chunks = hidden_states.chunk(self.model.num_aux_hidden_states, dim=-1)
+            hidden_states = torch.cat(
+                [norm(chunk) for norm, chunk in zip(self.model.fc_norm, chunks)],
+                dim=-1,
+            )
         return self.model.fc(hidden_states)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

vllm/model_executor/models/internlm2.py

@@ -41,6 +41,7 @@ from vllm.sequence import IntermediateTensors
 from .interfaces import SupportsLoRA, SupportsPP
 from .interfaces_base import default_pooling_type
 from .utils import (
+    AutoWeightsLoader,
     StageMissingLayer,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
@@ -308,6 +309,42 @@ class InternLM2Model(nn.Module):
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("gate_up_proj", "w1", 0),
+            ("gate_up_proj", "w3", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            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:
+                # 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 = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
 
 class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
     packed_modules_mapping = {
@@ -368,40 +405,11 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("gate_up_proj", "w1", 0),
-            ("gate_up_proj", "w3", 1),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-            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:
-                # 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 = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["output."] if self.config.tie_word_embeddings else None),
+        )
+        return loader.load_weights(weights)
 
 
 @default_pooling_type(tok_pooling_type="ALL")

vllm/model_executor/models/nemotron_h.py

@@ -875,6 +875,7 @@ class NemotronHForCausalLM(
         self.lm_head = ParallelLMHead(
             config.vocab_size,
             config.hidden_size,
+            quant_config=self.quant_config,
             prefix=maybe_prefix(prefix, "lm_head"),
         )
 

vllm/model_executor/models/qwen3_5.py

@@ -477,6 +477,7 @@ class Qwen3_5ForCausalLMBase(
                 self.lm_head = ParallelLMHead(
                     config.vocab_size,
                     config.hidden_size,
+                    quant_config=self.quant_config,
                     prefix=maybe_prefix(prefix, "lm_head"),
                 )
         else:

vllm/model_executor/models/qwen3_5_mtp.py

@@ -381,6 +381,7 @@ class Qwen3_5MTP(nn.Module, SupportsMultiModal):
                 self.lm_head = ParallelLMHead(
                     config.vocab_size,
                     config.hidden_size,
+                    quant_config=self.quant_config,
                     prefix=maybe_prefix(prefix, "lm_head"),
                 )
         else:

vllm/model_executor/models/registry.py

@@ -980,7 +980,7 @@ class _ModelRegistry:
             raise TypeError(msg)
 
         if model_arch in self.models:
-            logger.warning(
+            logger.debug(
                 "Model architecture %s is already registered, and will be "
                 "overwritten by the new model class %s.",
                 model_arch,

vllm/models/deepseek_v4/amd/model.py

@@ -54,7 +54,6 @@ from vllm.models.deepseek_v4.attention import (
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.utils.import_utils import has_tilelang
 
 
 class DeepseekV4MLP(nn.Module):
@@ -474,7 +473,6 @@ class DeepseekV4DecoderLayer(nn.Module):
         self.mhc_pre = MHCPreOp()
         self.mhc_post = MHCPostOp()
         self.mhc_fused_post_pre = MHCFusedPostPreOp()
-        self.has_tilelang = has_tilelang()
 
     def hc_pre(
         self,
@@ -505,7 +503,7 @@ class DeepseekV4DecoderLayer(nn.Module):
     ):
         return self.mhc_post(x, residual, post, comb)
 
-    def _forward_fused_post_pre(
+    def _forward_cuda(
         self,
         x: torch.Tensor,
         positions: torch.Tensor,
@@ -557,7 +555,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         x = self.ffn(x, input_ids)
         return x, residual, post_mix, res_mix
 
-    def _forward_unfused_post_pre(
+    def _forward_rocm(
         self,
         x: torch.Tensor,
         positions: torch.Tensor,
@@ -596,13 +594,12 @@ class DeepseekV4DecoderLayer(nn.Module):
     ) -> tuple[
         torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
     ]:
-        if not self.has_tilelang:
-            return self._forward_unfused_post_pre(
+        if current_platform.is_rocm():
+            return self._forward_rocm(
                 x, positions, input_ids, post_mix, res_mix, residual
             )
-        return self._forward_fused_post_pre(
-            x, positions, input_ids, post_mix, res_mix, residual
-        )
+
+        return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual)
 
 
 @support_torch_compile
@@ -685,7 +682,6 @@ class DeepseekV4Model(nn.Module):
             requires_grad=False,
         )
         self.hc_head_op = HCHeadOp()
-        self.has_tilelang = has_tilelang()
         # Pre-hc_head residual stream buffer for the MTP draft. Stable
         # address (outside the cudagraph pool) so the copy_ in forward()
         # refreshes it correctly across captured shapes.
@@ -752,7 +748,7 @@ class DeepseekV4Model(nn.Module):
                 res_mix,
                 residual,
             )
-        if layer is not None and self.has_tilelang:
+        if layer is not None and current_platform.is_cuda():
             hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix)
 
         if not get_pp_group().is_last_rank:

vllm/models/deepseek_v4/amd/mtp.py

@@ -39,7 +39,6 @@ from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weigh
 from vllm.model_executor.models.utils import maybe_prefix
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.utils.import_utils import has_tilelang
 
 from .model import DeepseekV4DecoderLayer
 
@@ -119,7 +118,6 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         )
 
         self.hc_head_op = HCHeadOp()
-        self.has_tilelang = has_tilelang()
 
     def forward(
         self,
@@ -146,7 +144,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
         hidden_states, residual, post_mix, res_mix = self.mtp_block(
             positions=positions, x=hidden_states, input_ids=None
         )
-        if self.has_tilelang:
+        if current_platform.is_cuda():
             hidden_states = self.mtp_block.hc_post(
                 hidden_states, residual, post_mix, res_mix
             )

vllm/models/deepseek_v4/common/ops/__init__.py

@@ -10,7 +10,6 @@ from .cache_utils import (
 from .fused_indexer_q import MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant
 from .fused_inv_rope_fp8_quant import fused_inv_rope_fp8_quant
 from .fused_qk_rmsnorm import fused_q_kv_rmsnorm
-from .save_partial_states import save_partial_states
 
 __all__ = [
     "MXFP4_BLOCK_SIZE",
@@ -21,5 +20,4 @@ __all__ = [
     "fused_inv_rope_fp8_quant",
     "fused_q_kv_rmsnorm",
     "quantize_and_insert_k_cache",
-    "save_partial_states",
 ]

vllm/models/deepseek_v4/common/ops/fused_compress_quant_cache.py

@@ -11,6 +11,12 @@ Three specialized kernels:
   - _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn:
         head=128, MXFP4 (block=32), 4 ue8m0 bytes
 
+Additional cutedsl kernels:
+  - _compress_kv_sparse_attn_cutedsl / _norm_rope_insert_sparse_attn_cutedsl:
+        CuTe DSL split kernels for C128
+  - _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl:
+        CuTe DSL fused kernels for C4
+
 RoPE is register-based via tl.reshape -> tl.split -> tl.interleave (or the
 even/odd halves are consumed directly for MXFP4, no interleave needed).
 FP8 UE8M0 quant uses tl.reshape to tile [N_QUANT_BLOCKS, QUANT_BLOCK] for
@@ -19,92 +25,42 @@ even/odd halves, producing (N_QUANT_BLOCKS, MXFP4_BLOCK/2) packed nibbles
 and N_QUANT_BLOCKS ue8m0 bytes.
 """
 
-from typing import Any
-
-import torch
+from functools import cache
 
 from vllm.triton_utils import tl, triton
 
 from .fused_indexer_q import _fp32x2_to_fp4x2
 
 
-def compress_norm_rope_store_triton(
-    state_cache: torch.Tensor,
-    num_actual: int,
-    token_to_req_indices: torch.Tensor,
-    positions: torch.Tensor,
-    slot_mapping: torch.Tensor,
-    block_table: torch.Tensor,
-    block_size: int,
-    state_width: int,
-    cos_sin_cache: torch.Tensor,
-    kv_cache: torch.Tensor,
-    k_cache_metadata: Any,
-    pdl_kwargs: dict,
-    head_dim: int,
-    rope_head_dim: int,
-    compress_ratio: int,
-    overlap: bool,
-    use_fp4_cache: bool,
-    rms_norm_weight: torch.Tensor,
-    rms_norm_eps: float,
-    quant_block: int,
-    token_stride: int,
-    scale_dim: int,
-) -> None:
-    """Shared triton launcher for the fused compress+norm+RoPE+insert path.
-
-    Picks one of the three kernels in this module based on ``head_dim`` and
-    ``use_fp4_cache``. Identical launch signature for all three.
-    """
-    if head_dim == 512:
-        kernel = _fused_kv_compress_norm_rope_insert_sparse_attn
-        num_warps = 4
-    elif use_fp4_cache:
-        kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
-        num_warps = 1
-    else:
-        kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
-        num_warps = 1
-
-    kernel[(num_actual,)](
-        # state cache
-        state_cache,
-        state_cache.stride(0),
-        state_cache.stride(1),
-        # metadata
-        token_to_req_indices,
-        positions,
-        slot_mapping,
-        block_table,
-        block_table.stride(0),
-        block_size,
-        # RMSNorm
-        rms_norm_weight,
-        rms_norm_eps,
-        # RoPE
-        cos_sin_cache,
-        cos_sin_cache.stride(0),
-        # KV cache
-        kv_cache,
-        k_cache_metadata.slot_mapping,
-        kv_cache.shape[1],  # paged KV cache block size (tokens per block)
-        # constexprs
-        HEAD_SIZE=head_dim,
-        TRITON_BLOCK_SIZE=triton.next_power_of_2(head_dim),
-        STATE_WIDTH=state_width,
-        COMPRESS_RATIO=compress_ratio,
-        OVERLAP=overlap,
-        ROPE_HEAD_DIM=rope_head_dim,
-        FP8_MAX=448.0,
-        QUANT_BLOCK=quant_block,
-        TOKEN_STRIDE=token_stride,
-        SCALE_DIM=scale_dim,
-        KV_BLOCK_STRIDE=kv_cache.stride(0),
-        num_warps=num_warps,
-        **pdl_kwargs,
+@cache
+def _get_sparse_attn_cutedsl_impls():
+    from .sparse_attn_compress_cutedsl import (
+        _compress_kv_sparse_attn_cutedsl,
+        _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
+        _norm_rope_insert_sparse_attn_cutedsl,
     )
 
+    return (
+        _compress_kv_sparse_attn_cutedsl,
+        _norm_rope_insert_sparse_attn_cutedsl,
+        _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
+    )
+
+
+def _compress_kv_sparse_attn_cutedsl(*args, **kwargs):
+    """CuTe DSL sparse-attention compress wrapper."""
+    return _get_sparse_attn_cutedsl_impls()[0](*args, **kwargs)
+
+
+def _norm_rope_insert_sparse_attn_cutedsl(*args, **kwargs):
+    """CuTe DSL RMSNorm/RoPE/FP8-store wrapper."""
+    return _get_sparse_attn_cutedsl_impls()[1](*args, **kwargs)
+
+
+def _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(*args, **kwargs):
+    """CuTe DSL fused C4 sparse-attention compressor wrapper."""
+    return _get_sparse_attn_cutedsl_impls()[2](*args, **kwargs)
+
 
 # =============================================================================
 # DeepseekV4 Attention path (head=512, nope=448 FP8 + rope=64 bf16)

vllm/models/deepseek_v4/common/ops/save_partial_states.py

@@ -1,101 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-
-from vllm.triton_utils import tl, triton
-
-
-def save_partial_states(
-    kv: torch.Tensor,
-    score: torch.Tensor,
-    ape: torch.Tensor,
-    positions: torch.Tensor,
-    state_cache: torch.Tensor,
-    slot_mapping: torch.Tensor,
-    block_size: int,
-    state_width: int,
-    compress_ratio: int,
-    pdl_kwargs: dict | None = None,
-) -> None:
-    """Write packed [kv, score+ape] partial states into the compressor cache.
-
-    One program per token; pads (slot_id == -1) are skipped.
-    """
-    num_actual = slot_mapping.shape[0]
-    head_size = kv.shape[-1]
-    _save_partial_states_kernel[(num_actual,)](
-        kv,
-        kv.stride(0),
-        score,
-        score.stride(0),
-        ape,
-        ape.stride(0),
-        positions,
-        state_cache,
-        state_cache.stride(0),
-        state_cache.stride(1),
-        slot_mapping,
-        block_size,
-        HEAD_SIZE=head_size,
-        TRITON_BLOCK_SIZE=triton.next_power_of_2(head_size),
-        STATE_WIDTH=state_width,
-        COMPRESS_RATIO=compress_ratio,
-        **(pdl_kwargs or {}),
-    )
-
-
-@triton.jit
-def _save_partial_states_kernel(
-    kv_ptr,
-    kv_stride,
-    score_ptr,
-    score_stride,
-    ape_ptr,
-    ape_stride,
-    positions_ptr,
-    state_cache_ptr,
-    state_cache_stride0,
-    state_cache_stride1,
-    slot_mapping_ptr,
-    block_size,
-    HEAD_SIZE: tl.constexpr,
-    TRITON_BLOCK_SIZE: tl.constexpr,
-    # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
-    STATE_WIDTH: tl.constexpr,
-    COMPRESS_RATIO: tl.constexpr,
-):
-    token_idx = tl.program_id(0)
-    slot_id = tl.load(slot_mapping_ptr + token_idx)
-
-    # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
-    # by vLLM).  During CUDA graph replay the batch may contain padding
-    # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
-    # illegal memory access.
-    if slot_id < 0:
-        return
-
-    block_idx = slot_id // block_size
-    pos_in_block = slot_id % block_size
-    base_ptr = (
-        state_cache_ptr
-        + block_idx * state_cache_stride0
-        + pos_in_block * state_cache_stride1
-    )
-
-    block = tl.arange(0, TRITON_BLOCK_SIZE)
-    mask = block < HEAD_SIZE
-
-    kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
-    tl.store(base_ptr + block, kv, mask=mask)
-
-    # Fused: score += ape[position % compress_ratio]
-    position = tl.load(positions_ptr + token_idx)
-    ape_row = position % COMPRESS_RATIO
-    ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
-    score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
-    tl.store(
-        base_ptr + STATE_WIDTH + block,
-        score + ape,
-        mask=mask,
-    )

vllm/models/deepseek_v4/common/ops/sparse_attn_compress_cutedsl.py

@@ -8,7 +8,6 @@ The public wrappers provide the C4 fused and C128 split kernels.
 from __future__ import annotations
 
 from functools import cache
-from typing import Any
 
 import cutlass
 import cutlass.cute as cute
@@ -1087,7 +1086,7 @@ class SparseAttnNormRopeStoreKernel:
         )
 
 
-def compress_kv_sparse_attn_cutedsl(
+def _compress_kv_sparse_attn_cutedsl(
     state_cache: torch.Tensor,
     token_to_req_indices: torch.Tensor,
     positions: torch.Tensor,
@@ -1119,7 +1118,7 @@ def compress_kv_sparse_attn_cutedsl(
     )
 
 
-def norm_rope_insert_sparse_attn_cutedsl(
+def _norm_rope_insert_sparse_attn_cutedsl(
     compressed_kv: torch.Tensor,
     positions: torch.Tensor,
     slot_mapping: torch.Tensor,
@@ -1175,7 +1174,7 @@ def norm_rope_insert_sparse_attn_cutedsl(
     )
 
 
-def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
+def _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
     state_cache: torch.Tensor,
     token_to_req_indices: torch.Tensor,
     positions: torch.Tensor,
@@ -1239,94 +1238,3 @@ def fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
         kv_slot_mapping,
         kv_cache_block_size,
     )
-
-
-def compress_norm_rope_store_cutedsl(
-    state_cache: torch.Tensor,
-    num_actual: int,
-    token_to_req_indices: torch.Tensor,
-    positions: torch.Tensor,
-    slot_mapping: torch.Tensor,
-    block_table: torch.Tensor,
-    block_size: int,
-    state_width: int,
-    cos_sin_cache: torch.Tensor,
-    kv_cache: torch.Tensor,
-    k_cache_metadata: Any,
-    pdl_kwargs: dict,
-    head_dim: int,
-    rope_head_dim: int,
-    compress_ratio: int,
-    overlap: bool,
-    use_fp4_cache: bool,
-    rms_norm_weight: torch.Tensor,
-    rms_norm_eps: float,
-    quant_block: int,
-    token_stride: int,
-    scale_dim: int,
-) -> None:
-    if compress_ratio == 4:
-        # For C4A, the single fused kernel is faster than the two-kernel version.
-        fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(
-            state_cache,
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_size,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            kv_cache,
-            k_cache_metadata.slot_mapping,
-            kv_cache.shape[1],  # paged KV cache block size
-            kv_cache.stride(0),
-            head_size=head_dim,
-            state_width=state_width,
-            rope_head_dim=rope_head_dim,
-            fp8_max=448.0,
-            quant_block=quant_block,
-            token_stride=token_stride,
-            scale_dim=scale_dim,
-            compress_ratio=compress_ratio,
-            overlap=overlap,
-        )
-    else:
-        # For C128, the two-kernel version is faster than the single fused kernel.
-        compressed_kv = torch.empty(
-            (num_actual, head_dim),
-            dtype=torch.float32,
-            device=state_cache.device,
-        )
-        compress_kv_sparse_attn_cutedsl(
-            state_cache,
-            token_to_req_indices,
-            positions,
-            slot_mapping,
-            block_table,
-            block_size,
-            compressed_kv,
-            head_size=head_dim,
-            state_width=state_width,
-            compress_ratio=compress_ratio,
-            overlap=overlap,
-        )
-        norm_rope_insert_sparse_attn_cutedsl(
-            compressed_kv,
-            positions,
-            slot_mapping,
-            rms_norm_weight,
-            rms_norm_eps,
-            cos_sin_cache,
-            kv_cache,
-            k_cache_metadata.slot_mapping,
-            kv_cache.shape[1],  # paged KV cache block size
-            kv_cache.stride(0),
-            head_size=head_dim,
-            rope_head_dim=rope_head_dim,
-            fp8_max=448.0,
-            quant_block=quant_block,
-            token_stride=token_stride,
-            scale_dim=scale_dim,
-            compress_ratio=compress_ratio,
-        )

vllm/models/deepseek_v4/compressor.py

@@ -13,13 +13,15 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import MergedColumnParallelLinear
 from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
-    compress_norm_rope_store_triton,
+    _compress_kv_sparse_attn_cutedsl,
+    _fused_kv_compress_norm_rope_insert_indexer_attn,
+    _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
+    _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl,
+    _norm_rope_insert_sparse_attn_cutedsl,
 )
 from vllm.models.deepseek_v4.common.ops.fused_indexer_q import MXFP4_BLOCK_SIZE
-from vllm.models.deepseek_v4.common.ops.save_partial_states import (
-    save_partial_states,
-)
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -171,16 +173,6 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
 
 
 class DeepseekCompressor(nn.Module):
-    """DeepSeek V4 KV/score compressor.
-
-    Owns the linear / norm / state-cache / ape state and the shared forward
-    prologue (kv/score split, save_partial_states launch). The
-    compress → norm → RoPE → store step is dispatched to a triton kernel
-    (``compress_norm_rope_store_triton``) by default, except for the NVIDIA
-    head_dim=128 indexer path which uses the cutedsl kernel
-    (``compress_norm_rope_store_cutedsl``) for better performance.
-    """
-
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -250,18 +242,32 @@ class DeepseekCompressor(nn.Module):
             assert not use_fp4_cache, (
                 "MXFP4 cache is only supported for indexer (head=128)"
             )
+            self._use_cutedsl_sparse_compressor = True
+            self._use_cutedsl_fused_sparse_compressor = self.compress_ratio == 4
+            self._compress_kernel = _compress_kv_sparse_attn_cutedsl
+            self._norm_rope_store_kernel = _norm_rope_insert_sparse_attn_cutedsl
+            self._fused_sparse_kernel = (
+                _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl
+            )
             self._quant_block = 64
             self._token_stride = self.nope_head_dim + self.rope_head_dim * 2
             self._scale_dim = self.nope_head_dim // 64 + 1  # 7 real + 1 pad
+            self._num_warps = 4
         elif self.head_dim == 128:
+            self._use_cutedsl_sparse_compressor = False
             if use_fp4_cache:
+                self._fused_kernel = (
+                    _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
+                )
                 self._quant_block = MXFP4_BLOCK_SIZE
                 self._token_stride = self.head_dim // 2
                 self._scale_dim = self.head_dim // MXFP4_BLOCK_SIZE
             else:
+                self._fused_kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
                 self._quant_block = 128
                 self._token_stride = self.head_dim
                 self._scale_dim = 4  # single float32 scale
+            self._num_warps = 1
         else:
             raise ValueError(
                 f"Unsupported head_dim for fused quant+cache: {self.head_dim}"
@@ -306,22 +312,29 @@ class DeepseekCompressor(nn.Module):
         )
 
         # Store the KV and score (with fused APE addition) in the state.
-        # NOTE: PDL is disabled — both this kernel and the compress kernels
-        # below depend on preceding kernel outputs (kv/score from the cublas
-        # GEMM; state_cache from this kernel) but neither emits/waits on PDL
-        # grid dependency primitives, so launch_pdl=True caused a
-        # read-after-write race and non-deterministic output.
-        save_partial_states(
-            kv=kv,
-            score=score,
-            ape=self.ape,
-            positions=positions,
-            state_cache=state_cache,
-            slot_mapping=slot_mapping,
-            block_size=block_size,
-            state_width=state_width,
-            compress_ratio=self.compress_ratio,
-            pdl_kwargs=pdl_kwargs,
+        # NOTE: PDL is disabled — both this kernel and _fused_kernel below
+        # depend on preceding kernel outputs (kv/score from the cublas GEMM;
+        # state_cache from this kernel) but neither emits/waits on PDL grid
+        # dependency primitives, so launch_pdl=True caused a read-after-write
+        # race and non-deterministic output.
+        _save_partial_states_kernel[(num_actual,)](
+            kv,
+            kv.stride(0),
+            score,
+            score.stride(0),
+            self.ape,
+            self.ape.stride(0),
+            positions,
+            state_cache,
+            state_cache.stride(0),
+            state_cache.stride(1),
+            slot_mapping,
+            block_size,
+            HEAD_SIZE=kv.shape[-1],
+            TRITON_BLOCK_SIZE=triton.next_power_of_2(kv.shape[-1]),
+            STATE_WIDTH=state_width,
+            COMPRESS_RATIO=self.compress_ratio,
+            **pdl_kwargs,
         )
 
         # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
@@ -335,46 +348,161 @@ class DeepseekCompressor(nn.Module):
         k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
         kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache
 
-        if current_platform.is_cuda():
-            # NVIDIA GPUs.
-            if self.head_dim == 512:
-                from .nvidia.ops.sparse_attn_compress_cutedsl import (
-                    compress_norm_rope_store_cutedsl,
+        if self._use_cutedsl_sparse_compressor:
+            if self._use_cutedsl_fused_sparse_compressor:
+                self._fused_sparse_kernel(
+                    state_cache,
+                    token_to_req_indices,
+                    positions,
+                    slot_mapping,
+                    block_table,
+                    block_size,
+                    self.norm.weight,
+                    self.rms_norm_eps,
+                    cos_sin_cache,
+                    kv_cache,
+                    k_cache_metadata.slot_mapping,
+                    kv_cache.shape[1],  # paged KV cache block size
+                    kv_cache.stride(0),
+                    head_size=self.head_dim,
+                    state_width=state_width,
+                    rope_head_dim=self.rope_head_dim,
+                    fp8_max=448.0,
+                    quant_block=self._quant_block,
+                    token_stride=self._token_stride,
+                    scale_dim=self._scale_dim,
+                    compress_ratio=self.compress_ratio,
+                    overlap=self.overlap,
                 )
-
-                # Main compressor path.
-                # Use a cutedsl kernel for better performance.
-                compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
             else:
-                # Indexer path (head_dim == 128).
-                # Use a triton kernel.
-                compress_norm_rope_store_fn = compress_norm_rope_store_triton
+                compressed_kv = torch.empty(
+                    (num_actual, self.head_dim),
+                    dtype=torch.float32,
+                    device=state_cache.device,
+                )
+                self._compress_kernel(
+                    state_cache,
+                    token_to_req_indices,
+                    positions,
+                    slot_mapping,
+                    block_table,
+                    block_size,
+                    compressed_kv,
+                    head_size=self.head_dim,
+                    state_width=state_width,
+                    compress_ratio=self.compress_ratio,
+                    overlap=self.overlap,
+                )
+                self._norm_rope_store_kernel(
+                    compressed_kv,
+                    positions,
+                    slot_mapping,
+                    self.norm.weight,
+                    self.rms_norm_eps,
+                    cos_sin_cache,
+                    kv_cache,
+                    k_cache_metadata.slot_mapping,
+                    kv_cache.shape[1],  # paged KV cache block size
+                    kv_cache.stride(0),
+                    head_size=self.head_dim,
+                    rope_head_dim=self.rope_head_dim,
+                    fp8_max=448.0,
+                    quant_block=self._quant_block,
+                    token_stride=self._token_stride,
+                    scale_dim=self._scale_dim,
+                    compress_ratio=self.compress_ratio,
+                )
         else:
-            # AMD GPUs.
-            # Always use a triton kernel.
-            compress_norm_rope_store_fn = compress_norm_rope_store_triton
+            self._fused_kernel[(num_actual,)](
+                # state cache
+                state_cache,
+                state_cache.stride(0),
+                state_cache.stride(1),
+                # metadata
+                token_to_req_indices,
+                positions,
+                slot_mapping,
+                block_table,
+                block_table.stride(0),
+                block_size,
+                # RMSNorm
+                self.norm.weight,
+                self.rms_norm_eps,
+                # RoPE
+                cos_sin_cache,
+                cos_sin_cache.stride(0),
+                # KV cache
+                kv_cache,
+                k_cache_metadata.slot_mapping,
+                kv_cache.shape[1],  # paged KV cache block size (tokens per block)
+                # constexprs
+                HEAD_SIZE=self.head_dim,
+                TRITON_BLOCK_SIZE=triton.next_power_of_2(self.head_dim),
+                STATE_WIDTH=state_width,
+                COMPRESS_RATIO=self.compress_ratio,
+                OVERLAP=self.overlap,
+                ROPE_HEAD_DIM=self.rope_head_dim,
+                FP8_MAX=448.0,
+                QUANT_BLOCK=self._quant_block,
+                TOKEN_STRIDE=self._token_stride,
+                SCALE_DIM=self._scale_dim,
+                KV_BLOCK_STRIDE=kv_cache.stride(0),
+                num_warps=self._num_warps,
+                **pdl_kwargs,
+            )
 
-        compress_norm_rope_store_fn(
-            state_cache=state_cache,
-            num_actual=num_actual,
-            token_to_req_indices=token_to_req_indices,
-            positions=positions,
-            slot_mapping=slot_mapping,
-            block_table=block_table,
-            block_size=block_size,
-            state_width=state_width,
-            cos_sin_cache=cos_sin_cache,
-            kv_cache=kv_cache,
-            k_cache_metadata=k_cache_metadata,
-            pdl_kwargs=pdl_kwargs,
-            head_dim=self.head_dim,
-            rope_head_dim=self.rope_head_dim,
-            compress_ratio=self.compress_ratio,
-            overlap=self.overlap,
-            use_fp4_cache=self.use_fp4_cache,
-            rms_norm_weight=self.norm.weight,
-            rms_norm_eps=self.rms_norm_eps,
-            quant_block=self._quant_block,
-            token_stride=self._token_stride,
-            scale_dim=self._scale_dim,
-        )
+
+@triton.jit
+def _save_partial_states_kernel(
+    kv_ptr,
+    kv_stride,
+    score_ptr,
+    score_stride,
+    ape_ptr,
+    ape_stride,
+    positions_ptr,
+    state_cache_ptr,
+    state_cache_stride0,
+    state_cache_stride1,
+    slot_mapping_ptr,
+    block_size,
+    HEAD_SIZE: tl.constexpr,
+    TRITON_BLOCK_SIZE: tl.constexpr,
+    # state_cache last dim packs [kv_state, score_state], each STATE_WIDTH wide.
+    STATE_WIDTH: tl.constexpr,
+    COMPRESS_RATIO: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    slot_id = tl.load(slot_mapping_ptr + token_idx)
+
+    # Skip padded / invalid tokens (slot_id == -1 is the PAD sentinel used
+    # by vLLM).  During CUDA graph replay the batch may contain padding
+    # tokens whose slot_mapping is -1; writing to kv_state[-1] would be an
+    # illegal memory access.
+    if slot_id < 0:
+        return
+
+    block_idx = slot_id // block_size
+    pos_in_block = slot_id % block_size
+    base_ptr = (
+        state_cache_ptr
+        + block_idx * state_cache_stride0
+        + pos_in_block * state_cache_stride1
+    )
+
+    block = tl.arange(0, TRITON_BLOCK_SIZE)
+    mask = block < HEAD_SIZE
+
+    kv = tl.load(kv_ptr + token_idx * kv_stride + block, mask=mask)
+    tl.store(base_ptr + block, kv, mask=mask)
+
+    # Fused: score += ape[position % compress_ratio]
+    position = tl.load(positions_ptr + token_idx)
+    ape_row = position % COMPRESS_RATIO
+    ape = tl.load(ape_ptr + ape_row * ape_stride + block, mask=mask)
+    score = tl.load(score_ptr + token_idx * score_stride + block, mask=mask)
+    tl.store(
+        base_ptr + STATE_WIDTH + block,
+        score + ape,
+        mask=mask,
+    )

vllm/models/deepseek_v4/nvidia/ops/__init__.py

@@ -5,9 +5,4 @@
 These modules import ``cutlass``/``cutedsl`` at module top level, so they must
 not be imported on non-CUDA platforms. Callers should gate on
 ``vllm.utils.import_utils.has_cutedsl()`` before importing from here.
-
-This ``__init__`` deliberately imports nothing: re-exporting the cutedsl
-modules here would eagerly ``import cutlass`` (initializing the CUDA driver) for
-anyone who imports ``vllm.models.deepseek_v4``, breaking forked subprocesses.
-Import the leaf modules directly under a ``has_cutedsl()``/``is_cuda()`` gate.
 """

vllm/platforms/cpu.py

@@ -14,6 +14,7 @@ from vllm.logger import init_logger
 from vllm.utils.cpu_resource_utils import (
     DEVICE_CONTROL_ENV_VAR,
     get_memory_node_info,
+    get_visible_memory_node,
 )
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -135,9 +136,13 @@ class CpuPlatform(Platform):
         scheduler_config.async_scheduling = False
 
         parallel_config = vllm_config.parallel_config
-        # 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":
+        if (
+            os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING", "1") == "1"
+            and parallel_config.distributed_executor_backend == "uni"
+        ):
+            # OMP requires the MP executor to function correctly, UniProc
+            # is not supported as it is not possible to set the OMP
+            # environment correctly
             parallel_config.distributed_executor_backend = "mp"
 
         if parallel_config.worker_cls == "auto":
@@ -481,3 +486,15 @@ class CpuPlatform(Platform):
             slot_mapping,
             isa,
         )
+
+    @classmethod
+    def get_current_memory_usage(
+        cls, device: torch.types.Device | None = None
+    ) -> float:
+        allowed_mem_node_list = get_visible_memory_node()
+        mem_status_list = [get_memory_node_info(i) for i in allowed_mem_node_list]
+        memory_usage = 0
+        for s in mem_status_list:
+            memory_usage += s.total_memory - s.available_memory
+
+        return memory_usage

vllm/platforms/cuda.py

@@ -592,15 +592,6 @@ class CudaPlatformBase(Platform):
             default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm
         )
 
-    @classmethod
-    def is_arch_support_pdl(cls) -> bool:
-        try:
-            device = torch.cuda.current_device()
-            major, _ = torch.cuda.get_device_capability(device)
-        except Exception:
-            return False
-        return major >= 9
-
 
 # NVML utils
 # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,

vllm/platforms/interface.py

@@ -1016,13 +1016,6 @@ class Platform:
         # Native always used by default. Platforms can override this behavior.
         return IrOpPriorityConfig.with_default(["native"])
 
-    @classmethod
-    def is_arch_support_pdl(cls) -> bool:
-        """
-        Does the current platform support PDL (Programmatic Dependent Launch)?
-        """
-        return False
-
 
 class UnspecifiedPlatform(Platform):
     _enum = PlatformEnum.UNSPECIFIED

vllm/sampling_params.py

@@ -7,9 +7,10 @@ import json as json_mod
 from dataclasses import field
 from enum import Enum, IntEnum
 from functools import cached_property
-from typing import Any
+from typing import Annotated, Any
 
 import msgspec
+from pydantic import BeforeValidator
 from pydantic.dataclasses import dataclass
 
 import vllm.envs as envs
@@ -30,6 +31,35 @@ MAX_LOGPROB_TOKEN_IDS = 128
 the per-request row width allocated by the sampler's `LogprobTokenIdsState`."""
 
 
+def validate_thinking_token_budget(value: int | float | bool | None) -> int | None:
+    """Validate ``thinking_token_budget``; return ``None`` if unset."""
+    if value is None:
+        return None
+    if isinstance(value, (bool, float)) or not isinstance(value, int):
+        raise VLLMValidationError(
+            "`thinking_token_budget` must be a non-negative integer "
+            "or -1 for unlimited.",
+            parameter="thinking_token_budget",
+            value=value,
+        )
+    if value == -1:
+        return None
+    if value < 0:
+        raise VLLMValidationError(
+            "`thinking_token_budget` must be a non-negative integer "
+            "or -1 for unlimited.",
+            parameter="thinking_token_budget",
+            value=value,
+        )
+    return value
+
+
+ThinkingTokenBudget = Annotated[
+    int | None,
+    BeforeValidator(validate_thinking_token_budget),
+]
+
+
 class SamplingType(IntEnum):
     GREEDY = 0
     RANDOM = 1
@@ -409,6 +439,10 @@ class SamplingParams(
         if self.seed == -1:
             self.seed = None
 
+        self.thinking_token_budget = validate_thinking_token_budget(
+            self.thinking_token_budget
+        )
+
         if self.stop is None:
             self.stop = []
         elif isinstance(self.stop, str):

vllm/triton_utils/importing.py

@@ -3,6 +3,7 @@
 
 import os
 import types
+from importlib.metadata import version
 from importlib.util import find_spec
 
 from vllm.logger import init_logger
@@ -48,6 +49,17 @@ if HAS_TRITON:
                 len(active_drivers),
             )
             HAS_TRITON = False
+
+        # Check Triton CPU
+        if "cpu" in version("vllm"):
+            if "cpu" in backends:
+                HAS_TRITON = True
+            else:
+                logger.warning(
+                    "Triton is installed, but doesn't include CPU backend. "
+                    "Disabling Triton."
+                )
+                HAS_TRITON = False
     except ImportError:
         # This can occur if Triton is partially installed or triton.backends
         # is missing.

vllm/utils/flashinfer.py

@@ -8,6 +8,7 @@ Users of vLLM should always import **only** these wrappers.
 import contextlib
 import functools
 import importlib
+import importlib.metadata
 import importlib.util
 import os
 import shutil
@@ -16,6 +17,7 @@ from typing import Any, NoReturn
 
 import requests
 import torch
+from packaging.version import Version
 
 import vllm.envs as envs
 from vllm.logger import init_logger
@@ -63,6 +65,43 @@ def has_flashinfer() -> bool:
     return True
 
 
+@functools.cache
+def get_flashinfer_version() -> Version | None:
+    if not has_flashinfer():
+        return None
+
+    return Version(importlib.metadata.version("flashinfer-python"))
+
+
+@functools.cache
+def flashinfer_trtllm_fp4_8x4_is_safe() -> bool:
+    """Return whether FlashInfer TRTLLM FP4 8x4 is safe to use.
+
+    FlashInfer documents TRTLLM FP4 as allowing 8x4 only for the activation
+    (A) scale layout while keeping the weight (B) path in 128x4. However,
+    released FlashInfer versions up to and including 0.6.11.* are known to
+    have correctness bugs on this path; see flashinfer-ai/flashinfer#2861.
+    """
+    version = get_flashinfer_version()
+    if version is None:
+        logger.warning_once(
+            "Disabling FlashInfer TRTLLM FP4 8x4 path because the installed "
+            "FlashInfer version could not be determined."
+        )
+        return False
+
+    if Version(version.base_version) <= Version("0.6.11"):
+        logger.warning_once(
+            "Disabling FlashInfer TRTLLM FP4 8x4 path for flashinfer==%s due "
+            "to known upstream correctness bugs in released versions up to "
+            "0.6.11.* (see flashinfer-ai/flashinfer#2861).",
+            version,
+        )
+        return False
+
+    return True
+
+
 def _missing(*_: Any, **__: Any) -> NoReturn:
     """Placeholder for unavailable FlashInfer backend."""
     raise RuntimeError(
@@ -742,7 +781,9 @@ def flashinfer_scaled_fp4_mm(
         block_scale_a = block_scale_a.view(torch.uint8)
         block_scale_b = block_scale_b.view(torch.uint8)
 
-    use_8x4_sf_layout = True if backend == "trtllm" and a.shape[0] <= 32 else False  # noqa: SIM210
+    use_8x4_sf_layout = (
+        backend == "trtllm" and a.shape[0] <= 32 and flashinfer_trtllm_fp4_8x4_is_safe()
+    )
 
     return flashinfer_mm_fp4(
         a,
@@ -954,6 +995,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
 
 __all__ = [
     "has_flashinfer",
+    "get_flashinfer_version",
     "flashinfer_trtllm_fp8_block_scale_moe",
     "flashinfer_cutlass_fused_moe",
     "flashinfer_cutedsl_grouped_gemm_nt_masked",
@@ -982,6 +1024,7 @@ __all__ = [
     "use_trtllm_attention",
     "flashinfer_mxfp4_quantize",
     "flashinfer_scaled_fp4_mm",
+    "flashinfer_trtllm_fp4_8x4_is_safe",
     "flashinfer_scaled_fp4_mm_out",
     "flashinfer_scaled_fp8_mm",
     "flashinfer_scaled_fp8_mm_out",

vllm/utils/import_utils.py

@@ -430,7 +430,6 @@ def has_triton_kernels() -> bool:
     return is_available
 
 
-@cache
 def has_tilelang() -> bool:
     """Whether the optional `tilelang` package is available."""
     return _has_module("tilelang")

vllm/utils/platform_utils.py

@@ -50,8 +50,10 @@ def is_pin_memory_available() -> bool:
 def is_uva_available() -> bool:
     """Check if Unified Virtual Addressing (UVA) is available."""
     # UVA requires pinned memory.
+    from vllm.platforms import current_platform
+
     # TODO: Add more requirements for UVA if needed.
-    return is_pin_memory_available()
+    return is_pin_memory_available() or current_platform.is_cpu()
 
 
 @cache

vllm/v1/engine/async_llm.py

@@ -81,7 +81,7 @@ class AsyncLLM(EngineClient):
         start_engine_loop: bool = True,
         stat_loggers: list[StatLoggerFactory] | None = None,
         aggregate_engine_logging: bool = False,
-        client_addresses: dict[str, Any] | None = None,
+        client_addresses: dict[str, str] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ) -> None:
@@ -209,7 +209,7 @@ class AsyncLLM(EngineClient):
         enable_log_requests: bool = False,
         aggregate_engine_logging: bool = False,
         disable_log_stats: bool = False,
-        client_addresses: dict[str, Any] | None = None,
+        client_addresses: dict[str, str] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ) -> "AsyncLLM":

vllm/v1/engine/coordinator.py

@@ -11,7 +11,7 @@ import zmq
 
 from vllm.config import ParallelConfig
 from vllm.logger import init_logger
-from vllm.utils.network_utils import make_zmq_socket
+from vllm.utils.network_utils import get_tcp_uri, make_zmq_socket
 from vllm.utils.system_utils import get_mp_context, set_process_title
 from vllm.v1.engine import EngineCoreOutputs, EngineCoreRequestType
 from vllm.v1.serial_utils import MsgpackDecoder
@@ -91,9 +91,16 @@ class DPCoordinator:
         if parallel_config.enable_elastic_ep:
             local_only_eng = False
 
-        front_publish_address = get_engine_client_zmq_addr(local_only, host=host)
-        back_publish_address = get_engine_client_zmq_addr(local_only_eng, host=host)
-        back_output_address = get_engine_client_zmq_addr(local_only_eng, host=host)
+        def bind_address(local_only: bool) -> str:
+            return (
+                get_engine_client_zmq_addr(local_only=True, host=host)
+                if local_only
+                else get_tcp_uri(host, 0)
+            )
+
+        front_publish_address = bind_address(local_only)
+        back_publish_address = bind_address(local_only_eng)
+        back_output_address = bind_address(local_only_eng)
 
         context = get_mp_context()
         parent_zmq_addr_pipe, child_zmq_addr_pipe = context.Pipe(duplex=False)

vllm/v1/engine/core_client.py

@@ -11,7 +11,6 @@ from collections import defaultdict, deque
 from collections.abc import Awaitable, Callable, Sequence
 from concurrent.futures import Future
 from dataclasses import dataclass
-from multiprocessing.connection import Connection
 from multiprocessing.queues import Queue
 from threading import Thread
 from typing import Any, TypeAlias, TypeVar
@@ -109,7 +108,7 @@ class EngineCoreClient(ABC):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, Any] | None = None,
+        client_addresses: dict[str, str] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ) -> "AsyncMPClient":
@@ -477,7 +476,7 @@ class MPClient(EngineCoreClient):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, Any] | None = None,
+        client_addresses: dict[str, str] | None = None,
     ):
         self.vllm_config = vllm_config
 
@@ -508,7 +507,7 @@ class MPClient(EngineCoreClient):
                 output_address = client_addresses["output_address"]
                 self.stats_update_address = client_addresses.get("stats_update_address")
                 # Tensor queues passed via client_addresses for multi-API-server case
-                tensor_queue = client_addresses.get("tensor_queue")
+                tensor_queue = client_addresses.get("tensor_queue")  # type: ignore[assignment]
                 self.input_socket = self.resources.input_socket = make_zmq_socket(
                     self.ctx,
                     input_address,
@@ -519,28 +518,6 @@ class MPClient(EngineCoreClient):
                 self.resources.output_socket = make_zmq_socket(
                     self.ctx, output_address, zmq.PULL
                 )
-
-                # Report bound endpoints back so the parent can forward
-                # them to engines (mirrors the DPCoordinator pattern).
-                actual_address_pipe: Connection | None = client_addresses.get(
-                    "actual_address_pipe"
-                )
-                if actual_address_pipe is not None:
-                    try:
-                        actual_input = self.input_socket.getsockopt(
-                            zmq.LAST_ENDPOINT
-                        ).decode()
-                        actual_output = self.resources.output_socket.getsockopt(
-                            zmq.LAST_ENDPOINT
-                        ).decode()
-                        actual_address_pipe.send(
-                            {
-                                "input_address": actual_input,
-                                "output_address": actual_output,
-                            }
-                        )
-                    finally:
-                        actual_address_pipe.close()
             else:
                 # Engines are managed by this client.
                 addresses = get_engine_zmq_addresses(vllm_config)
@@ -555,15 +532,6 @@ class MPClient(EngineCoreClient):
                     self.ctx, addresses.outputs[0], zmq.PULL
                 )
 
-                # Resolve ``tcp://host:0`` placeholders to bound endpoints
-                # before engines DEALER-connect. No-op for IPC.
-                addresses.inputs[0] = self.input_socket.getsockopt(
-                    zmq.LAST_ENDPOINT
-                ).decode()
-                addresses.outputs[0] = self.resources.output_socket.getsockopt(
-                    zmq.LAST_ENDPOINT
-                ).decode()
-
                 with launch_core_engines(
                     vllm_config, executor_class, log_stats, addresses
                 ) as (engine_manager, coordinator, addresses, tensor_queue):
@@ -925,7 +893,7 @@ class AsyncMPClient(MPClient):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, Any] | None = None,
+        client_addresses: dict[str, str] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ):
@@ -1175,7 +1143,7 @@ class DPAsyncMPClient(AsyncMPClient):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, Any] | None = None,
+        client_addresses: dict[str, str] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ):
@@ -1355,7 +1323,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient):
         vllm_config: VllmConfig,
         executor_class: type[Executor],
         log_stats: bool,
-        client_addresses: dict[str, Any] | None = None,
+        client_addresses: dict[str, str] | None = None,
         client_count: int = 1,
         client_index: int = 0,
     ):

vllm/v1/engine/utils.py

@@ -23,12 +23,7 @@ from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.ray.ray_env import get_env_vars_to_copy
 from vllm.utils import numa_utils
-from vllm.utils.network_utils import (
-    get_open_port,
-    get_open_zmq_ipc_path,
-    get_tcp_uri,
-    zmq_socket_ctx,
-)
+from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx
 from vllm.utils.system_utils import get_mp_context
 from vllm.v1.engine.coordinator import DPCoordinator
 from vllm.v1.executor import Executor
@@ -960,19 +955,8 @@ class CoreEngineActorManager:
 def get_engine_zmq_addresses(
     vllm_config: VllmConfig,
     num_api_servers: int = 1,
-    *,
-    defer_api_server_ports: bool = True,
 ) -> EngineZmqAddresses:
-    """Allocate ZMQ addresses for engine-client communication.
-
-    By default each TCP address is a ``tcp://host:0`` placeholder; the
-    consumer (API-server child or single-process ``MPClient``) binds, then
-    recovers the kernel-assigned port via ``getsockopt(zmq.LAST_ENDPOINT)``
-    and writes it back into ``addresses`` before the engine handshake.
-
-    Set ``defer_api_server_ports=False`` only when the consumer cannot
-    report a bound port back (e.g. the Rust front-end). IPC paths are
-    unaffected."""
+    """Allocate ZMQ addresses for engine-client communication."""
     parallel_config = vllm_config.parallel_config
     local_engine_count = parallel_config.data_parallel_size_local
     local_start_index = parallel_config.data_parallel_rank_local
@@ -994,14 +978,15 @@ def get_engine_zmq_addresses(
     if parallel_config.enable_elastic_ep:
         client_local_only = False
 
-    def _addr() -> str:
-        if client_local_only:
-            return get_open_zmq_ipc_path()
-        return get_tcp_uri(host, 0 if defer_api_server_ports else get_open_port())
-
     return EngineZmqAddresses(
-        inputs=[_addr() for _ in range(num_api_servers)],
-        outputs=[_addr() for _ in range(num_api_servers)],
+        inputs=[
+            get_engine_client_zmq_addr(client_local_only, host)
+            for _ in range(num_api_servers)
+        ],
+        outputs=[
+            get_engine_client_zmq_addr(client_local_only, host)
+            for _ in range(num_api_servers)
+        ],
     )
 
 
@@ -1110,11 +1095,9 @@ def launch_core_engines(
     if parallel_config.enable_elastic_ep:
         handshake_local_only = False
 
-    # Preserve "port=0 means auto-pick" for the handshake address, which
-    # is consumed by engines spawned in this process and so cannot defer
-    # port resolution to bind time.
-    rpc_port = parallel_config.data_parallel_rpc_port or get_open_port()
-    handshake_address = get_engine_client_zmq_addr(handshake_local_only, host, rpc_port)
+    handshake_address = get_engine_client_zmq_addr(
+        handshake_local_only, host, parallel_config.data_parallel_rpc_port
+    )
 
     if local_engines_only and dp_rank > 0:
         assert not handshake_local_only