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

[DSv4] Drop _get_compressed_kv_buffer in DeepseekCompressor (#43690 )
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
2026-05-26 14:48:26 -04:00 · 2026-05-26 10:11:25 -07:00 · 2026-05-26 09:31:23 -07:00 · 2026-05-26 09:21:05 -07:00 · 2026-05-26 08:56:04 -07:00 · 2026-05-26 11:15:42 -04:00
70 changed files with 3825 additions and 1302 deletions
@@ -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  --------------------------------------------------------#

@@ -38,6 +38,28 @@ steps:
  commands:
    - pytest -v -s kernels/core/test_minimax_reduce_rms.py

+- label: Deepseek V4 Kernel Test (H100)
+  key: deepseek-v4-kernel-test-h100
+  timeout_in_minutes: 15
+  device: h100
+  source_file_dependencies:
+  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+  - vllm/models/deepseek_v4/common/ops/
+  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  commands:
+    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+
+- label: Deepseek V4 Kernel Test (B200)
+  key: deepseek-v4-kernel-test-b200
+  timeout_in_minutes: 15
+  device: b200-k8s
+  source_file_dependencies:
+  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
+  - vllm/models/deepseek_v4/common/ops/
+  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  commands:
+    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+
 - label: Kernels Attention Test %N
  key: kernels-attention-test
  timeout_in_minutes: 35
@@ -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:
@@ -122,18 +122,45 @@ __device__ __forceinline__ float warpSum(float val) {
 // Per-slot inner pipeline
 // ────────────────────────────────────────────────────────────────────────────
 // Shared by both kernel variants: 1 CTA per (token, head) pair vs. 1 CTA per
-// token
-template <typename scalar_t_in>
+// token.  Templated on `kNumHeadsQPadded` so the KV-sentinel comparison and
+// q_out stride fold to compile-time constants.
+//
+// Slot layout (per token):
+//   slot < num_heads_q                          → live-Q   (RMSNorm + RoPE,
+//                                                           read q_in →
+//                                                           write q_out)
+//   num_heads_q <= slot < kNumHeadsQPadded      → pad-Q    (zero-fill q_out;
+//                                                           v0/v1 unused)
+//   slot == kNumHeadsQPadded                    → KV       (RoPE + UE8M0 quant
+//                                                           + paged-cache
+//                                                           insert)
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __device__ __forceinline__ void processDeepseekV4Slot(
    uint4 v0, uint4 v1, int const tokenIdx, int const slotIdx,
    int const dim_base, int const laneId, int const num_heads_q,
-    float const eps, scalar_t_in* __restrict__ q_inout,
+    float const eps, scalar_t_in* __restrict__ q_out,
    uint8_t* __restrict__ k_cache, int64_t const* __restrict__ slot_mapping,
    int64_t const* __restrict__ position_ids,
    float const* __restrict__ cos_sin_cache, int const cache_block_size,
    int const kv_block_stride) {
  using Converter = vllm::_typeConvert<scalar_t_in>;
-  bool const isKV = (slotIdx == num_heads_q);
+  bool const isKV = (slotIdx == kNumHeadsQPadded);
+  bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
+
+  // ── Pad-Q branch: write 32 B of zeros and exit. ─────────────────────────
+  // FlashMLA reads these slots; bf16 +0.0 is bit pattern 0x0000, so a uint4
+  // zero literal is correct.  Matches the live-Q branch's vectorized store.
+  if (isPadQ) {
+    scalar_t_in* dst =
+        q_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
+        dim_base;
+    uint4 const zero4 = {0u, 0u, 0u, 0u};
+    *reinterpret_cast<uint4*>(dst) = zero4;
+    *reinterpret_cast<uint4*>(dst + 8) = zero4;
+    return;
+  }

  // ── Decode the bf16 → 16 fp32 registers ─────────────────────────────
  float elements[kElemsPerLane];
@@ -207,7 +234,7 @@ __device__ __forceinline__ void processDeepseekV4Slot(
  // triggering and per-iteration buffer rotation.
  // ═══════════════════════════════════════════════════════════════════
  if (!isKV) {
-    // ── Q: cast back to bf16 and store. ────────────────────────────
+    // ── Live-Q: cast back to bf16 and store into the padded q_out. ─────
    uint4 out0, out1;
    typename Converter::packed_hip_type* po0 =
        reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
@@ -224,8 +251,9 @@ __device__ __forceinline__ void processDeepseekV4Slot(
          make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
    }
    scalar_t_in* dst =
-        q_inout +
-        (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
+        q_out +
+        (static_cast<int64_t>(tokenIdx) * kNumHeadsQPadded + slotIdx) *
+            kHeadDim +
        dim_base;
    *reinterpret_cast<uint4*>(dst) = out0;
    *reinterpret_cast<uint4*>(dst + 8) = out1;
@@ -313,20 +341,32 @@ __device__ __forceinline__ void processDeepseekV4Slot(
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
 //
-// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
+// Grid: 1D, gridDim.x = ceil(num_tokens_full * (kNumHeadsQPadded + 1) /
 // warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
-// warp handles one (token, head_slot) pair. head_slot < num_heads_q          →
-// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q         → KV
-// branch (RoPE + UE8M0 quant + insert)
+// warp handles one (token, head_slot) pair.
+//   slot < num_heads_q                              → live-Q branch
+//                                                     (RMSNorm + RoPE,
+//                                                      read q_in → write q_out)
+//   num_heads_q <= slot < kNumHeadsQPadded          → pad-Q branch
+//                                                     (zero-fill q_out)
+//   slot == kNumHeadsQPadded                        → KV branch
+//                                                     (RoPE + UE8M0 quant +
+//                                                      paged-cache insert)
+//
+// `kNumHeadsQPadded` is a template parameter (compile-time constant) so the
+// divisions in the grid math and the KV-sentinel comparison fold to fast
+// constant operations.  The launch wrapper dispatches the runtime value to
+// the matching instantiation.
 //
 // With DP padding, q/kv/position_ids can have more rows than slot_mapping.
-// The Q branch covers all `num_tokens_full` rows (downstream attention uses
-// them).  The KV branch only inserts the first `num_tokens_insert` tokens
-// (= slot_mapping length) into the paged cache.
+// The live-Q and pad-Q branches cover all `num_tokens_full` rows (downstream
+// attention uses them).  The KV branch only inserts the first
+// `num_tokens_insert` tokens (= slot_mapping length) into the paged cache.
 //
-template <typename scalar_t_in>
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
-    scalar_t_in* __restrict__ q_inout,         // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ q_in,      // [N, num_heads_q,      512]
+    scalar_t_in* __restrict__ q_out,           // [N, kNumHeadsQPadded, 512]
    scalar_t_in const* __restrict__ kv_in,     // [N, 512] bf16
    uint8_t* __restrict__ k_cache,             // [num_blocks, block_stride]
    int64_t const* __restrict__ slot_mapping,  // [num_tokens_insert] i64
@@ -335,7 +375,7 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
    float const eps,
    int const num_tokens_full,    // = q.size(0) = kv.size(0)
    int const num_tokens_insert,  // = slot_mapping.size(0), ≤ num_tokens_full
-    int const num_heads_q,        // H
+    int const num_heads_q,        // live Q heads (input layout)
    int const cache_block_size,   // tokens per paged-cache block
    int const kv_block_stride) {  // bytes per paged-cache block
 #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
@@ -351,12 +391,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
    int const laneId = threadIdx.x % 32;
    int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;

-    int const total_slots_per_token = num_heads_q + 1;
-    int const tokenIdx = globalWarpIdx / total_slots_per_token;
-    int const slotIdx = globalWarpIdx % total_slots_per_token;
+    constexpr int kTotalSlotsPerToken = kNumHeadsQPadded + 1;
+    int const tokenIdx = globalWarpIdx / kTotalSlotsPerToken;
+    int const slotIdx = globalWarpIdx % kTotalSlotsPerToken;
    if (tokenIdx >= num_tokens_full) return;

-    bool const isKV = (slotIdx == num_heads_q);
+    bool const isKV = (slotIdx == kNumHeadsQPadded);
+    bool const isPadQ = !isKV && (slotIdx >= num_heads_q);
    // KV branch: skip DP-padded tokens (no slot reserved for them).
    if (isKV && tokenIdx >= num_tokens_insert) return;

@@ -371,22 +412,26 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
    // Dim range this lane owns within the 512-wide head.
    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16

-    // Two 16-byte loads per thread (8 bf16 each).  Use uint4 as the vector
-    // type; the shared per-slot helper bitcasts to scalar_t_in packed pairs.
-    scalar_t_in const* src_ptr;
-    if (isKV) {
-      src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
-    } else {
-      int64_t const q_row_offset =
-          (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
-          dim_base;
-      src_ptr = q_inout + q_row_offset;
+    // Load only for live-Q and KV slots; pad-Q skips the read (q_in beyond
+    // num_heads_q is out of bounds) and the helper zero-fills its output.
+    uint4 v0, v1;
+    if (!isPadQ) {
+      scalar_t_in const* src_ptr;
+      if (isKV) {
+        src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+      } else {
+        int64_t const q_row_offset =
+            (static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) *
+                kHeadDim +
+            dim_base;
+        src_ptr = q_in + q_row_offset;
+      }
+      v0 = *reinterpret_cast<uint4 const*>(src_ptr);
+      v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
    }
-    uint4 const v0 = *reinterpret_cast<uint4 const*>(src_ptr);
-    uint4 const v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);

-    processDeepseekV4Slot<scalar_t_in>(
-        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_inout,
+    processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
+        v0, v1, tokenIdx, slotIdx, dim_base, laneId, num_heads_q, eps, q_out,
        k_cache, slot_mapping, position_ids, cos_sin_cache, cache_block_size,
        kv_block_stride);

@@ -408,9 +453,9 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
 // Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q
 // KV branch (RoPE + UE8M0 quant + insert)
 //
-template <typename scalar_t_in>
+template <typename scalar_t_in, int kNumHeadsQPadded>
 __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
-    scalar_t_in* __restrict__ q_inout,  // [N, H, 512] bf16, in place
+    scalar_t_in const* __restrict__ q_in, scalar_t_in* __restrict__ q_out,
    scalar_t_in const* __restrict__ kv_in, uint8_t* __restrict__ k_cache,
    int64_t const* __restrict__ slot_mapping,
    int64_t const* __restrict__ position_ids,
@@ -435,25 +480,32 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 #endif

    int const dim_base = laneId * kElemsPerLane;  // in [0, 512) step 16
-    int const slot_end =
-        (tokenIdx >= num_tokens_insert) ? num_heads_q : (num_heads_q + 1);
+    // Slot enumeration: live-Q + pad-Q + (KV if this token has a slot).
+    int const slot_end = (tokenIdx >= num_tokens_insert)
+                             ? kNumHeadsQPadded
+                             : (kNumHeadsQPadded + 1);

-    auto src_for_slot = [&](int s) -> scalar_t_in const* {
-      if (s == num_heads_q) {
-        return kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+    auto load_slot = [&](int s, uint4& va, uint4& vb) {
+      // pad-Q slots skip the load — q_in beyond num_heads_q is OOB.
+      if (s >= num_heads_q && s < kNumHeadsQPadded) return;
+      scalar_t_in const* src;
+      if (s == kNumHeadsQPadded) {
+        src = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
+      } else {
+        src = q_in +
+              (static_cast<int64_t>(tokenIdx) * num_heads_q +
+               static_cast<int64_t>(s)) *
+                  kHeadDim +
+              dim_base;
      }
-      return q_inout +
-             (static_cast<int64_t>(tokenIdx) * num_heads_q +
-              static_cast<int64_t>(s)) *
-                 kHeadDim +
-             dim_base;
+      va = *reinterpret_cast<uint4 const*>(src);
+      vb = *reinterpret_cast<uint4 const*>(src + 8);
    };

    if (warpId < slot_end) {
      int curr_slot = warpId;
-      scalar_t_in const* src_curr = src_for_slot(curr_slot);
-      uint4 v0_curr = *reinterpret_cast<uint4 const*>(src_curr);
-      uint4 v1_curr = *reinterpret_cast<uint4 const*>(src_curr + 8);
+      uint4 v0_curr, v1_curr;
+      load_slot(curr_slot, v0_curr, v1_curr);

      while (curr_slot < slot_end) {
        int const next_slot = curr_slot + warpsPerBlock;
@@ -462,14 +514,12 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
        // Prefetch src for the next slot
        uint4 v0_next, v1_next;
        if (has_next) {
-          scalar_t_in const* src_next = src_for_slot(next_slot);
-          v0_next = *reinterpret_cast<uint4 const*>(src_next);
-          v1_next = *reinterpret_cast<uint4 const*>(src_next + 8);
+          load_slot(next_slot, v0_next, v1_next);
        }

-        processDeepseekV4Slot<scalar_t_in>(
+        processDeepseekV4Slot<scalar_t_in, kNumHeadsQPadded>(
            v0_curr, v1_curr, tokenIdx, curr_slot, dim_base, laneId,
-            num_heads_q, eps, q_inout, k_cache, slot_mapping, position_ids,
+            num_heads_q, eps, q_out, k_cache, slot_mapping, position_ids,
            cos_sin_cache, cache_block_size, kv_block_stride);

        // ── Buffer rotation: hand the prefetched LDGs to the next iter.
@@ -490,10 +540,10 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid(
 // ────────────────────────────────────────────────────────────────────────────
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t_in>
-void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
-    scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
-    int64_t const* slot_mapping, int64_t const* position_ids,
+template <typename scalar_t_in, int kNumHeadsQPadded>
+static void launchFusedDeepseekV4Templated(
+    scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
+    uint8_t* k_cache, int64_t const* slot_mapping, int64_t const* position_ids,
    float const* cos_sin_cache, float const eps, int const num_tokens_full,
    int const num_tokens_insert, int const num_heads_q,
    int const cache_block_size, int const kv_block_stride,
@@ -501,7 +551,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
  constexpr int kBlockSize = 256;
  constexpr int kWarpsPerBlock = kBlockSize / 32;
  int64_t const total_warps =
-      static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
+      static_cast<int64_t>(num_tokens_full) * (kNumHeadsQPadded + 1);
  int const grid =
      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);

@@ -532,46 +582,85 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(

  if (num_tokens_full < NUM_TOKEN_CUTOFF) {
    cudaLaunchKernelEx(
-        &config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
-        q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-        num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        &config,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in,
+                                                        kNumHeadsQPadded>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
        kv_block_stride);
  } else {
    config.gridDim = dim3(num_tokens_full);
    cudaLaunchKernelEx(
        &config,
-        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<scalar_t_in>,
-        q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
-        num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
+        fusedDeepseekV4QNormRopeKVRopeQuantInsertKernelReducedGrid<
+            scalar_t_in, kNumHeadsQPadded>,
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
+        eps, num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
        kv_block_stride);
  }
-
 #else
  // ROCm: use standard kernel launch syntax (no PDL/stream serialization)
  // clang-format off
-  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
+  fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in, kNumHeadsQPadded>
      <<<grid, kBlockSize, 0, stream>>>(
-          q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
-          eps, num_tokens_full, num_tokens_insert, num_heads_q,
+          q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,
+          cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q,
          cache_block_size, kv_block_stride);
 #endif
 }

+// Runtime dispatch into one of the precompiled `kNumHeadsQPadded`
+// instantiations.  Supported padded head counts: 8, 16, 32, 64, 128.
+template <typename scalar_t_in>
+void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
+    scalar_t_in const* q_in, scalar_t_in* q_out, scalar_t_in const* kv_in,
+    uint8_t* k_cache, int64_t const* slot_mapping,
+    int64_t const* position_ids, float const* cos_sin_cache, float const eps,
+    int const num_tokens_full, int const num_tokens_insert,
+    int const num_heads_q, int const num_heads_q_padded,
+    int const cache_block_size, int const kv_block_stride,
+    cudaStream_t stream) {
+#define DISPATCH(N)                                                         \
+  case N:                                                                   \
+    launchFusedDeepseekV4Templated<scalar_t_in, N>(                         \
+        q_in, q_out, kv_in, k_cache, slot_mapping, position_ids,            \
+        cos_sin_cache, eps, num_tokens_full, num_tokens_insert, num_heads_q, \
+        cache_block_size, kv_block_stride, stream);                         \
+    return;
+
+  switch (num_heads_q_padded) {
+    DISPATCH(8)
+    DISPATCH(16)
+    DISPATCH(32)
+    DISPATCH(64)
+    DISPATCH(128)
+    default:
+      TORCH_CHECK(false,
+                  "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert: "
+                  "unsupported num_heads_q_padded=",
+                  num_heads_q_padded,
+                  " (compiled instantiations: 8, 16, 32, 64, 128).");
+  }
+#undef DISPATCH
+}
+
 }  // namespace deepseek_v4_fused_ops
 }  // namespace vllm

 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
 // ────────────────────────────────────────────────────────────────────────────
-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q,                    // [N, H, 512] bf16, in place
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in,           // [N, num_heads_q, 512] bf16
    torch::Tensor const& kv,             // [N, 512] bf16 (read-only)
    torch::Tensor& k_cache,              // [num_blocks, block_bytes] uint8
    torch::Tensor const& slot_mapping,   // [N] int64
    torch::Tensor const& position_ids,   // [N] int64
    torch::Tensor const& cos_sin_cache,  // [max_pos, rope_dim] bf16
+    int64_t q_head_padded,               // padded Q head count for output
    double eps, int64_t cache_block_size) {
-  TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
+  TORCH_CHECK(q_in.is_cuda() && q_in.is_contiguous(),
+              "q_in must be contiguous CUDA");
  TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
  TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
  TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
@@ -579,9 +668,12 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
  TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
              "position_ids must be int64 CUDA");
  TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
-  TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
+  TORCH_CHECK(q_in.dim() == 3 && q_in.size(2) == 512,
+              "q_in shape [N, num_heads_q, 512]");
  TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
-  TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
+  TORCH_CHECK(q_in.dtype() == kv.dtype(), "q_in and kv dtype must match");
+  TORCH_CHECK(q_head_padded >= q_in.size(1),
+              "q_head_padded must be >= q_in.size(1) (num_heads_q)");
  TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
  TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
              "cos_sin_cache shape [max_pos, 64]");
@@ -591,32 +683,41 @@ void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
  // With DP padding, slot_mapping can be shorter than q/kv/positions.
  // Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
  // KV quant+insert runs only on the first slot_mapping.size(0) rows.
-  int const num_tokens_full = static_cast<int>(q.size(0));
+  int const num_tokens_full = static_cast<int>(q_in.size(0));
  int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
  TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
                  static_cast<int>(position_ids.size(0)) == num_tokens_full,
              "q/kv/position_ids row counts must match");
  TORCH_CHECK(num_tokens_insert <= num_tokens_full,
              "slot_mapping must not exceed q row count");
-  int const num_heads_q = static_cast<int>(q.size(1));
+  int const num_heads_q = static_cast<int>(q_in.size(1));
+  int const num_heads_q_padded = static_cast<int>(q_head_padded);
  int const cache_block_size_i = static_cast<int>(cache_block_size);
  int const kv_block_stride = static_cast<int>(k_cache.stride(0));

-  at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+  at::cuda::OptionalCUDAGuard device_guard(device_of(q_in));
  auto stream = at::cuda::getCurrentCUDAStream();

+  // Allocate the padded q output.  The kernel writes every element (live
+  // region gets RMSNorm+RoPE; pad region gets zeros), so `empty` is safe.
+  torch::Tensor q_out = torch::empty(
+      {q_in.size(0), q_head_padded, q_in.size(2)}, q_in.options());
+
  VLLM_DISPATCH_HALF_TYPES(
-      q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
+      q_in.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
        using qkv_scalar_t = scalar_t;
        vllm::deepseek_v4_fused_ops::
            launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
-                reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
+                reinterpret_cast<qkv_scalar_t const*>(q_in.data_ptr()),
+                reinterpret_cast<qkv_scalar_t*>(q_out.data_ptr()),
                reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
                reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
                reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
                reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
                cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
                num_tokens_full, num_tokens_insert, num_heads_q,
-                cache_block_size_i, kv_block_stride, stream);
+                num_heads_q_padded, cache_block_size_i, kv_block_stride,
+                stream);
      });
-}
+  return q_out;
+}
@@ -70,10 +70,11 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
                        torch::Tensor& weight, double epsilon);

-void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-    torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
+torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+    torch::Tensor const& q_in, torch::Tensor const& kv, torch::Tensor& k_cache,
    torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
-    torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
+    torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps,
+    int64_t cache_block_size);

 void apply_repetition_penalties_(torch::Tensor& logits,
                                 const torch::Tensor& prompt_mask,
@@ -99,9 +99,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // kernel launch.
  ops.def(
      "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
-      "Tensor! q, Tensor kv, Tensor! k_cache, "
+      "Tensor q_in, Tensor kv, Tensor! k_cache, "
      "Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
-      "float eps, int cache_block_size) -> ()");
+      "int q_head_padded, float eps, int cache_block_size) -> Tensor");
  ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
           &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);

@@ -220,7 +220,8 @@ COPY use_existing_torch.py use_existing_torch.py
 COPY pyproject.toml pyproject.toml
 RUN --mount=type=cache,target=/opt/uv/cache \
    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
-        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' requirements/cuda.txt; \
+        sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' requirements/cuda.txt; \
    fi \
    && if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
        echo "Installing torch nightly..." \
@@ -746,7 +747,8 @@ COPY requirements/common.txt /tmp/common.txt
 COPY requirements/cuda.txt /tmp/requirements-cuda.txt
 RUN --mount=type=cache,target=/opt/uv/cache \
    if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
-        sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^nvidia-cutlass-dsl\[cu13\]/nvidia-cutlass-dsl/' /tmp/requirements-cuda.txt; \
+        sed -i 's/^humming-kernels\[cu13\]/humming-kernels[cu12]/' /tmp/requirements-cuda.txt; \
    fi && \
    uv pip install --system -r /tmp/requirements-cuda.txt \
        --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
@@ -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
@@ -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
@@ -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).
@@ -25,4 +25,7 @@ nvidia-cutlass-dsl[cu13]==4.5.0
 quack-kernels>=0.3.3

 # Tokenspeed_MLA for faster mla with spec decode
-tokenspeed-mla==0.1.2
+tokenspeed-mla==0.1.2
+
+# Humming kernels for quantization gemm
+humming-kernels[cu13]==0.1.2
@@ -1,3 +1,6 @@
 lmcache >= 0.3.9
+# CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0
+# until a fixed newer release is verified for runtime images.
+cupy-cuda13x < 14.1.0
 nixl >= 1.1.0 # Required for disaggregated prefill
 mooncake-transfer-engine >= 0.3.8
@@ -1017,6 +1017,8 @@ def get_requirements() -> list[str]:
            if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
                # [cu13] extra is the default; strip it on CUDA 12 builds.
                req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl")
+            if "humming-kernels[cu13]" in req and cuda_major == "12":
+                req = req.replace("humming-kernels[cu13]", "humming-kernels[cu12]")
            modified_requirements.append(req)
        requirements = modified_requirements
    elif _is_hip():
@@ -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,
+        )
@@ -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
@@ -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")
@@ -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
@@ -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"]))
@@ -120,9 +120,19 @@ pytestmark = pytest.mark.skipif(
 )


-def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs):
-    torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
-        q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+def _call_fused(
+    q_in, q_head_padded, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+):
+    return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        q_in,
+        kv,
+        k_cache,
+        slot_mapping,
+        positions,
+        cos_sin_cache,
+        q_head_padded,
+        eps,
+        bs,
    )


@@ -130,8 +140,23 @@ def _call_fused(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)


@pytest.mark.parametrize("num_tokens", [1, 4, 17, 64, 2048])
-@pytest.mark.parametrize("n_heads", [8, 64])
-def test_q_path_matches_reference(num_tokens: int, n_heads: int):
+@pytest.mark.parametrize(
+    "n_heads,padded_heads",
+    [
+        # Each supported padded_heads instantiation: padded (n_heads <
+        # padded_heads) and unpadded (n_heads == padded_heads).
+        (1, 8),
+        (8, 8),
+        (8, 16),
+        (16, 16),
+        (16, 32),
+        (32, 32),
+        (8, 64),
+        (64, 64),
+        (64, 128),
+    ],
+)
+def test_q_path_matches_reference(num_tokens: int, n_heads: int, padded_heads: int):
    torch.manual_seed(0)
    device = "cuda"
    dtype = torch.bfloat16
@@ -156,10 +181,16 @@ def test_q_path_matches_reference(num_tokens: int, n_heads: int):
        num_blocks, bs, HEAD_BYTES, dtype=torch.uint8, device=device
    ).view(num_blocks, -1)
    slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
-    q_fused = q.clone()
-    _call_fused(q_fused, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs)
+    q_out = _call_fused(
+        q, padded_heads, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, bs
+    )

-    torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
+    if n_heads < padded_heads:
+        pad_region = q_out[:, n_heads:padded_heads]
+        assert pad_region.abs().max().item() == 0.0, (
+            "padded head slots must be exact zero"
+        )


 # ── Test 2: KV path round-trip byte/value parity ─────────────────────────────
@@ -201,11 +232,12 @@ def test_kv_path_matches_reference(num_tokens: int, block_size: int):
        kv_ref, k_cache_ref, slot_mapping, block_size=block_size
    )

-    # ── Fused path (dummy q, single head) ──────────────────────────────────
+    # ── Fused path (dummy q, padded to FlashMLA's min head count 64) ───────
    k_cache_fused = torch.zeros_like(k_cache_ref)
    q_dummy = torch.zeros(num_tokens, 1, HEAD_DIM, dtype=dtype, device=device)
-    _call_fused(
+    _ = _call_fused(
        q_dummy,
+        64,
        kv,
        k_cache_fused,
        slot_mapping,
@@ -298,8 +330,9 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):
    # Fused: pass full-sized q/kv/positions, shorter slot_mapping.
    q_dummy = torch.zeros(total, 1, HEAD_DIM, dtype=dtype, device=device)
    k_cache_fused = torch.zeros_like(k_cache_ref)
-    _call_fused(
+    _ = _call_fused(
        q_dummy,
+        64,
        kv,
        k_cache_fused,
        slot_mapping,
@@ -316,9 +349,26 @@ def test_kv_path_with_dp_padding(num_tokens: int, pad: int, block_size: int):


@pytest.mark.parametrize("num_tokens", [1, 4, 17, 2048])
-@pytest.mark.parametrize("n_heads", [8, 64])
+@pytest.mark.parametrize(
+    "n_heads,padded_heads",
+    [
+        # Each supported padded_heads instantiation: padded (n_heads <
+        # padded_heads) and unpadded (n_heads == padded_heads).
+        (1, 8),
+        (8, 8),
+        (8, 16),
+        (16, 16),
+        (16, 32),
+        (32, 32),
+        (8, 64),
+        (64, 64),
+        (64, 128),
+    ],
+)
@pytest.mark.parametrize("block_size", [16, 64])
-def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
+def test_combined_q_and_kv(
+    num_tokens: int, n_heads: int, padded_heads: int, block_size: int
+):
    torch.manual_seed(2)
    device = "cuda"
    dtype = torch.bfloat16
@@ -345,10 +395,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
    )

    # Fused single call.
-    q_fused = q.clone()
    k_cache_fused = torch.zeros_like(k_cache_ref)
-    _call_fused(
-        q_fused,
+    q_out = _call_fused(
+        q,
+        padded_heads,
        kv,
        k_cache_fused,
        slot_mapping,
@@ -358,5 +408,10 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
        block_size,
    )

-    torch.testing.assert_close(q_fused, q_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(q_out[:, :n_heads], q_ref, rtol=1e-2, atol=1e-2)
+    if n_heads < padded_heads:
+        pad_region = q_out[:, n_heads:padded_heads]
+        assert pad_region.abs().max().item() == 0.0, (
+            "padded head slots must be exact zero"
+        )
    torch.testing.assert_close(k_cache_fused, k_cache_ref, rtol=0, atol=0)
@@ -8,7 +8,7 @@ the existing separate operations (inverse RoPE via rotate_neox + FP8 quant
 via per_token_group_quant_fp8).

 The reference faithfully reproduces the exact flow in
-deepseek_v4/nvidia/ops/attention.py:295-310:
+deepseek_v4/attention.py:295-310:
  1. Apply inverse RoPE (NeoX style, last rope_dim=64 dims of each head)
  2. Reshape [T, H, head_dim] -> [T, G, D]
  3. Transpose+flatten to [G*T, D], quantize, reshape back
@@ -668,7 +668,7 @@ def _unfused_inv_rope_fp8_quant(
    nope_dim: int = NOPE_DIM,
    rope_dim: int = ROPE_DIM,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Unfused path matching deepseek_v4/nvidia/ops/attention.py:295-310.
+    """Unfused path matching deepseek_v4/attention.py:295-310.

    Uses the production CUDA RoPE kernel + per_token_group_quant_fp8.
    """
@@ -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
@@ -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
@@ -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.",
@@ -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
@@ -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
+        )
@@ -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.
@@ -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.
@@ -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,
@@ -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
@@ -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)
@@ -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.")
@@ -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
@@ -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,
@@ -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
@@ -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
@@ -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,
        )

@@ -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] = []
@@ -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)
@@ -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
@@ -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

@@ -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)."
        ),
    )

@@ -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")
@@ -7,6 +7,7 @@ from vllm.distributed import (
    tensor_model_parallel_all_gather,
    tensor_model_parallel_all_reduce,
 )
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.triton_utils.allocation import set_triton_allocator
 from vllm.utils.torch_utils import direct_register_custom_op
@@ -406,7 +407,7 @@ def _run_fused_moe_lora_one_shot(

    # NPID_FACTOR heuristic: scale N-axis parallelism when base CTA count is
    # short of saturating the SM array. Cap by the cost of redundant shrink.
-    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    sm_count = current_platform.num_compute_units(device.index)
    base_programs = max(M_blocks * num_slices * grid_lora_dim, 1)
    shrink_ratio = K / max(K + N_per_slice, 1)
    max_npid_by_budget = max(1, int(1.5 / max(shrink_ratio, 1e-3)) + 1)
@@ -786,7 +787,7 @@ def _run_fused_moe_lora_small_batch(
    N_tiles = triton.cdiv(N_per_slice, BLOCK_N)
    pair_slices = M_grid * num_slices

-    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    sm_count = current_platform.num_compute_units(device.index)
    n_tiles_per_program = _pick_small_batch_chunk(pair_slices, N_tiles, sm_count)
    n_chunks = triton.cdiv(N_tiles, n_tiles_per_program)
    work_total = pair_slices * n_chunks
@@ -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),
    )


@@ -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),
+        )
@@ -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,
@@ -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,
    )

@@ -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
@@ -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,
        )
@@ -43,12 +43,9 @@ from vllm.model_executor.parameter import (
    RowvLLMParameter,
 )
 from vllm.model_executor.utils import set_weight_attrs
+from vllm.platforms import current_platform

-if TYPE_CHECKING:
-    from vllm.model_executor.models.utils import WeightsMapper
-
-
-try:
+if current_platform.is_cuda():
    from humming.dtypes import DataType
    from humming.layer import HummingMethod
    from humming.schema import (
@@ -65,16 +62,17 @@ try:
        HummingIndexedExperts,
        get_humming_moe_gemm_type,
    )
-except ModuleNotFoundError:
-    HummingMethod = None

-
-def assert_humming_available():
-    assert HummingMethod is not None, (
-        "humming is not available, please run "
-        "'pip install git+https://github.com/inclusionAI/humming' to install it."
+if TYPE_CHECKING:
+    from humming.schema import (
+        BaseInputSchema,
+        BaseWeightSchema,
+        HummingInputSchema,
+        HummingWeightSchema,
    )

+    from vllm.model_executor.models.utils import WeightsMapper
+

 def prepare_padded_shape(shape, x):
    padded_shape = math.ceil(shape / x) * x
@@ -186,7 +184,6 @@ class HummingConfig(QuantizationConfig):
    packed_modules_mapping: dict[str, list[str]] = {}

    def __init__(self, full_config: dict[str, Any] | None = None):
-        assert_humming_available()
        self.full_config: dict[str, Any] = full_config or {}

    @classmethod
@@ -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":
@@ -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.
@@ -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
@@ -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,
@@ -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
@@ -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,
@@ -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
@@ -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")
@@ -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"),
        )

@@ -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:
@@ -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:
@@ -11,17 +11,12 @@ import torch.nn as nn
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import (
-    get_ep_group,
    get_pp_group,
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
 )
-from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
 from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
-from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
-    fused_topk_bias,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
@@ -52,16 +47,13 @@ from vllm.model_executor.models.utils import (
    make_layers,
    maybe_prefix,
 )
-from vllm.model_executor.utils import set_weight_attrs
-from vllm.models.deepseek_v4.nvidia.ops.attention import (
+from vllm.models.deepseek_v4.attention import (
    DeepseekV4Indexer,
    DeepseekV4MLAModules,
    DeepseekV4MultiHeadLatentAttentionWrapper,
 )
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
-from vllm.triton_utils import tl, triton
-from vllm.utils.torch_utils import direct_register_custom_op


 class DeepseekV4MLP(nn.Module):
@@ -115,501 +107,6 @@ class DeepseekV4MLP(nn.Module):
        return x


-@triton.jit
-def _deepseek_v4_stage_mega_moe_inputs_kernel(
-    hidden_states,
-    x_fp8,
-    x_sf,
-    topk_ids,
-    topk_weights,
-    topk_idx_out,
-    topk_weights_out,
-    hidden_stride_m: tl.constexpr,
-    hidden_stride_k: tl.constexpr,
-    x_stride_m: tl.constexpr,
-    x_stride_k: tl.constexpr,
-    x_sf_stride_m: tl.constexpr,
-    x_sf_stride_k: tl.constexpr,
-    topk_ids_stride_m: tl.constexpr,
-    topk_ids_stride_k: tl.constexpr,
-    topk_weights_stride_m: tl.constexpr,
-    topk_weights_stride_k: tl.constexpr,
-    topk_idx_stride_m: tl.constexpr,
-    topk_idx_stride_k: tl.constexpr,
-    topk_weights_out_stride_m: tl.constexpr,
-    topk_weights_out_stride_k: tl.constexpr,
-    hidden_size: tl.constexpr,
-    top_k: tl.constexpr,
-    BLOCK_K: tl.constexpr,
-    GROUP_K: tl.constexpr,
-    BLOCK_TOPK: tl.constexpr,
-) -> None:
-    token_id = tl.program_id(0)
-    k_block_id = tl.program_id(1)
-
-    k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
-    k_mask = k_offsets < hidden_size
-    hidden = tl.load(
-        hidden_states + token_id * hidden_stride_m + k_offsets * hidden_stride_k,
-        mask=k_mask,
-        other=0.0,
-    ).to(tl.float32)
-
-    num_groups: tl.constexpr = BLOCK_K // GROUP_K
-    hidden_groups = tl.reshape(tl.abs(hidden), [num_groups, GROUP_K])
-    amax = tl.max(hidden_groups, axis=1)
-    amax = tl.maximum(amax, 1.0e-4)
-
-    scale = amax / 448.0
-    scale_bits = scale.to(tl.uint32, bitcast=True)
-    scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
-        tl.uint32
-    )
-    scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
-    rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
-
-    hidden_groups = tl.reshape(hidden, [num_groups, GROUP_K])
-    scaled = hidden_groups * (1.0 / rounded_scale)[:, None]
-    scaled = tl.reshape(scaled, [BLOCK_K])
-    fp8 = scaled.to(tl.float8e4nv)
-    tl.store(
-        x_fp8 + token_id * x_stride_m + k_offsets * x_stride_k,
-        fp8,
-        mask=k_mask,
-    )
-
-    scale_offsets = tl.arange(0, num_groups)
-    packed_scale = tl.sum(scale_exp << (scale_offsets * 8), axis=0).to(tl.int32)
-    tl.store(
-        x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
-        packed_scale,
-    )
-
-    if k_block_id == 0:
-        topk_offsets = tl.arange(0, BLOCK_TOPK)
-        topk_mask = topk_offsets < top_k
-
-        ids = tl.load(
-            topk_ids + token_id * topk_ids_stride_m + topk_offsets * topk_ids_stride_k,
-            mask=topk_mask,
-            other=0,
-        ).to(tl.int64)
-        tl.store(
-            topk_idx_out
-            + token_id * topk_idx_stride_m
-            + topk_offsets * topk_idx_stride_k,
-            ids,
-            mask=topk_mask,
-        )
-
-        weights = tl.load(
-            topk_weights
-            + token_id * topk_weights_stride_m
-            + topk_offsets * topk_weights_stride_k,
-            mask=topk_mask,
-            other=0.0,
-        )
-        tl.store(
-            topk_weights_out
-            + token_id * topk_weights_out_stride_m
-            + topk_offsets * topk_weights_out_stride_k,
-            weights,
-            mask=topk_mask,
-        )
-
-
-def _stage_deepseek_v4_mega_moe_inputs(
-    hidden_states: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    x_fp8: torch.Tensor,
-    x_sf: torch.Tensor,
-    topk_idx_out: torch.Tensor,
-    topk_weights_out: torch.Tensor,
-) -> None:
-    num_tokens, hidden_size = hidden_states.shape
-    if num_tokens == 0:
-        return
-    if hidden_size % 128 != 0:
-        raise ValueError(
-            "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
-            "a multiple of 128."
-        )
-    top_k = topk_ids.shape[1]
-    if topk_weights.shape != topk_ids.shape:
-        raise ValueError(
-            "DeepSeek V4 MegaMoE input staging requires topk_weights and "
-            "topk_ids to have the same shape."
-        )
-
-    block_k = 128
-    grid = (num_tokens, triton.cdiv(hidden_size, block_k))
-    block_topk = triton.next_power_of_2(top_k)
-    _deepseek_v4_stage_mega_moe_inputs_kernel[grid](
-        hidden_states,
-        x_fp8,
-        x_sf,
-        topk_ids,
-        topk_weights,
-        topk_idx_out,
-        topk_weights_out,
-        hidden_states.stride(0),
-        hidden_states.stride(1),
-        x_fp8.stride(0),
-        x_fp8.stride(1),
-        x_sf.stride(0),
-        x_sf.stride(1),
-        topk_ids.stride(0),
-        topk_ids.stride(1),
-        topk_weights.stride(0),
-        topk_weights.stride(1),
-        topk_idx_out.stride(0),
-        topk_idx_out.stride(1),
-        topk_weights_out.stride(0),
-        topk_weights_out.stride(1),
-        hidden_size,
-        top_k,
-        BLOCK_K=block_k,
-        GROUP_K=32,
-        BLOCK_TOPK=block_topk,
-        num_warps=4,
-    )
-
-
-def make_deepseek_v4_expert_params_mapping(
-    num_experts: int,
-) -> list[tuple[str, str, int, str]]:
-    return [
-        (
-            "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_",
-            f"experts.{expert_id}.{weight_name}.",
-            expert_id,
-            shard_id,
-        )
-        for expert_id in range(num_experts)
-        for shard_id, weight_name in [
-            ("w1", "w1"),
-            ("w2", "w2"),
-            ("w3", "w3"),
-        ]
-    ]
-
-
-class DeepseekV4MegaMoEExperts(nn.Module):
-    _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        *,
-        num_experts: int,
-        num_local_experts: int,
-        experts_start_idx: int,
-        top_k: int,
-        hidden_size: int,
-        intermediate_size: int,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.prefix = prefix
-        self.num_experts = num_experts
-        self.num_local_experts = num_local_experts
-        self.experts_start_idx = experts_start_idx
-        self.experts_end_idx = experts_start_idx + num_local_experts
-        self.top_k = top_k
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
-
-        weight_attrs = {"weight_loader": self.weight_loader}
-        self.w13_weight = nn.Parameter(
-            torch.zeros(
-                num_local_experts,
-                2 * intermediate_size,
-                hidden_size // 2,
-                dtype=torch.uint8,
-            ),
-            requires_grad=False,
-        )
-        set_weight_attrs(self.w13_weight, weight_attrs)
-
-        self.w13_weight_scale = nn.Parameter(
-            torch.zeros(
-                num_local_experts,
-                2 * intermediate_size,
-                hidden_size // 32,
-                dtype=torch.uint8,
-            ),
-            requires_grad=False,
-        )
-        set_weight_attrs(self.w13_weight_scale, weight_attrs)
-        self.w13_weight_scale.quant_method = "block"
-
-        self.w2_weight = nn.Parameter(
-            torch.zeros(
-                num_local_experts,
-                hidden_size,
-                intermediate_size // 2,
-                dtype=torch.uint8,
-            ),
-            requires_grad=False,
-        )
-        set_weight_attrs(self.w2_weight, weight_attrs)
-
-        self.w2_weight_scale = nn.Parameter(
-            torch.zeros(
-                num_local_experts,
-                hidden_size,
-                intermediate_size // 32,
-                dtype=torch.uint8,
-            ),
-            requires_grad=False,
-        )
-        set_weight_attrs(self.w2_weight_scale, weight_attrs)
-        self.w2_weight_scale.quant_method = "block"
-
-        self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
-        self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
-
-        # Register in the static forward context so the custom-op wrapper
-        # can look up this module by name from within a torch.compile graph.
-        compilation_config = vllm_config.compilation_config
-        if prefix in compilation_config.static_forward_context:
-            raise ValueError(f"Duplicate layer name: {prefix}")
-        compilation_config.static_forward_context[prefix] = self
-
-    def _map_global_expert_id(self, expert_id: int) -> int:
-        if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx:
-            return -1
-        return expert_id - self.experts_start_idx
-
-    def weight_loader(
-        self,
-        param: nn.Parameter,
-        loaded_weight: torch.Tensor,
-        weight_name: str,
-        shard_id: str,
-        expert_id: int,
-        return_success: bool = False,
-    ) -> bool | None:
-        local_expert_id = self._map_global_expert_id(expert_id)
-        if local_expert_id == -1:
-            return False if return_success else None
-
-        expert_data = param.data[local_expert_id]
-        if shard_id in ("w1", "w3"):
-            if "w13_" not in weight_name:
-                return False if return_success else None
-            shard_offset = 0 if shard_id == "w1" else self.intermediate_size
-            expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size)
-        elif shard_id == "w2":
-            if "w2_" not in weight_name:
-                return False if return_success else None
-        else:
-            raise ValueError(f"Unsupported expert shard id: {shard_id}")
-
-        if expert_data.shape != loaded_weight.shape:
-            raise ValueError(
-                f"DeepSeek V4 MegaMoE expert weight shape mismatch for "
-                f"{weight_name}: parameter shard {tuple(expert_data.shape)} "
-                f"vs checkpoint {tuple(loaded_weight.shape)}"
-            )
-        expert_data.copy_(loaded_weight)
-        return True if return_success else None
-
-    @staticmethod
-    def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor:
-        return (sf.to(torch.int32) << 23).view(torch.float32)
-
-    def _check_runtime_supported(self) -> None:
-        if not torch.cuda.is_available():
-            raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.")
-        device = self.w13_weight.device
-        if device.type != "cuda":
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA."
-            )
-        if torch.cuda.get_device_capability(device)[0] != 10:
-            raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.")
-        if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0:
-            raise ValueError(
-                "DeepGEMM MegaMoE requires hidden and intermediate sizes "
-                "to be multiples of 128."
-            )
-
-    def finalize_weights(self) -> None:
-        if self._transformed_l1_weights is not None:
-            return
-
-        self._check_runtime_supported()
-        import vllm.third_party.deep_gemm as deep_gemm
-
-        w13_scale = deep_gemm.transform_sf_into_required_layout(
-            self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
-            2 * self.intermediate_size,
-            self.hidden_size,
-            (1, 32),
-            self.num_local_experts,
-        )
-        w2_scale = deep_gemm.transform_sf_into_required_layout(
-            self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
-            self.hidden_size,
-            self.intermediate_size,
-            (1, 32),
-            self.num_local_experts,
-        )
-        self._transformed_l1_weights, self._transformed_l2_weights = (
-            deep_gemm.transform_weights_for_mega_moe(
-                (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
-                (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
-            )
-        )
-        # Drop the original loader-side parameters: the MegaMoE kernels only
-        # consume the transformed views above. transform_weights_for_mega_moe
-        # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
-        # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
-        # aliases the original storage, and _transformed_l2_weights still holds
-        # it, so the storage stays live after we drop the Parameter.
-        self.w13_weight = None
-        self.w13_weight_scale = None
-        self.w2_weight = None
-        self.w2_weight_scale = None
-
-    def get_symm_buffer(self):
-        import vllm.third_party.deep_gemm as deep_gemm
-
-        group = get_ep_group().device_group
-        device = torch.accelerator.current_device_index()
-        key = (
-            id(group),
-            device,
-            self.num_experts,
-            self.max_num_tokens,
-            self.top_k,
-            self.hidden_size,
-            self.intermediate_size,
-        )
-        symm_buffer = self._symm_buffer_cache.get(key)
-        if symm_buffer is None:
-            symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
-                group,
-                self.num_experts,
-                self.max_num_tokens,
-                self.top_k,
-                self.hidden_size,
-                self.intermediate_size,
-            )
-            self._symm_buffer_cache[key] = symm_buffer
-        return symm_buffer
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        *,
-        activation_clamp: float | None,
-        fast_math: bool = True,
-    ) -> torch.Tensor:
-        if hidden_states.shape[0] > self.max_num_tokens:
-            raise ValueError(
-                f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, "
-                f"but the symmetric buffer was sized for {self.max_num_tokens}."
-            )
-        y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
-        torch.ops.vllm.deepseek_v4_mega_moe_experts(
-            hidden_states,
-            topk_weights,
-            topk_ids,
-            y,
-            self.prefix,
-            activation_clamp,
-            fast_math,
-        )
-        return y
-
-    def _run_mega_moe(
-        self,
-        hidden_states: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        y: torch.Tensor,
-        activation_clamp: float | None,
-        fast_math: bool,
-    ) -> None:
-        import vllm.third_party.deep_gemm as deep_gemm
-
-        symm_buffer = self.get_symm_buffer()
-        num_tokens = hidden_states.shape[0]
-        _stage_deepseek_v4_mega_moe_inputs(
-            hidden_states,
-            topk_weights,
-            topk_ids,
-            symm_buffer.x[:num_tokens],
-            symm_buffer.x_sf[:num_tokens],
-            symm_buffer.topk_idx[:num_tokens],
-            symm_buffer.topk_weights[:num_tokens],
-        )
-
-        # This method must have been already called during the weight loading phase.
-        # We call it again here to cover the dummy weight loading case.
-        self.finalize_weights()
-
-        assert self._transformed_l1_weights is not None
-        assert self._transformed_l2_weights is not None
-        deep_gemm.fp8_fp4_mega_moe(
-            y,
-            self._transformed_l1_weights,
-            self._transformed_l2_weights,
-            symm_buffer,
-            activation_clamp=activation_clamp,
-            fast_math=fast_math,
-        )
-
-
-DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True  # type: ignore[attr-defined]
-
-
-def _deepseek_v4_mega_moe_experts_op(
-    hidden_states: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    out: torch.Tensor,
-    layer_name: str,
-    activation_clamp: float | None,
-    fast_math: bool,
-) -> None:
-    self = get_forward_context().no_compile_layers[layer_name]
-    self._run_mega_moe(
-        hidden_states,
-        topk_weights,
-        topk_ids,
-        out,
-        activation_clamp,
-        fast_math,
-    )
-
-
-def _deepseek_v4_mega_moe_experts_op_fake(
-    hidden_states: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    out: torch.Tensor,
-    layer_name: str,
-    activation_clamp: float | None,
-    fast_math: bool,
-) -> None:
-    return None
-
-
-direct_register_custom_op(
-    op_name="deepseek_v4_mega_moe_experts",
-    op_func=_deepseek_v4_mega_moe_experts_op,
-    mutates_args=["out"],
-    fake_impl=_deepseek_v4_mega_moe_experts_op_fake,
-)
-
-
 class DeepseekV4MoE(nn.Module):
    def __init__(
        self,
@@ -622,15 +119,6 @@ class DeepseekV4MoE(nn.Module):
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        self.prefix = prefix
-        self.use_mega_moe = (
-            vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
-        )
-        if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE currently requires expert parallel. "
-                "Enable it with --enable-expert-parallel, or pick a different "
-                "moe backend."
-            )

        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
        self.hidden_size = config.hidden_size
@@ -641,16 +129,6 @@ class DeepseekV4MoE(nn.Module):
        self.swiglu_limit = config.swiglu_limit
        self.renormalize = config.norm_topk_prob
        self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
-        if self.use_mega_moe and self.scoring_func != "sqrtsoftplus":
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."
-            )
-        if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="
-                f"{config.expert_dtype!r}. Drop --kernel-config moe_backend="
-                "deep_gemm_mega_moe for this checkpoint."
-            )

        self.gate = GateLinear(
            input_size=config.hidden_size,
@@ -663,7 +141,7 @@ class DeepseekV4MoE(nn.Module):
        self.gate.e_score_correction_bias = None
        self.gate.tid2eid = None
        is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
-        self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
+        self.hash_indices_dtype = torch.int32
        if is_hash_moe:
            # hash MoE doesn't use e_score_correction_bias
            # Use randint instead of empty to avoid garbage values causing
@@ -694,47 +172,10 @@ class DeepseekV4MoE(nn.Module):
                hidden_act=config.hidden_act,
                swiglu_limit=self.swiglu_limit,
                quant_config=quant_config,
-                reduce_results=self.use_mega_moe,
+                reduce_results=False,
                prefix=f"{prefix}.shared_experts",
            )

-        if self.use_mega_moe:
-            self._init_mega_moe_experts(vllm_config, config, prefix)
-        else:
-            self._init_fused_moe_experts(config, quant_config, prefix)
-
-    def _init_mega_moe_experts(
-        self,
-        vllm_config: VllmConfig,
-        config,
-        prefix: str,
-    ) -> None:
-        self.ep_group = get_ep_group()
-        self.ep_size = self.ep_group.world_size
-        self.ep_rank = self.ep_group.rank_in_group
-        assert config.n_routed_experts % self.ep_size == 0
-
-        self.n_local_experts = config.n_routed_experts // self.ep_size
-        self.experts_start_idx = self.ep_rank * self.n_local_experts
-        self.experts_end_idx = self.experts_start_idx + self.n_local_experts
-
-        self.experts = DeepseekV4MegaMoEExperts(
-            vllm_config,
-            num_experts=config.n_routed_experts,
-            num_local_experts=self.n_local_experts,
-            experts_start_idx=self.experts_start_idx,
-            top_k=config.num_experts_per_tok,
-            hidden_size=config.hidden_size,
-            intermediate_size=config.moe_intermediate_size,
-            prefix=f"{prefix}.experts",
-        )
-
-    def _init_fused_moe_experts(
-        self,
-        config,
-        quant_config,
-        prefix: str,
-    ) -> None:
        self.tp_rank = get_tensor_model_parallel_rank()
        assert config.n_routed_experts % self.tp_size == 0

@@ -766,44 +207,6 @@ class DeepseekV4MoE(nn.Module):
        if self.gate.tid2eid is not None and input_ids is None:
            raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")

-        if not self.use_mega_moe:
-            return self._forward_fused_moe(hidden_states, input_ids)
-
-        org_shape = hidden_states.shape
-        router_logits, _ = self.gate(hidden_states)
-        topk_weights, topk_ids = fused_topk_bias(
-            hidden_states=hidden_states,
-            gating_output=router_logits,
-            scoring_func=self.scoring_func,
-            e_score_correction_bias=self.gate.e_score_correction_bias.data
-            if self.gate.e_score_correction_bias is not None
-            else None,
-            topk=self.n_activated_experts,
-            renormalize=self.renormalize,
-            indices_type=self.hash_indices_dtype,
-            input_tokens=input_ids,
-            hash_indices_table=self.gate.tid2eid,
-            routed_scaling_factor=self.routed_scaling_factor,
-        )
-        activation_clamp = (
-            float(self.swiglu_limit) if self.swiglu_limit is not None else None
-        )
-        final_hidden_states = self.experts(
-            hidden_states,
-            topk_weights,
-            topk_ids,
-            activation_clamp=activation_clamp,
-        )
-
-        if self.shared_experts is not None:
-            shared_output = self.shared_experts(hidden_states)
-            final_hidden_states += shared_output
-
-        return final_hidden_states.view(org_shape)
-
-    def _forward_fused_moe(
-        self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
-    ) -> torch.Tensor:
        org_shape = hidden_states.shape
        if self.experts.is_internal_router:
            # In this case, the gate/router runs inside the FusedMoE class
@@ -822,10 +225,6 @@ class DeepseekV4MoE(nn.Module):

        return final_hidden_states.view(org_shape)

-    def finalize_mega_moe_weights(self) -> None:
-        if self.use_mega_moe:
-            self.experts.finalize_weights()
-

 class DeepseekV4Attention(nn.Module):
    def __init__(
@@ -1211,15 +610,6 @@ class DeepseekV4Model(nn.Module):
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        self.config = config
-        self.use_mega_moe = (
-            vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
-        )
-        if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:
-            raise NotImplementedError(
-                "DeepSeek V4 MegaMoE currently requires expert parallel. "
-                "Enable it with --enable-expert-parallel, or pick a different "
-                "moe backend."
-            )
        self.vocab_size = config.vocab_size
        self.hc_eps = config.hc_eps
        self.hc_mult = config.hc_mult
@@ -1348,9 +738,6 @@ class DeepseekV4Model(nn.Module):
            assert intermediate_tensors is not None
            hidden_states = intermediate_tensors["hidden_states"]

-        if self.use_mega_moe:
-            input_ids = input_ids.to(torch.int64)
-
        residual, post_mix, res_mix = None, None, None
        for layer in islice(self.layers, self.start_layer, self.end_layer):
            hidden_states, residual, post_mix, res_mix = layer(
@@ -1482,9 +869,6 @@ class DeepseekV4Model(nn.Module):
        return loaded_params

    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
-        if first_layer.ffn.use_mega_moe:
-            return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
        return FusedMoE.make_expert_params_mapping(
@@ -1495,10 +879,6 @@ class DeepseekV4Model(nn.Module):
            num_experts=self.config.n_routed_experts,
        )

-    def finalize_mega_moe_weights(self) -> None:
-        for layer in islice(self.layers, self.start_layer, self.end_layer):
-            layer.ffn.finalize_mega_moe_weights()
-

 def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
    if expert_dtype == "fp4":
@@ -1601,7 +981,6 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP):
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
        loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-        self.model.finalize_mega_moe_weights()
        return loaded_params

    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
@@ -40,10 +40,7 @@ from vllm.model_executor.models.utils import maybe_prefix
 from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors

-from .model import (
-    DeepseekV4DecoderLayer,
-    make_deepseek_v4_expert_params_mapping,
-)
+from .model import DeepseekV4DecoderLayer

 logger = init_logger(__name__)

@@ -330,19 +327,13 @@ class DeepSeekV4MTP(nn.Module):
        head_rank_end = n_local_head * (tp_rank + 1)

        # Pre-compute expert mapping ONCE.
-        first_layer = next(iter(self.model.layers.values()))
-        if first_layer.mtp_block.ffn.use_mega_moe:
-            expert_mapping = make_deepseek_v4_expert_params_mapping(
-                self.config.n_routed_experts
-            )
-        else:
-            expert_mapping = FusedMoE.make_expert_params_mapping(
-                self,
-                ckpt_gate_proj_name="w1",
-                ckpt_down_proj_name="w2",
-                ckpt_up_proj_name="w3",
-                num_experts=self.config.n_routed_experts,
-            )
+        expert_mapping = FusedMoE.make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="w1",
+            ckpt_down_proj_name="w2",
+            ckpt_up_proj_name="w3",
+            num_experts=self.config.n_routed_experts,
+        )

        # FP8 experts register ``..._weight_scale_inv`` (block_quant) while
        # FP4/MXFP4 experts register ``..._weight_scale``. Choose the suffix
@@ -465,14 +456,9 @@ class DeepSeekV4MTP(nn.Module):
                    f"Use a checkpoint that includes MTP layer weights, "
                    f"or disable speculative decoding."
                )
-        self.finalize_mega_moe_weights()
        logger.info_once("MTP draft model loaded: %d params", len(loaded_params))
        return loaded_params

-    def finalize_mega_moe_weights(self) -> None:
-        for layer in self.model.layers.values():
-            layer.mtp_block.ffn.finalize_mega_moe_weights()
-
    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
        """
        Rewrite the weight name to match the format of the original model.
@@ -32,7 +32,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
 from vllm.v1.worker.workspace import current_workspace_manager

 if TYPE_CHECKING:
-    from vllm.models.deepseek_v4.nvidia.ops.attention import (
+    from vllm.models.deepseek_v4.attention import (
        DeepseekV4MLAAttention,
    )

@@ -592,6 +592,10 @@ class DeepseekV4ROCMAiterMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):

    backend_cls = DeepseekV4ROCMAiterMLASparseBackend

+    @classmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        return num_heads
+
    @classmethod
    def forward_mqa(  # type: ignore[override]
        cls,
@@ -156,18 +156,6 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        self.head_dim = head_dim
        self.scale = scale

-        # FlashMLA sparse kernel only supports 64 or 128 heads; pad up to the
-        # next supported size. Must match DeepseekV4MLAAttention.padded_heads.
-        if num_heads <= 64:
-            self.padded_heads = 64
-        elif num_heads <= 128:
-            self.padded_heads = 128
-        else:
-            raise ValueError(
-                f"DeepseekV4 attention does not support {num_heads} heads "
-                "(must be <= 128)."
-            )
-
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.window_size = window_size
@@ -263,6 +251,9 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
            indexer=self.indexer,
            topk_indices_buffer=self.topk_indices_buffer,
        )
+        # Mirror the inner layer's padded head count (single source of truth).
+        self.padded_heads = self.mla_attn.padded_heads
+
        # Register this layer in the compilation config's static forward context
        # This allows the custom op to retrieve the layer during execution
        compilation_config = mla_modules.vllm_config.compilation_config
@@ -450,7 +441,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

            def wq_b_kv_insert() -> torch.Tensor:
                q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
-                self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+                q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
                return q

            # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs
@@ -484,7 +475,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):

            def wq_b_kv_insert() -> torch.Tensor:
                q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
-                self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
+                q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
                return q

            q, _ = maybe_execute_in_parallel(
@@ -497,12 +488,7 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        else:
            # SWA-only layer: no compressor, no overlap.
            q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
-            self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
-
-        # Pad q to FlashMLA-required head count (64 or 128)
-        if self.n_local_heads < self.padded_heads:
-            pad_size = self.padded_heads - self.n_local_heads
-            q = F.pad(q, (0, 0, 0, pad_size), value=0.0)
+            q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)

        # MLA attention writes into the pre-allocated `out` buffer
        # ([num_tokens, padded_heads, head_dim]).
@@ -516,9 +502,17 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        attn_metadata: (
            dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
        ),
-    ) -> None:
+    ) -> torch.Tensor:
        if not isinstance(attn_metadata, dict):
-            return
+            # Profile run: kernel doesn't fire; produce a padded tensor so
+            # downstream FlashMLA gets the right shape.
+            if self.n_local_heads < self.padded_heads:
+                return F.pad(
+                    q,
+                    (0, 0, 0, self.padded_heads - self.n_local_heads),
+                    value=0.0,
+                )
+            return q

        swa_metadata = cast(
            "DeepseekSparseSWAMetadata | None",
@@ -530,16 +524,19 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
        swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)

        # Horizontally fused:
-        #   Q side:  q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE
+        #   Q side:  q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE,
+        #            with zero-fill for the padding head slots.  The kernel
+        #            allocates and returns the padded q tensor.
        #   KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert
        # kv is unchanged; mla_attn reads kv solely via swa_kv_cache.
-        torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
+        return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
            q,
            kv,
            swa_kv_cache_2d,
            swa_metadata.slot_mapping,
            positions.to(torch.int64),
            self.rotary_emb.cos_sin_cache,
+            self.padded_heads,
            self.eps,
            swa_metadata.block_size,
        )
@@ -607,9 +604,6 @@ direct_register_custom_op(


 class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
-    # FlashMLA FP8 sparse only supports 64 or 128 heads
-    SUPPORTED_HEAD_COUNTS = (64, 128)
-
    def __init__(
        self,
        num_heads: int,
@@ -655,19 +649,8 @@ class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
        self.aux_stream = aux_stream
        self.ln_events = [torch.cuda.Event(), torch.cuda.Event()]

-        # Determine padded head count for FlashMLA
-        if num_heads not in self.SUPPORTED_HEAD_COUNTS:
-            if num_heads < 64:
-                self.padded_heads = 64
-            elif num_heads < 128:
-                self.padded_heads = 128
-            else:
-                raise ValueError(
-                    f"DeepseekV4MLAAttention does not support {num_heads} heads. "
-                    f"Supported: <= 128 (will be padded to 64 or 128)"
-                )
-        else:
-            self.padded_heads = num_heads
+        # Padded Q head count is dictated by the selected impl.
+        self.padded_heads = self.impl_cls.get_padded_num_q_heads(num_heads)

        # Store attention sink
        assert attn_sink is not None
@@ -173,33 +173,6 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):


 class DeepseekCompressor(nn.Module):
-    _compressed_kv_buffers: ClassVar[dict[tuple[str, int, int], torch.Tensor]] = {}
-
-    @classmethod
-    def _get_compressed_kv_buffer(
-        cls,
-        device: str,
-        max_num_tokens: int,
-        head_dim: int,
-    ) -> torch.Tensor:
-        if device == "cuda" and torch.accelerator.is_available():
-            device_key = f"cuda:{torch.accelerator.current_device_index()}"
-            alloc_device = torch.device(device_key)
-        else:
-            device_key = str(device)
-            alloc_device = torch.device(device)
-
-        key = (device_key, max_num_tokens, head_dim)
-        buffer = cls._compressed_kv_buffers.get(key)
-        if buffer is None:
-            buffer = torch.empty(
-                (max_num_tokens, head_dim),
-                dtype=torch.float32,
-                device=alloc_device,
-            )
-            cls._compressed_kv_buffers[key] = buffer
-        return buffer
-
    def __init__(
        self,
        vllm_config: VllmConfig,
@@ -276,11 +249,6 @@ class DeepseekCompressor(nn.Module):
            self._fused_sparse_kernel = (
                _fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl
            )
-            self._compressed_kv_buffer = self._get_compressed_kv_buffer(
-                self.device,
-                vllm_config.scheduler_config.max_num_batched_tokens,
-                self.head_dim,
-            )
            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
@@ -407,7 +375,11 @@ class DeepseekCompressor(nn.Module):
                    overlap=self.overlap,
                )
            else:
-                compressed_kv = self._compressed_kv_buffer[:num_actual]
+                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,
@@ -28,7 +28,7 @@ from vllm.v1.attention.ops.flashmla import (
 from vllm.v1.worker.workspace import current_workspace_manager

 if TYPE_CHECKING:
-    from vllm.models.deepseek_v4.nvidia.ops.attention import (
+    from vllm.models.deepseek_v4.attention import (
        DeepseekV4MLAAttention,
    )
    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -63,6 +63,18 @@ class DeepseekV4SparseMLAAttentionImpl(SparseMLAAttentionImpl[FlashMLASparseMeta
    ) -> None:
        raise NotImplementedError

+    @classmethod
+    @abstractmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        """Q head count the backend wants q allocated at.
+
+        The MLA wrapper allocates the q/output buffers at
+        ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must
+        satisfy ``result >= num_heads``. Backends with no padding constraint
+        return ``num_heads``.
+        """
+        raise NotImplementedError
+

 class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
    @staticmethod
@@ -104,6 +116,16 @@ class DeepseekV4FlashMLASparseImpl(DeepseekV4SparseMLAAttentionImpl):

    backend_cls = DeepseekV4FlashMLASparseBackend

+    @classmethod
+    def get_padded_num_q_heads(cls, num_heads: int) -> int:
+        # FP8 decode kernel only supports h_q = 64 or 128.
+        if num_heads > 128:
+            raise ValueError(
+                f"DeepseekV4 FlashMLA does not support {num_heads} heads "
+                "(FP8 decode kernel requires h_q in {64, 128})."
+            )
+        return 64 if num_heads <= 64 else 128
+
    @classmethod
    def forward_mqa(  # type: ignore[override]
        cls,
@@ -54,7 +54,7 @@ from vllm.model_executor.models.utils import (
    maybe_prefix,
 )
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.models.deepseek_v4.nvidia.ops.attention import (
+from vllm.models.deepseek_v4.attention import (
    DeepseekV4Indexer,
    DeepseekV4MLAModules,
    DeepseekV4MultiHeadLatentAttentionWrapper,
@@ -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):
@@ -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",
Author	SHA1	Message	Date
Wentao YeandGitHub	bfb22a7d1d	Merge branch 'main' into wentao-fix-flashinfer-layout	2026-05-26 14:48:26 -04:00
Woosuk KwonandGitHub	193ce8812e	[DSv4] Drop _get_compressed_kv_buffer in DeepseekCompressor (#43690 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-05-26 10:11:25 -07:00
Woosuk Kwon GitHub Mark McLoughlin	3aea37d28e	[Doc] Add line limit to AGENTS.md (#43635 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>	2026-05-26 09:31:23 -07:00
Wei-Ming ChenandGitHub	6f5b533241	Add LM head quantization support for ModelOpt (#42124 ) Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>	2026-05-26 09:21:05 -07:00
Woosuk KwonandGitHub	c8414a8271	[ROCm] Remove MegaMoE integration in deepseek v4 (#43629 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-05-26 08:56:04 -07:00
bnellnm GitHub Robert Shaw	f51bbc694d	[MoE Refactor] W4a8 int8 oracle (#42789 ) Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-05-26 11:15:42 -04:00
bnellnm GitHub Robert Shaw	b226ddacfd	[MoE Refactor] Migrate ModelOptMxFp8FusedMoE to oracle (#42768 ) Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>	2026-05-26 11:14:14 -04:00
Yongye ZhuandGitHub	6ab6ffb428	[Feat][DSV4] Fuse q pad into deepseek v4 fused kernel (#43162 )	2026-05-26 05:12:54 -10:00
Andreas KaratzasandGitHub	445ded18c1	[ROCm][CI] Extend ROCm quick reduce coverage (#40990 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-05-26 21:57:13 +08:00
Simon Danielsson GitHubgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	d565357a90	[Docs][ROCm] MoRI-IO Connector Usage Guide (#43603 ) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Signed-off-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-05-26 21:52:30 +08:00
Mohammad Miadh AngkadandGitHub	a970fb5a1a	Fix CuPy runtime deps and restore humming (#43530 ) Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>	2026-05-26 05:59:40 -07:00
Chaojun ZhangandGitHub	861b97765d	[XPU] Fix fused MoE LoRA kernel crash on XPU by using platform-agnos num_compute_units (#43646 ) Signed-off-by: Chaojun,Zhang <chaojun.zhang@intel.com>	2026-05-26 03:40:32 -07:00
Javier De Jesus GitHub Cyrus Leung	ebd0692f80	[Model] Use AutoWeightsLoader for InternLM2 (#38278 ) Signed-off-by: Jesus De Jesus <dejesus.9297@gmail.com> Signed-off-by: javierdejesusda <javier.dejesusj9@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>	2026-05-26 03:39:26 -07:00
linzm1007 GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	739af5c7e1	[Reasoning] [Bugfix] Reject invalid thinking_token_budget values (#43402 ) Signed-off-by: linzm1007 <linzm1007@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-05-26 03:37:30 -07:00
Thibault CastellsandGitHub	5d09f471f4	[Misc] Support interleaved custom image benchmark datasets (#43636 ) Signed-off-by: ThibaultCastells <thib.castells@icloud.com>	2026-05-26 03:37:25 -07:00
Simon Danielsson GitHub TJian	681d7dd38b	[Misc][Refactor][ROCm] Convert MoRI-related envvars to extra config args (#43303 ) Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>	2026-05-26 03:33:35 -07:00
Ethan FengandGitHub	755043cf3c	[KV Transfer] Enable HMA by default for connectors that support it (#41847 ) Signed-off-by: Ethan Feng <ethan.fengch@gmail.com>	2026-05-26 12:28:51 +02:00
yewentao256	8ea74c05c8	fix flashinfer layout Signed-off-by: yewentao256 <zhyanwentao@126.com>	2026-05-19 11:30:36 -04:00